diff --git a/docs/advanced-setup.md b/docs/advanced-setup.md index e447b7eef0..8f096f55c9 100644 --- a/docs/advanced-setup.md +++ b/docs/advanced-setup.md @@ -478,6 +478,7 @@ host. Without this variable the behavior is unchanged. | `CODEX_AUTH_JSON_PATH` | Codex only | Path to a Codex CLI `auth.json` file | | `CODEX_HOME` | Codex only | Alternative Codex home directory | | `OPENCLAUDE_MAX_RETRIES` | No | Maximum retry attempts for retryable API failures, capped at 100 (default: 10). Set to `0` to disable retries after the initial request. If unset, deprecated `CLAUDE_CODE_MAX_RETRIES` is still honored for compatibility. | +| `OPENCLAUDE_MAX_TURNS` | No | Per-prompt **local** interactive REPL turn cap for the in-process query loop. Defaults to `50`. Set a larger positive integer for long autonomous local interactive sessions (for example models that take many small tool steps). CLI `--max-turns 0` explicitly disables this cap and prints a cautionary warning. Precedence for a valid override: CLI `--max-turns` → this env var → legacy `CLAUDE_CODE_MAX_TURNS` (only when this var is unset/empty) → `/config` → Max turns (interactive) → `50`. If this env var is set but invalid (zero, negative, non-integer), the default `50` is used and lower layers are not consulted — same pattern as `OPENCLAUDE_MAX_RETRIES`. Does not apply to remote-backed interactive sessions (`connect` / `ssh` / `--remote`). | | `OPENCLAUDE_RETRY_DELAY_MS` | No | Base retry delay in milliseconds for APIs that do not send `Retry-After`; exponential backoff starts from this value, capped at 60000 (default: 500) | | `OPENCLAUDE_QUERY_HARD_MAX_MS` | No | Foreground query hard maximum in milliseconds. Defaults to 1800000 (30 minutes). Use a larger positive integer for long autonomous sessions; invalid, zero, negative, fractional, or timer-overflow values are ignored with a warning. | | `OPENCLAUDE_DISABLE_CO_AUTHORED_BY` | No | Suppress the default `Co-Authored-By` trailer in generated git commits | diff --git a/src/cli/printMaxTurns.test.ts b/src/cli/printMaxTurns.test.ts new file mode 100644 index 0000000000..2153661918 --- /dev/null +++ b/src/cli/printMaxTurns.test.ts @@ -0,0 +1,223 @@ +import { afterEach, describe, expect, spyOn, test } from 'bun:test' + +import type { Command } from '../commands.js' +import type { SDKMessage } from '../entrypoints/agentSdkTypes.js' +import * as queryEngineModule from '../QueryEngine.js' +import * as queryModule from '../query.js' +import { type QueryParams } from '../query.js' +import type { QueryDeps } from '../query/deps.js' +import { getDefaultAppState, type AppState } from '../state/AppStateStore.js' +import type { Tools } from '../Tool.js' +import { + createAssistantMessage, + createUserMessage, +} from '../utils/messages.js' +import * as processModule from '../utils/process.js' +import { asSystemPrompt } from '../utils/systemPromptType.js' +import { parseMaxTurnsCli, resolveReplMaxTurns } from '../utils/replMaxTurns.js' +import { runHeadless } from './print.js' + +function makeHeadlessQueryParams(maxTurns: number | undefined): QueryParams { + return { + messages: [createUserMessage({ content: 'headless prompt' })], + systemPrompt: asSystemPrompt([]), + userContext: {}, + systemContext: {}, + canUseTool: async () => ({ behavior: 'allow' }), + maxTurns, + querySource: 'sdk', + toolUseContext: { + abortController: new AbortController(), + agentId: 'agent-test', + getAppState: () => ({ + fastMode: false, + mcp: { tools: [], clients: [] }, + toolPermissionContext: { + mode: 'default', + additionalWorkingDirectories: new Map(), + alwaysAllowRules: {}, + alwaysDenyRules: {}, + alwaysAskRules: {}, + isBypassPermissionsModeAvailable: false, + }, + sessionHooks: new Map(), + mainLoopModel: 'gpt-4o', + effortValue: undefined, + advisorModel: undefined, + }), + options: { + commands: [], + debug: false, + thinkingConfig: { type: 'disabled' }, + tools: [], + verbose: false, + mcpClients: [], + mcpResources: {}, + isNonInteractiveSession: true, + agentDefinitions: { activeAgents: [], allAgents: [] }, + appendSystemPrompt: undefined, + providerOverride: undefined, + mainLoopModel: 'gpt-4o', + }, + addNotification: () => {}, + messages: [], + setInProgressToolUseIDs: () => {}, + setResponseLength: () => {}, + updateAttributionState: () => {}, + } as unknown as QueryParams['toolUseContext'], + deps: { + callModel: async function* () { + yield createAssistantMessage({ content: 'done' }) + }, + microcompact: async messages => ({ messages }), + autocompact: async () => ({ + compactionResult: null, + consecutiveFailures: undefined, + }), + uuid: () => '00000000-0000-4000-8000-000000000000', + } as unknown as QueryDeps, + } +} + +function createHeadlessSuccessResult(): SDKMessage { + return { + type: 'result', + subtype: 'success', + duration_ms: 0, + duration_api_ms: 0, + is_error: false, + num_turns: 1, + result: 'done', + stop_reason: 'end_turn', + total_cost_usd: 0, + usage: { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + }, + modelUsage: {}, + permission_denials: [], + uuid: '00000000-0000-4000-8000-000000000001', + session_id: 'test-session', + } +} + +function createHeadlessRunOptions(maxTurns: number | undefined) { + return { + continue: undefined, + resume: undefined, + fromPr: undefined, + resumeSessionAt: undefined, + verbose: false, + outputFormat: 'text', + jsonSchema: undefined, + permissionPromptToolName: undefined, + allowedTools: undefined, + thinkingConfig: { type: 'disabled' as const }, + maxTurns, + maxBudgetUsd: undefined, + taskBudget: undefined, + systemPrompt: undefined, + appendSystemPrompt: undefined, + userSpecifiedModel: undefined, + fallbackModel: undefined, + teleport: undefined, + sdkUrl: undefined, + replayUserMessages: undefined, + includePartialMessages: undefined, + forkSession: undefined, + rewindFiles: undefined, + enableAuthStatus: undefined, + agent: undefined, + workload: undefined, + } +} + +async function waitForAskCall( + askSpy: ReturnType>, +): Promise { + const deadline = Date.now() + 15_000 + while (askSpy.mock.calls.length === 0) { + if (Date.now() >= deadline) { + throw new Error('Timed out waiting for runHeadless to call ask()') + } + await new Promise(resolve => setTimeout(resolve, 25)) + } +} + +describe('headless --print max-turns', () => { + const savedSimple = process.env.CLAUDE_CODE_SIMPLE + + afterEach(() => { + if (savedSimple === undefined) { + delete process.env.CLAUDE_CODE_SIMPLE + } else { + process.env.CLAUDE_CODE_SIMPLE = savedSimple + } + }) + + test('forwards parsed --max-turns 0 into query params without interactive resolution', async () => { + const headlessMaxTurns = parseMaxTurnsCli('0') + expect(headlessMaxTurns).toBe(0) + expect(resolveReplMaxTurns(headlessMaxTurns)).toBeUndefined() + + const querySpy = spyOn(queryModule, 'query') + const params = makeHeadlessQueryParams(headlessMaxTurns) + + const generator = queryModule.query(params) + let terminal + while (true) { + const next = await generator.next() + if (next.done) { + terminal = next.value + break + } + } + + expect(querySpy.mock.calls[0]?.[0]?.maxTurns).toBe(0) + expect(terminal?.reason).toBe('completed') + querySpy.mockRestore() + }) + + test('forwards maxTurns 0 through runHeadless into ask()', async () => { + process.env.CLAUDE_CODE_SIMPLE = '1' + + const stdoutSpy = spyOn(processModule, 'writeToStdout').mockImplementation( + () => {}, + ) + const askSpy = spyOn(queryEngineModule, 'ask').mockImplementation( + async function* () { + yield createHeadlessSuccessResult() + }, + ) + + let state = getDefaultAppState() + const getAppState = () => state + const setAppState = (update: (previous: AppState) => AppState) => { + state = update(state) + } + + let runPromise: Promise | undefined + try { + runPromise = runHeadless( + 'headless prompt', + getAppState, + setAppState, + [] as Command[], + [] as Tools, + {}, + [], + createHeadlessRunOptions(0), + ) + + await waitForAskCall(askSpy) + expect(askSpy.mock.calls[0]?.[0]?.maxTurns).toBe(0) + await runPromise + } finally { + await runPromise?.catch(() => {}) + askSpy.mockRestore() + stdoutSpy.mockRestore() + } + }) +}) diff --git a/src/components/SessionBackgroundHint.tsx b/src/components/SessionBackgroundHint.tsx index bd5e224c18..5be8328e53 100644 --- a/src/components/SessionBackgroundHint.tsx +++ b/src/components/SessionBackgroundHint.tsx @@ -47,6 +47,8 @@ export function SessionBackgroundHint(t0) { saveGlobalConfig(_temp2); } } else { + // Session backgrounding remains intentionally disabled until the + // handoff path is ready to ship; shell/agent backgrounding still runs. if (isEnvTruthy("false") && isLoading) { handleDoublePress(); } diff --git a/src/components/Settings/Config.tsx b/src/components/Settings/Config.tsx index 815c58e09a..89d246efbe 100644 --- a/src/components/Settings/Config.tsx +++ b/src/components/Settings/Config.tsx @@ -11,6 +11,7 @@ import { type GlobalConfig, saveGlobalConfig, getCurrentProjectConfig, type Outp import { normalizeApiKeyForConfig } from '../../utils/authPortable.js'; import { getGlobalConfig, getAutoUpdaterDisabledReason, formatAutoUpdaterDisabledReason, getRemoteControlAtStartup } from '../../utils/config.js'; import { normalizeCompactTailTurns } from '../../utils/relevancePruning.js'; +import { normalizeReplMaxTurns, REPL_MAX_TURNS_OPTIONS } from '../../utils/replMaxTurns.js'; import chalk from 'chalk'; import { getModeColor, permissionModeTitle, permissionModeFromString, toExternalPermissionMode, isExternalPermissionMode, PERMISSION_MODES, type ExternalPermissionMode, type PermissionMode } from '../../utils/permissions/PermissionMode.js'; import { getAutoModeEnabledState, hasAutoModeOptInAnySource, transitionPlanAutoMode } from '../../utils/permissions/permissionSetup.js'; @@ -335,6 +336,29 @@ export function Config({ value: compactTailTurnsValue as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS }); } + }, { + id: 'replMaxTurns', + label: 'Max turns (interactive)', + // Display/persist the saved preference (normalized). Effective runtime cap + // may still be overridden by CLI `--max-turns` or OPENCLAUDE_MAX_TURNS. + value: String(normalizeReplMaxTurns(globalConfig.replMaxTurns)), + // Include a hand-edited config value so it round-trips through the picker. + options: [...new Set([...REPL_MAX_TURNS_OPTIONS.map(String), String(normalizeReplMaxTurns(globalConfig.replMaxTurns))])], + type: 'enum' as const, + onChange(replMaxTurnsValue: string) { + const replMaxTurns = normalizeReplMaxTurns(replMaxTurnsValue); + saveGlobalConfig(current => ({ + ...current, + replMaxTurns + })); + setGlobalConfig({ + ...getGlobalConfig(), + replMaxTurns + }); + logEvent('tengu_repl_max_turns_changed', { + value: replMaxTurnsValue as AnalyticsMetadata_I_VERIFIED_THIS_IS_NOT_CODE_OR_FILEPATHS + }); + } }, { id: 'toolHistoryCompressionEnabled', label: 'Tool history compression', @@ -1272,6 +1296,9 @@ export function Config({ if (globalConfig.compactTailTurns !== initialConfig.current.compactTailTurns) { formattedChanges.push(`Set compaction recent messages kept to ${normalizeCompactTailTurns(globalConfig.compactTailTurns)}`); } + if (globalConfig.replMaxTurns !== initialConfig.current.replMaxTurns) { + formattedChanges.push(`Set interactive max turns to ${normalizeReplMaxTurns(globalConfig.replMaxTurns)}`); + } if (globalConfig.toolHistoryCompressionEnabled !== initialConfig.current.toolHistoryCompressionEnabled) { formattedChanges.push(`${globalConfig.toolHistoryCompressionEnabled ? 'Enabled' : 'Disabled'} tool history compression`); } diff --git a/src/main.tsx b/src/main.tsx index a7832be9b6..be9171f38c 100644 --- a/src/main.tsx +++ b/src/main.tsx @@ -123,6 +123,7 @@ import { getModelDeprecationWarning } from './utils/model/deprecation.js'; import { getDefaultMainLoopModel, getUserSpecifiedModelSetting, normalizeModelStringForAPI, parseUserSpecifiedModel } from './utils/model/model.js'; import { ensureModelStringsInitialized } from './utils/model/modelStrings.js'; import { PERMISSION_MODES } from './utils/permissions/PermissionMode.js'; +import { MAX_TURNS_CLI_DESCRIPTION, parseMaxTurnsCommanderArgument } from './utils/replMaxTurns.js'; import { checkAndDisableBypassPermissions, getAutoModeEnabledStateIfCached, initializeToolPermissionContext, initialPermissionModeFromCLI, isDefaultPermissionModeAuto, parseToolListFromCLI, stripDangerousPermissionsForAutoMode, verifyAutoModeGateAccess } from './utils/permissions/permissionSetup.js'; import { cleanupOrphanedPluginVersionsInBackground } from './utils/plugins/cacheUtils.js'; import { initializeVersionedPlugins } from './utils/plugins/installedPluginsManager.js'; @@ -945,7 +946,9 @@ async function run(): Promise { } catch (error) { throw new InvalidArgumentError(errorMessage(error)); } - })).option('--bare', 'Minimal mode: skip hooks, LSP, plugin sync, attribution, auto-memory, background prefetches, keychain reads, and CLAUDE.md auto-discovery. Sets CLAUDE_CODE_SIMPLE=1. Anthropic auth is strictly ANTHROPIC_API_KEY or apiKeyHelper via --settings (OAuth and keychain are never read). 3P providers (Bedrock/Vertex/Foundry) use their own credentials. Skills still resolve via /skill-name. Explicitly provide context via: --system-prompt[-file], --append-system-prompt[-file], --add-dir (CLAUDE.md dirs), --mcp-config, --settings, --agents, --plugin-dir.', () => true).addOption(new Option('--init', 'Run Setup hooks with init trigger, then continue').hideHelp()).addOption(new Option('--init-only', 'Run Setup and SessionStart:startup hooks, then exit').hideHelp()).addOption(new Option('--maintenance', 'Run Setup hooks with maintenance trigger, then continue').hideHelp()).addOption(new Option('--output-format ', 'Output format (only works with --print): "text" (default), "json" (single result), or "stream-json" (realtime streaming)').choices(['text', 'json', 'stream-json'])).addOption(new Option('--json-schema ', 'JSON Schema for structured output validation. ' + 'Example: {"type":"object","properties":{"name":{"type":"string"}},"required":["name"]}').argParser(String)).option('--include-hook-events', 'Include all hook lifecycle events in the output stream (only works with --output-format=stream-json)', () => true).option('--include-partial-messages', 'Include partial message chunks as they arrive (only works with --print and --output-format=stream-json)', () => true).addOption(new Option('--input-format ', 'Input format (only works with --print): "text" (default), or "stream-json" (realtime streaming input)').choices(['text', 'stream-json'])).option('--mcp-debug', '[DEPRECATED. Use --debug instead] Enable MCP debug mode (shows MCP server errors)', () => true).option('--dangerously-skip-permissions', 'Bypass all permission checks. Recommended only for sandboxes with no internet access.', () => true).option('--allow-dangerously-skip-permissions', 'Enable bypassing all permission checks as an option, without it being enabled by default. Recommended only for sandboxes with no internet access.', () => true).addOption(new Option('--thinking ', 'Thinking mode: enabled (equivalent to adaptive), disabled').choices(['enabled', 'adaptive', 'disabled']).hideHelp()).addOption(new Option('--max-thinking-tokens ', '[DEPRECATED. Use --thinking instead for newer models] Maximum number of thinking tokens (only works with --print)').argParser(Number).hideHelp()).addOption(new Option('--max-turns ', 'Maximum number of agentic turns in non-interactive mode. This will early exit the conversation after the specified number of turns. (only works with --print)').argParser(Number).hideHelp()).addOption(new Option('--max-budget-usd ', 'Maximum dollar amount to spend on API calls (only works with --print)').argParser(value => { + })).option('--bare', 'Minimal mode: skip hooks, LSP, plugin sync, attribution, auto-memory, background prefetches, keychain reads, and CLAUDE.md auto-discovery. Sets CLAUDE_CODE_SIMPLE=1. Anthropic auth is strictly ANTHROPIC_API_KEY or apiKeyHelper via --settings (OAuth and keychain are never read). 3P providers (Bedrock/Vertex/Foundry) use their own credentials. Skills still resolve via /skill-name. Explicitly provide context via: --system-prompt[-file], --append-system-prompt[-file], --add-dir (CLAUDE.md dirs), --mcp-config, --settings, --agents, --plugin-dir.', () => true).addOption(new Option('--init', 'Run Setup hooks with init trigger, then continue').hideHelp()).addOption(new Option('--init-only', 'Run Setup and SessionStart:startup hooks, then exit').hideHelp()).addOption(new Option('--maintenance', 'Run Setup hooks with maintenance trigger, then continue').hideHelp()).addOption(new Option('--output-format ', 'Output format (only works with --print): "text" (default), "json" (single result), or "stream-json" (realtime streaming)').choices(['text', 'json', 'stream-json'])).addOption(new Option('--json-schema ', 'JSON Schema for structured output validation. ' + 'Example: {"type":"object","properties":{"name":{"type":"string"}},"required":["name"]}').argParser(String)).option('--include-hook-events', 'Include all hook lifecycle events in the output stream (only works with --output-format=stream-json)', () => true).option('--include-partial-messages', 'Include partial message chunks as they arrive (only works with --print and --output-format=stream-json)', () => true).addOption(new Option('--input-format ', 'Input format (only works with --print): "text" (default), or "stream-json" (realtime streaming input)').choices(['text', 'stream-json'])).option('--mcp-debug', '[DEPRECATED. Use --debug instead] Enable MCP debug mode (shows MCP server errors)', () => true).option('--dangerously-skip-permissions', 'Bypass all permission checks. Recommended only for sandboxes with no internet access.', () => true).option('--allow-dangerously-skip-permissions', 'Enable bypassing all permission checks as an option, without it being enabled by default. Recommended only for sandboxes with no internet access.', () => true).addOption(new Option('--thinking ', 'Thinking mode: enabled (equivalent to adaptive), disabled').choices(['enabled', 'adaptive', 'disabled']).hideHelp()).addOption(new Option('--max-thinking-tokens ', '[DEPRECATED. Use --thinking instead for newer models] Maximum number of thinking tokens (only works with --print)').argParser(Number).hideHelp()).addOption(new Option('--max-turns ', MAX_TURNS_CLI_DESCRIPTION).argParser(value => { + return parseMaxTurnsCommanderArgument(value); + })).addOption(new Option('--max-budget-usd ', 'Maximum dollar amount to spend on API calls (only works with --print)').argParser(value => { const amount = Number(value); if (isNaN(amount) || amount <= 0) { throw new Error('--max-budget-usd must be a positive number greater than 0'); @@ -3049,6 +3052,8 @@ async function run(): Promise { systemPrompt, appendSystemPrompt, thinkingConfig, + // Interactive REPL default is 50 via resolveReplMaxTurns; CLI wins over env. + maxTurns: options.maxTurns, }; // Shared context for processResumedConversation calls @@ -3149,7 +3154,8 @@ async function run(): Promise { mainThreadAgentDefinition, disableSlashCommands, directConnectConfig, - thinkingConfig + thinkingConfig, + maxTurns: options.maxTurns, }, renderAndRun); return; } else if (feature('SSH_REMOTE') && _pendingSSH?.host) { @@ -3215,7 +3221,8 @@ async function run(): Promise { mainThreadAgentDefinition, disableSlashCommands, sshSession, - thinkingConfig + thinkingConfig, + maxTurns: options.maxTurns, }, renderAndRun); return; } else if (feature('KAIROS') && _pendingAssistantChat && (_pendingAssistantChat.sessionId || _pendingAssistantChat.discover)) { @@ -3311,7 +3318,8 @@ async function run(): Promise { mainThreadAgentDefinition, disableSlashCommands, remoteSessionConfig, - thinkingConfig + thinkingConfig, + maxTurns: options.maxTurns, }, renderAndRun); return; } else if (options.resume || options.fromPr || teleport || remote !== null) { @@ -3460,7 +3468,8 @@ async function run(): Promise { mainThreadAgentDefinition, disableSlashCommands, remoteSessionConfig, - thinkingConfig + thinkingConfig, + maxTurns: options.maxTurns, }, renderAndRun); return; } else if (teleport) { diff --git a/src/query.ts b/src/query.ts index 1e45817655..8dc317f78e 100644 --- a/src/query.ts +++ b/src/query.ts @@ -52,6 +52,7 @@ import { logAntError, logForDebugging } from './utils/debug.js' import { getMissingToolResultAbortMessage, getQueryAbortSystemMessage, + normalizeAbortReason, shouldCreateUserInterruptionMessage, } from './utils/abortReasons.js' import { @@ -165,6 +166,71 @@ const taskSummaryModule = feature('BG_SESSIONS') : null /* eslint-enable @typescript-eslint/no-require-imports */ +async function cleanupComputerUseAtTerminal( + toolUseContext: ToolUseContext, +): Promise { + // feature() must remain the direct condition so external builds eliminate + // the native Computer Use dependency at bundle time. + if (feature('CHICAGO_MCP')) { + if (toolUseContext.agentId) return + try { + const { cleanupComputerUseAfterTurn } = await import( + './utils/computerUse/cleanup.js' + ) + await cleanupComputerUseAfterTurn(toolUseContext) + } catch { + // Failures are silent — this is dogfooding cleanup, not critical path. + } + } +} + +async function* emitAbortedStreaming( + signal: AbortSignal, + toolUseContext: ToolUseContext, +): AsyncGenerator< + Message, + Extract +> { + await cleanupComputerUseAtTerminal(toolUseContext) + const abortReason = signal.reason + const abortSystemMessage = getQueryAbortSystemMessage(abortReason) + if (abortSystemMessage) { + yield createSystemMessage(abortSystemMessage, 'warning') + } + if (shouldCreateUserInterruptionMessage(abortReason)) { + yield createUserInterruptionMessage({ toolUse: false }) + } + return { reason: 'aborted_streaming' } +} + +function* emitAbortedToolsAfterCleanup( + signal: AbortSignal, + maxTurns: number | undefined, + nextTurnCount: number, + hasSharedTurnBudget: boolean, +): Generator> { + const abortReason = signal.reason + const abortSystemMessage = getQueryAbortSystemMessage(abortReason) + if (abortSystemMessage) { + yield createSystemMessage(abortSystemMessage, 'warning') + } + if (shouldCreateUserInterruptionMessage(abortReason)) { + yield createUserInterruptionMessage({ toolUse: true }) + } + if ( + maxTurns && + nextTurnCount > maxTurns && + (!hasSharedTurnBudget || normalizeAbortReason(abortReason) !== 'background') + ) { + yield createAttachmentMessage({ + type: 'max_turns_reached', + maxTurns, + turnCount: nextTurnCount, + }) + } + return { reason: 'aborted_tools' } +} + function* yieldMissingToolResultBlocks( assistantMessages: AssistantMessage[], errorMessage: string, @@ -435,6 +501,8 @@ export type QueryParams = { /** Called around each outbound model request, including retries. */ onModelRequestStart?: () => void onModelRequestEnd?: () => void + /** Called once provider dispatch is accepted for the current attempt. */ + onProviderDispatchAccepted?: () => void systemPrompt: SystemPrompt userContext: { [k: string]: string } systemContext: { [k: string]: string } @@ -444,6 +512,11 @@ export type QueryParams = { querySource: QuerySource maxOutputTokensOverride?: number maxTurns?: number + /** + * Mutable per-prompt budget shared by query() calls that continue the same + * logical prompt (for example, when a local REPL query is backgrounded). + */ + turnBudget?: QueryTurnBudget skipCacheWrite?: boolean autoCompactTracking?: AutoCompactTrackingState onAutoCompactTrackingChange?: ( @@ -458,6 +531,17 @@ export type QueryParams = { deps?: QueryDeps } +export type QueryTurnBudget = { + readonly maxTurns: number | undefined + turnsStarted: number +} + +export function createQueryTurnBudget( + maxTurns?: number, +): QueryTurnBudget { + return { maxTurns, turnsStarted: 0 } +} + /** * `ultrathink_effort` is emitted while processing the current user input. * Its attachment is deliberately transient, but the request-level effort @@ -595,9 +679,14 @@ async function* queryLoop( canUseTool, fallbackModel, querySource, - maxTurns, skipCacheWrite, } = params + const maxTurns = params.turnBudget + ? params.turnBudget.maxTurns + : params.maxTurns + const initialTurnCount = params.turnBudget + ? params.turnBudget.turnsStarted + 1 + : 1 const deps = params.deps ?? productionDeps() const ultrathinkEffortForCurrentTurn = hasUltrathinkEffortForCurrentTurn( params.messages, @@ -617,7 +706,7 @@ async function* queryLoop( hasAttemptedReactiveCompact: false, hasAttemptedContextOverflowRecovery: false, hasAttemptedProviderFallback: false, - turnCount: 1, + turnCount: initialTurnCount, continuationNudgeCount: 0, pendingToolUseSummary: undefined, transition: undefined, @@ -659,11 +748,42 @@ async function* queryLoop( // at the new endpoint — KTD6 in the plan. let pinnedRouteProviderId: string | undefined = undefined const toolFailureGuardState = createToolFailureLoopGuardState() + // Identifies the turn this queryLoop invocation claimed in a shared budget. + // Retries for that turn are allowed; a different invocation that snapped the + // same next turn is stale and must not dispatch a duplicate provider call. + let reservedTurnCount: number | undefined = undefined // Snapshot immutable env/statsig/session state once at entry. See QueryConfig // for what's included and why feature() gates are intentionally excluded. const config = buildQueryConfig() + // Ctrl+B can abort the foreground while it is still preparing query + // context. Let the background continuation reserve the turn in that race; + // the aborted invocation never reached a provider request and must not + // consume the shared prompt budget. + if ( + params.turnBudget && + state.toolUseContext.abortController.signal.aborted + ) { + return yield* emitAbortedStreaming( + state.toolUseContext.abortController.signal, + state.toolUseContext, + ) + } + + // Reject an invocation that cannot start another provider turn. Do not + // reserve it here: context preparation below can await, and Ctrl+B may hand + // the prompt off before any provider request is dispatched. + if (maxTurns && state.turnCount > maxTurns) { + await cleanupComputerUseAtTerminal(state.toolUseContext) + yield createAttachmentMessage({ + type: 'max_turns_reached', + maxTurns, + turnCount: state.turnCount, + }) + return { reason: 'max_turns', turnCount: state.turnCount } + } + // Fired once per user turn — the prompt is invariant across loop iterations, // so per-iteration firing would ask sideQuery the same question N times. // Consume point polls settledAt (never blocks). `using` disposes on all @@ -1369,8 +1489,17 @@ async function* queryLoop( try { let streamingFallbackOccured = false queryCheckpoint('query_api_streaming_start') - params.onModelRequestStart?.() + // queryModel performs provider-specific asynchronous preparation. + // Claim the turn from its callback immediately before the actual + // request is dispatched, not merely when callModel is entered. + let providerDispatchRejected = false + let providerDispatchAccepted = false + let modelRequestLifecycleStarted = false try { + // Arm interruption correction and other per-attempt hooks before + // callModel performs async provider preparation. + params.onModelRequestStart?.() + modelRequestLifecycleStarted = true for await (const message of deps.callModel({ messages: prependUserContext( injectRequestOnlyMessages( @@ -1413,6 +1542,36 @@ async function* queryLoop( ), queryTracking, queryLifecycle: toolUseContext.queryLifecycle, + onProviderRequestStart: () => { + if (toolUseContext.abortController.signal.aborted) { + providerDispatchRejected = true + return false + } + // Retries reuse this turn's reservation, but they must still + // prove the foreground owns dispatch after any asynchronous + // credential refresh or client recreation. + if (providerDispatchAccepted) return true + if ( + params.turnBudget && + params.turnBudget.turnsStarted >= turnCount && + reservedTurnCount !== turnCount + ) { + providerDispatchRejected = true + return false + } + // Fallback attempts reuse turnCount. The local reservation + // makes retries idempotent while rejecting a stale claimant. + if ( + params.turnBudget && + params.turnBudget.turnsStarted < turnCount + ) { + params.turnBudget.turnsStarted = turnCount + reservedTurnCount = turnCount + } + providerDispatchAccepted = true + params.onProviderDispatchAccepted?.() + return true + }, // Explicit /effort selection wins. When it is unset, carry the // current turn's ultrathink attachment through to the API client // so OpenAI-compatible providers receive reasoning_effort=high @@ -1624,8 +1783,17 @@ async function* queryLoop( } } } + if (providerDispatchRejected) { + if (toolUseContext.abortController.signal.aborted) { + return yield* emitAbortedStreaming( + toolUseContext.abortController.signal, + toolUseContext, + ) + } + return { reason: 'aborted_streaming' } + } } finally { - params.onModelRequestEnd?.() + if (modelRequestLifecycleStarted) params.onModelRequestEnd?.() } queryCheckpoint('query_api_streaming_end') @@ -1852,16 +2020,7 @@ async function* queryLoop( // chicago MCP: auto-unhide + lock release on interrupt. Same cleanup // as the natural turn-end path in stopHooks.ts. Main thread only — // see stopHooks.ts for the subagent-releasing-main's-lock rationale. - if (feature('CHICAGO_MCP') && !toolUseContext.agentId) { - try { - const { cleanupComputerUseAfterTurn } = await import( - './utils/computerUse/cleanup.js' - ) - await cleanupComputerUseAfterTurn(toolUseContext) - } catch { - // Failures are silent — this is dogfooding cleanup, not critical path - } - } + await cleanupComputerUseAtTerminal(toolUseContext) const abortSystemMessage = getQueryAbortSystemMessage(abortReason) if (abortSystemMessage) { @@ -2624,40 +2783,16 @@ async function* queryLoop( // We were aborted during tool calls if (toolUseContext.abortController.signal.aborted) { - const abortReason = toolUseContext.abortController.signal.reason // chicago MCP: auto-unhide + lock release when aborted mid-tool-call. // This is the most likely Ctrl+C path for CU (e.g. slow screenshot). // Main thread only — see stopHooks.ts for the subagent rationale. - if (feature('CHICAGO_MCP') && !toolUseContext.agentId) { - try { - const { cleanupComputerUseAfterTurn } = await import( - './utils/computerUse/cleanup.js' - ) - await cleanupComputerUseAfterTurn(toolUseContext) - } catch { - // Failures are silent — this is dogfooding cleanup, not critical path - } - } - const abortSystemMessage = getQueryAbortSystemMessage(abortReason) - if (abortSystemMessage) { - yield createSystemMessage(abortSystemMessage, 'warning') - } - - if (shouldCreateUserInterruptionMessage(abortReason)) { - yield createUserInterruptionMessage({ - toolUse: true, - }) - } - // Check maxTurns before returning when aborted - const nextTurnCountOnAbort = turnCount + 1 - if (maxTurns && nextTurnCountOnAbort > maxTurns) { - yield createAttachmentMessage({ - type: 'max_turns_reached', - maxTurns, - turnCount: nextTurnCountOnAbort, - }) - } - return { reason: 'aborted_tools' } + await cleanupComputerUseAtTerminal(toolUseContext) + return yield* emitAbortedToolsAfterCleanup( + toolUseContext.abortController.signal, + maxTurns, + turnCount + 1, + params.turnBudget !== undefined, + ) } // If a hook indicated to prevent continuation, stop here @@ -2977,6 +3112,18 @@ async function* queryLoop( nextTurnCount > maxTurns && !nextAgentStepLimit?.summaryRequested ) { + await cleanupComputerUseAtTerminal(toolUseContext) + // Attachment/memory/skill collection above can await after the earlier + // post-tool abort check. Re-check immediately before emitting the cap so + // a Ctrl+B handoff cannot make both owners persist the terminal record. + if (toolUseContext.abortController.signal.aborted) { + return yield* emitAbortedToolsAfterCleanup( + toolUseContext.abortController.signal, + maxTurns, + nextTurnCount, + params.turnBudget !== undefined, + ) + } yield createAttachmentMessage({ type: 'max_turns_reached', maxTurns, diff --git a/src/query/agentStepLimit.test.ts b/src/query/agentStepLimit.test.ts index ee9a8f639a..6501d936a1 100644 --- a/src/query/agentStepLimit.test.ts +++ b/src/query/agentStepLimit.test.ts @@ -1,9 +1,14 @@ import { describe, expect, test } from 'bun:test' import { z } from 'zod/v4' -import { query, type QueryParams } from '../query.js' +import { + createQueryTurnBudget, + query, + type QueryParams, +} from '../query.js' import { buildTool, type Tools } from '../Tool.js' import type { QueryDeps } from './deps.js' +import { FallbackTriggeredError } from '../services/api/withRetry.js' import { createAssistantMessage, createUserMessage, @@ -12,6 +17,9 @@ import { import { asSystemPrompt } from '../utils/systemPromptType.js' import { countToolUses } from '../tools/AgentTool/agentToolUtils.js' import { AGENT_STEP_LIMIT_TOOL_RESULT_PREFIX } from './agentStepLimit.js' +import type { Terminal } from './transitions.js' +import type { Message } from '../types/message.js' +import { dequeueAll, enqueue } from '../utils/messageQueueManager.js' const echoCalls: string[] = [] @@ -97,6 +105,10 @@ function makeParams( tools: Tools = [], agentStepLimit?: { maxSteps: number; agentType: string }, ): QueryParams { + const dispatchingCallModel: QueryDeps['callModel'] = async function* (input) { + if (input.options.onProviderRequestStart?.() === false) return + yield* callModel(input) + } return { messages: [createUserMessage({ content: 'inspect' })], systemPrompt: asSystemPrompt([]), @@ -107,7 +119,7 @@ function makeParams( querySource: 'agent:builtin:general-purpose', ...(agentStepLimit ? { agentStepLimit } : {}), deps: { - callModel, + callModel: dispatchingCallModel, microcompact: async messages => ({ messages }), autocompact: async () => ({ compactionResult: null, @@ -132,6 +144,428 @@ async function drain(params: QueryParams): Promise<{ } describe('agent step limits', () => { + test('shares one turn budget across query calls for the same prompt', async () => { + let modelCalls = 0 + const turnBudget = createQueryTurnBudget(1) + const callModel = async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'done' }) + } + + await drain({ + ...makeParams(callModel), + turnBudget, + }) + + const { yielded, returned } = await drain({ + ...makeParams(callModel), + turnBudget, + }) + + expect(modelCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + expect(returned).toEqual({ reason: 'max_turns', turnCount: 2 }) + expect( + yielded.some( + message => + message.type === 'attachment' && + message.attachment.type === 'max_turns_reached' && + message.attachment.turnCount === 2, + ), + ).toBe(true) + }) + + test('provider fallback retries do not consume another shared turn', async () => { + let modelCalls = 0 + const turnBudget = createQueryTurnBudget(1) + const params = makeParams(async function* () { + modelCalls++ + if (modelCalls === 1) { + throw new FallbackTriggeredError('primary-model', 'fallback-model') + } + yield createAssistantMessage({ content: 'completed on fallback' }) + }) + params.fallbackModel = 'fallback-model' + + const first = await drain({ ...params, turnBudget }) + const second = await drain({ + ...makeParams(async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'must not run' }) + }), + turnBudget, + }) + + expect(first.returned).toEqual({ reason: 'completed' }) + expect(second.returned).toEqual({ reason: 'max_turns', turnCount: 2 }) + expect(modelCalls).toBe(2) + expect(turnBudget.turnsStarted).toBe(1) + }) + + test('a retry cannot dispatch after foreground ownership is aborted', async () => { + let providerCalls = 0 + const turnBudget = createQueryTurnBudget(1) + const params = makeParams(async function* () {}) + const abortController = params.toolUseContext.abortController + params.deps!.callModel = async function* ({ options }) { + if (options.onProviderRequestStart?.() === false) return + providerCalls++ + + // Simulate an abort during asynchronous retry preparation. The second + // ownership check must reject even though this turn is already reserved. + abortController.abort('background') + if (options.onProviderRequestStart?.() === false) return + providerCalls++ + yield createAssistantMessage({ content: 'stale retry' }) + } + + const result = await drain({ ...params, turnBudget }) + + expect(result.returned).toEqual({ reason: 'aborted_streaming' }) + expect(providerCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + }) + + test('concurrent query calls cannot claim the same shared turn', async () => { + let modelCalls = 0 + let preparationArrivals = 0 + let releasePreparation!: () => void + const preparationBarrier = new Promise(resolve => { + releasePreparation = resolve + }) + const turnBudget = createQueryTurnBudget(1) + const callModel = async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'one claimant' }) + } + const makeConcurrentParams = (): QueryParams => { + const params = makeParams(callModel) + params.deps!.autocompact = async () => { + preparationArrivals++ + if (preparationArrivals === 2) releasePreparation() + let timeout: ReturnType | undefined + try { + await Promise.race([ + preparationBarrier, + new Promise((_, reject) => { + timeout = setTimeout( + () => reject(new Error('preparation barrier timed out')), + 5_000, + ) + timeout.unref?.() + }), + ]) + } finally { + if (timeout !== undefined) clearTimeout(timeout) + } + return { + wasCompacted: false, + compactionResult: undefined, + consecutiveFailures: undefined, + } + } + return { ...params, turnBudget } + } + + const results = await Promise.all([ + drain(makeConcurrentParams()), + drain(makeConcurrentParams()), + ]) + + expect(modelCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + expect(results.map(result => result.returned.reason).sort()).toEqual([ + 'aborted_streaming', + 'completed', + ]) + }) + + test('does not charge a handoff aborted before its first request', async () => { + let modelCalls = 0 + const turnBudget = createQueryTurnBudget(1) + const callModel = async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'continued in background' }) + } + const foregroundParams = makeParams(callModel) + foregroundParams.toolUseContext.abortController.abort('background') + + const foreground = await drain({ + ...foregroundParams, + turnBudget, + }) + const background = await drain({ + ...makeParams(callModel), + turnBudget, + }) + + expect(foreground.returned).toEqual({ reason: 'aborted_streaming' }) + expect(background.returned).toEqual({ reason: 'completed' }) + expect(modelCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + }) + + test('does not charge a handoff aborted while preparing its first request', async () => { + let modelCalls = 0 + const turnBudget = createQueryTurnBudget(1) + const callModel = async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'continued in background' }) + } + const foregroundParams = makeParams(callModel) + const foregroundAbortController = + foregroundParams.toolUseContext.abortController + foregroundParams.deps!.autocompact = async () => { + foregroundAbortController.abort('background') + return { + wasCompacted: false, + compactionResult: undefined, + consecutiveFailures: undefined, + } + } + + const foreground = await drain({ + ...foregroundParams, + turnBudget, + }) + const background = await drain({ + ...makeParams(callModel), + turnBudget, + }) + + expect(foreground.returned).toEqual({ reason: 'aborted_streaming' }) + expect(background.returned).toEqual({ reason: 'completed' }) + expect(modelCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + }) + + test('does not charge a handoff aborted during provider preparation', async () => { + let providerCalls = 0 + let releaseProviderPreparation!: () => void + let providerPreparationStarted!: () => void + const providerPreparation = new Promise(resolve => { + releaseProviderPreparation = resolve + }) + const providerPreparationEntry = new Promise(resolve => { + providerPreparationStarted = resolve + }) + const turnBudget = createQueryTurnBudget(1) + const foregroundParams = makeParams(async function* () {}) + const foregroundAbortController = + foregroundParams.toolUseContext.abortController + foregroundParams.deps!.callModel = async function* (input) { + providerPreparationStarted() + await providerPreparation + if (input.options.onProviderRequestStart?.() === false) return + providerCalls++ + yield createAssistantMessage({ content: 'stale foreground request' }) + } + + const foregroundPromise = drain({ + ...foregroundParams, + turnBudget, + }) + await providerPreparationEntry + foregroundAbortController.abort('background') + releaseProviderPreparation() + const foreground = await foregroundPromise + const background = await drain({ + ...makeParams(async function* () { + providerCalls++ + yield createAssistantMessage({ content: 'continued in background' }) + }), + turnBudget, + }) + + expect(foreground.returned).toEqual({ reason: 'aborted_streaming' }) + expect(background.returned).toEqual({ reason: 'completed' }) + expect(providerCalls).toBe(1) + expect(turnBudget.turnsStarted).toBe(1) + }) + + test('background handoff emits the final-turn cap only once', async () => { + let modelCalls = 0 + let foregroundAbortController: AbortController + const backgroundingTool = buildTool({ + name: 'Background', + inputSchema: z.object({}), + maxResultSizeChars: Infinity, + async description() { + return 'Background the query' + }, + async prompt() { + return '' + }, + async call() { + foregroundAbortController.abort('background') + return { data: 'backgrounded' } + }, + mapToolResultToToolResultBlockParam(content, toolUseID) { + return { + type: 'tool_result', + tool_use_id: toolUseID, + content: String(content), + } + }, + renderToolUseMessage() { + return null + }, + renderToolResultMessage() { + return null + }, + }) + const foregroundParams = makeParams( + async function* () { + modelCalls++ + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_background_1', + name: 'Background', + input: {}, + }, + ], + }) + }, + [backgroundingTool], + ) + foregroundAbortController = + foregroundParams.toolUseContext.abortController + const turnBudget = createQueryTurnBudget(1) + + const foreground = await drain({ + ...foregroundParams, + turnBudget, + }) + const background = await drain({ + ...makeParams(async function* () { + modelCalls++ + yield createAssistantMessage({ content: 'must not run' }) + }), + turnBudget, + }) + const capAttachments = [...foreground.yielded, ...background.yielded].filter( + message => + message.type === 'attachment' && + message.attachment.type === 'max_turns_reached', + ) + + expect(foreground.returned).toEqual({ reason: 'aborted_tools' }) + expect(background.returned).toEqual({ + reason: 'max_turns', + turnCount: 2, + }) + expect(modelCalls).toBe(1) + expect(capAttachments).toHaveLength(1) + expect(background.yielded).toContain(capAttachments[0]) + }) + + test('late background handoff after attachments still emits one final-turn cap', async () => { + dequeueAll() + let foregroundAbortController: AbortController + const queueingTool = buildTool({ + name: 'QueuePrompt', + inputSchema: z.object({}), + maxResultSizeChars: Infinity, + async description() { + return 'Queue a prompt during tool execution' + }, + async prompt() { + return '' + }, + async call() { + enqueue({ value: 'queued during tool execution', mode: 'prompt' }) + return { data: 'queued' } + }, + mapToolResultToToolResultBlockParam(content, toolUseID) { + return { + type: 'tool_result', + tool_use_id: toolUseID, + content: String(content), + } + }, + renderToolUseMessage() { + return null + }, + renderToolResultMessage() { + return null + }, + }) + const foregroundParams = makeParams( + async function* () { + yield createAssistantMessage({ + content: [ + { + type: 'tool_use', + id: 'toolu_queue_prompt_1', + name: 'QueuePrompt', + input: {}, + }, + ], + }) + }, + [queueingTool], + ) + foregroundParams.querySource = 'repl_main_thread' + foregroundAbortController = + foregroundParams.toolUseContext.abortController + const turnBudget = createQueryTurnBudget(1) + const foregroundGenerator = query({ ...foregroundParams, turnBudget }) + type ForegroundNext = Awaited< + ReturnType + > + type ForegroundYield = ForegroundNext extends IteratorResult + ? Y + : never + const foregroundYielded: ForegroundYield[] = [] + + try { + let foregroundReturned: Terminal | undefined + while (true) { + const next = await foregroundGenerator.next() + if (next.done) { + foregroundReturned = next.value + break + } + foregroundYielded.push(next.value) + if ( + next.value.type === 'attachment' && + next.value.attachment.type === 'queued_command' + ) { + // This yield is after the post-tool abort check and immediately + // before the terminal max-turn check that used to double-emit. + foregroundAbortController.abort('background') + } + } + + const background = await drain({ + ...makeParams(async function* () { + yield createAssistantMessage({ content: 'must not run' }) + }), + turnBudget, + }) + const capAttachments = [ + ...foregroundYielded, + ...background.yielded, + ].filter( + message => + message.type === 'attachment' && + message.attachment.type === 'max_turns_reached', + ) + + expect(foregroundReturned).toEqual({ reason: 'aborted_tools' }) + expect(background.returned).toEqual({ + reason: 'max_turns', + turnCount: 2, + }) + expect(capAttachments).toHaveLength(1) + expect(background.yielded).toContain(capAttachments[0]) + } finally { + dequeueAll() + } + }) + test('without a configured limit, tool use behavior is unchanged', async () => { echoCalls.length = 0 let modelCalls = 0 diff --git a/src/query/requestOnlyMessages.test.ts b/src/query/requestOnlyMessages.test.ts index 37da355769..268c9a8a57 100644 --- a/src/query/requestOnlyMessages.test.ts +++ b/src/query/requestOnlyMessages.test.ts @@ -292,10 +292,11 @@ test('keeps an explicit effort selection over ultrathink', async () => { expect(requestEffort).toBe('low') }) -test('scopes model-request lifecycle callbacks to callModel', async () => { +test('scopes model-request lifecycle callbacks to provider dispatch', async () => { const events: string[] = [] const params = baseParams( - async function* () { + async function* ({ options }) { + if (options.onProviderRequestStart?.() === false) return yield createAssistantMessage({ content: 'done' }) }, async () => ({ wasCompacted: false }), @@ -308,6 +309,29 @@ test('scopes model-request lifecycle callbacks to callModel', async () => { expect(events).toEqual(['start', 'end']) }) +test('closes the model-request lifecycle when its start callback aborts', async () => { + const events: string[] = [] + let providerDispatched = false + const params = baseParams( + async function* ({ options }) { + if (options.onProviderRequestStart?.() === false) return + providerDispatched = true + yield createAssistantMessage({ content: 'must not dispatch' }) + }, + async () => ({ wasCompacted: false }), + ) + params.onModelRequestStart = () => { + events.push('start') + params.toolUseContext.abortController.abort('background') + } + params.onModelRequestEnd = () => events.push('end') + + await collect(params) + + expect(providerDispatched).toBe(false) + expect(events).toEqual(['start', 'end']) +}) + for (const preserveCorrection of [false, true]) { test(`reapplies request-only context after ${ preserveCorrection ? 'suffix-preserving' : 'full' diff --git a/src/screens/REPL.tsx b/src/screens/REPL.tsx index c299f7cb44..a7afdd9348 100644 --- a/src/screens/REPL.tsx +++ b/src/screens/REPL.tsx @@ -8,6 +8,7 @@ import { count } from '../utils/array.js'; import { dirname, join } from 'path'; import { tmpdir } from 'os'; import figures from 'figures'; +import chalk from 'chalk'; // eslint-disable-next-line custom-rules/prefer-use-keybindings -- / n N Esc [ v are bare letters in transcript modal context, same class as g/G/j/k in ScrollKeybindingHandler import { useInput } from '../ink.js'; import { useSearchInput } from '../hooks/useSearchInput.js'; @@ -38,7 +39,7 @@ import { logForDebugging } from '../utils/debug.js'; import { QueryGuard } from '../utils/QueryGuard.js'; import { getQueryGuardOptionsFromEnv } from '../utils/queryGuardConfig.js'; import { QueryLifecycleOperationTracker, formatQueryLifecycleAbortSignalReason, formatQueryLifecycleLogMessage, getQueryTerminalReason, type QueryActiveOperationSnapshot, type QueryGuardTimeoutInfo, type QueryLifecycleContext, type QueryTerminalReason } from '../utils/queryLifecycle.js'; -import { resolveReplMaxTurns } from './replMaxTurns.js'; +import { claimBackgroundTurnBudget, canRestoreDeferredMaxTurnsCap, computeDeferredMaxTurnsCapForBackgroundHandoff, createForegroundTurnBudgetHandoff, getReplMaxTurnsWarning, releaseForegroundTurnBudget, resolveReplMaxTurnsForSession, shouldShowReplMaxTurnsUnlimitedWarning, shouldContinueBackgroundAfterForegroundQuery, waitForForegroundTurnBudgetSettlement, type ForegroundTurnBudgetHandoff } from './replMaxTurns.js'; import { createCombinedAbortSignal } from '../utils/combinedAbortSignal.js'; import { isEnvTruthy } from '../utils/envUtils.js'; import { formatTokens, truncateToWidth } from '../utils/format.js'; @@ -142,7 +143,8 @@ import { useQueueProcessor } from '../hooks/useQueueProcessor.js'; import { useMailboxBridge } from '../hooks/useMailboxBridge.js'; import { queryCheckpoint, logQueryProfileReport, clearQueryProfile } from '../utils/queryProfiler.js'; import type { Message as MessageType, UserMessage, ProgressMessage, HookResultMessage, PartialCompactDirection } from '../types/message.js'; -import { query } from '../query.js'; +import { query, type QueryTurnBudget } from '../query.js'; +import type { Terminal as QueryTerminal } from '../query/transitions.js'; import type { AutoCompactTrackingState } from '../services/compact/autoCompact.js'; import { mergeClients, useMergedClients } from '../hooks/useMergedClients.js'; import { getQuerySourceForREPL } from '../utils/promptCategory.js'; @@ -208,7 +210,7 @@ import { useIDEIntegration } from '../hooks/useIDEIntegration.js'; import exit from '../commands/exit/index.js'; import { ExitFlow } from '../components/ExitFlow.js'; import { getCurrentWorktreeSession } from '../utils/worktree.js'; -import { popAllEditable, enqueue, type SetAppState, getCommandQueue, getCommandQueueLength, removeByFilter } from '../utils/messageQueueManager.js'; +import { popAllEditable, enqueue, prepend, type SetAppState, getCommandQueue, getCommandQueueLength, removeByFilter } from '../utils/messageQueueManager.js'; import { useCommandQueue } from '../hooks/useCommandQueue.js'; import { SessionBackgroundHint } from '../components/SessionBackgroundHint.js'; import { startBackgroundSession } from '../tasks/LocalMainSessionTask.js'; @@ -289,6 +291,7 @@ import { useMessageActions, MessageActionsKeybindings, MessageActionsBar, type M import { setClipboard } from '../ink/termio/osc.js'; import type { ScrollBoxHandle } from '../ink/components/ScrollBox.js'; import { createAttachmentMessage, getQueuedCommandAttachments } from '../utils/attachments.js'; +import { dedupeQueuedTaskNotifications, filterClaimedTaskNotificationsForRestore, pendingCommandsForEmbeddedNotifications } from '../utils/taskNotificationIdentity.js'; // Stable empty array for hooks that accept MCPServerConnection[] — avoids // creating a new [] literal on every render in remote mode, which would @@ -572,7 +575,10 @@ function logQueryLifecycle(event: string, context: QueryLifecycleContext, extras logForDebugging(formatQueryLifecycleLogMessage(event, context, extras)); } // Default per-prompt cap for every local interactive REPL entrypoint. Headless -// and SDK callers retain their explicit maxTurns contracts. +// and SDK callers retain their explicit maxTurns contracts. Local interactive +// callers can raise the cap via --max-turns, OPENCLAUDE_MAX_TURNS / +// CLAUDE_CODE_MAX_TURNS, or `/config` → Max turns (interactive). +// Remote-backed sessions are not capped here. export type Props = { commands: Command[]; debug: boolean; @@ -655,8 +661,27 @@ export function REPL({ fallbackModel, maxTurns: maxTurnsProp }: Props): React.ReactNode { - const maxTurns = resolveReplMaxTurns(maxTurnsProp) + // Resolve at query time so `/config` changes apply on the next prompt + // without requiring a REPL remount. CLI prop still wins over env/config. const isRemoteSession = !!remoteSessionConfig; + const foregroundTurnBudgetRef = useRef(null); + const backgroundHandoffStartedRef = useRef(false); + const [backgroundHandoffPreparing, setBackgroundHandoffPreparing] = useState(false); + + useEffect(() => { + if ( + shouldShowReplMaxTurnsUnlimitedWarning(maxTurnsProp, { + isRemoteSession, + directConnectConfig, + sshSession, + }) + ) { + const warning = getReplMaxTurnsWarning(maxTurnsProp) + if (warning) { + process.stderr.write(chalk.yellow(`${warning}\n`)) + } + } + }, [maxTurnsProp, isRemoteSession, directConnectConfig, sshSession]); // Env-var gates hoisted to mount-time — isEnvTruthy does toLowerCase+trim+ // includes, and these were on the render path (hot during PageUp spam). @@ -1035,7 +1060,7 @@ export function REPL({ // Derived: any loading source active. Read-only — no setter. Local query // loading is driven by queryGuard (reserve/tryStart/end/cancelReservation), // external loading by setIsExternalLoading. - const isLoading = isQueryActive || isExternalLoading; + const isLoading = isQueryActive || isExternalLoading || backgroundHandoffPreparing; // Elapsed time is computed by SpinnerWithVerb from these refs on each // animation frame, avoiding a useInterval that re-renders the entire REPL. @@ -2831,60 +2856,162 @@ export function REPL({ // Session backgrounding (Ctrl+B to background/foreground) const handleBackgroundQuery = useCallback(() => { const backgroundSessionId = getSessionId(); - // Stop the foreground query so the background one takes over - abortController?.abort('background'); - // Aborting subagents may produce task-completed notifications. - // Clear task notifications so the queue processor doesn't immediately - // start a new foreground query; forward them to the background session. - const removedNotifications = removeByFilter(cmd => cmd.mode === 'task-notification'); - void (async () => { - const toolUseContext = getToolUseContext(messagesRef.current, [], new AbortController(), mainLoopModel); - const [defaultSystemPrompt, userContext, systemContext] = await Promise.all([getSystemPrompt(toolUseContext.options.tools, mainLoopModel, Array.from(toolPermissionContext.additionalWorkingDirectories.keys()), toolUseContext.options.mcpClients), getUserContext(), getSystemContext()]); - const systemPrompt = buildEffectiveSystemPrompt({ - mainThreadAgentDefinition, - toolUseContext, - customSystemPrompt, - defaultSystemPrompt, - appendSystemPrompt - }); - toolUseContext.renderedSystemPrompt = systemPrompt; - const notificationAttachments = await getQueuedCommandAttachments(removedNotifications).catch(() => []); - const notificationMessages = notificationAttachments.map(createAttachmentMessage); - - // Deduplicate: if the query loop already yielded a notification into - // messagesRef before we removed it from the queue, skip duplicates. - // We use prompt text for dedup because source_uuid is not set on - // task-notification QueuedCommands (enqueuePendingNotification callers - // don't pass uuid), so it would always be undefined. - const existingPrompts = new Set(); - for (const m of messagesRef.current) { - if (m.type === 'attachment' && m.attachment.type === 'queued_command' && m.attachment.commandMode === 'task-notification' && typeof m.attachment.prompt === 'string') { - existingPrompts.add(m.attachment.prompt); + const backgroundSessionTitle = terminalTitle; + setBackgroundHandoffPreparing(true); + // Transfer the exact per-prompt budget before aborting. Re-resolving the + // limit or copying a callback-maintained count here can reset or skew the + // cap while the foreground query is winding down. + const backgroundHandoff = claimBackgroundTurnBudget(foregroundTurnBudgetRef, backgroundHandoffStartedRef); + if (!backgroundHandoff) return; + const restoreDeferredMaxTurnsCap = () => { + const attemptRestore = (): boolean => { + const cap = backgroundHandoff.deferredMaxTurnsCap; + if (!cap || !canRestoreDeferredMaxTurnsCap(backgroundHandoff, messagesRef.current)) { + return false; } + backgroundHandoff.deferredMaxTurnsCap = undefined; + setMessages(prev => [...prev, createAttachmentMessage({ + type: 'max_turns_reached', + maxTurns: cap.maxTurns, + turnCount: cap.turnCount + })]); + return true; + }; + if (!attemptRestore()) { + void backgroundHandoff.settled.then(() => { + attemptRestore(); + }); } - const uniqueNotifications = notificationMessages.filter(m => m.attachment.type === 'queued_command' && (typeof m.attachment.prompt !== 'string' || !existingPrompts.has(m.attachment.prompt))); - startBackgroundSession({ - messages: [...messagesRef.current, ...uniqueNotifications], - queryParams: { - systemPrompt, - userContext, - systemContext, - canUseTool, + }; + const backgroundSession = startBackgroundSession({ + prepare: async backgroundAbortController => { + // The foreground owns transcript completion. Wait until its abort path + // has appended terminal tool results before snapshotting continuation + // state, while still honoring a task stop during that wait. + const shouldContinue = await waitForForegroundTurnBudgetSettlement(backgroundHandoff, backgroundAbortController.signal); + if (shouldContinue === null) { + throw backgroundAbortController.signal.reason; + } + if (!shouldContinue) return null; + + // The foreground is settled, but its QueryGuard will shortly release + // and allow a new prompt. Capture this continuation's transcript before + // any preparation await can observe that later foreground turn. + const settledMessages = [...messagesRef.current]; + backgroundHandoff.settledTranscriptTailUuid = + settledMessages.at(-1)?.uuid ?? null; + // Claim main-thread notifications only. Subagent-addressed entries keep + // their owner-scoped drain path (QueuedCommand.agentId); issue #2079 is + // about configurable interactive turn caps, not queue isolation. + const pendingNotifications = removeByFilter( + cmd => + cmd.mode === 'task-notification' && cmd.agentId === undefined, + ); + let restorableNotifications = filterClaimedTaskNotificationsForRestore( + pendingNotifications, + settledMessages, + ); + let notificationOwnershipActive = true; + const restoreNotificationsIfUnsent = () => { + if (!notificationOwnershipActive) return; + notificationOwnershipActive = false; + if (restorableNotifications.length > 0) { + prepend(restorableNotifications); + } + }; + try { + const toolUseContext = getToolUseContext(settledMessages, [], backgroundAbortController, mainLoopModel); + const [defaultSystemPrompt, userContext, systemContext] = await Promise.all([getSystemPrompt(toolUseContext.options.tools, mainLoopModel, Array.from(toolPermissionContext.additionalWorkingDirectories.keys()), toolUseContext.options.mcpClients), getUserContext(), getSystemContext()]).catch(error => { + restoreNotificationsIfUnsent(); + throw error; + }); + if (backgroundAbortController.signal.aborted) { + restoreNotificationsIfUnsent(); + throw backgroundAbortController.signal.reason; + } + const systemPrompt = buildEffectiveSystemPrompt({ + mainThreadAgentDefinition, toolUseContext, - fallbackModel, - maxTurns, - querySource: getQuerySourceForREPL(), - autoCompactTracking: getAutoCompactTrackingForSession(backgroundSessionId), - onAutoCompactTrackingChange: tracking => { - setAutoCompactTrackingForSession(backgroundSessionId, tracking); + customSystemPrompt, + defaultSystemPrompt, + appendSystemPrompt + }); + toolUseContext.renderedSystemPrompt = systemPrompt; + const notificationAttachments = await getQueuedCommandAttachments(pendingNotifications).catch(error => { + restoreNotificationsIfUnsent(); + throw error; + }); + if (backgroundAbortController.signal.aborted) { + restoreNotificationsIfUnsent(); + throw backgroundAbortController.signal.reason; + } + const notificationMessages = notificationAttachments.map(createAttachmentMessage); + + // Deduplicate against settled transcript keys and within the claimed batch. + const uniqueNotifications = dedupeQueuedTaskNotifications( + settledMessages, + notificationMessages, + ); + restorableNotifications = pendingCommandsForEmbeddedNotifications( + pendingNotifications, + uniqueNotifications, + ); + return { + messages: [...settledMessages, ...uniqueNotifications], + restoreNotificationsIfUnsent, + commitNotificationOwnership: () => { + notificationOwnershipActive = false; + }, + queryParams: { + systemPrompt, + userContext, + systemContext, + canUseTool, + toolUseContext, + fallbackModel, + turnBudget: backgroundHandoff.budget, + querySource: getQuerySourceForREPL(), + autoCompactTracking: getAutoCompactTrackingForSession(backgroundSessionId), + onAutoCompactTrackingChange: tracking => { + setAutoCompactTrackingForSession(backgroundSessionId, tracking); + } } - }, - description: terminalTitle, - setAppState, - agentDefinition: mainThreadAgentDefinition - }); - })(); - }, [abortController, mainLoopModel, toolPermissionContext, mainThreadAgentDefinition, getToolUseContext, customSystemPrompt, appendSystemPrompt, canUseTool, setAppState, getAutoCompactTrackingForSession, setAutoCompactTrackingForSession, fallbackModel, maxTurns]); + }; + } catch (error) { + restoreNotificationsIfUnsent(); + throw error; + } + }, + description: backgroundSessionTitle, + setAppState, + agentDefinition: mainThreadAgentDefinition, + onPreparationError: () => { + restoreDeferredMaxTurnsCap(); + addNotification({ + key: 'background-session-start-failed', + text: 'Could not start the background session. The current request was cancelled.', + priority: 'high', + }); + }, + onContinuationCancelled: restoreDeferredMaxTurnsCap, + onRegistered: controller => { + setAbortController(current => + current === controller ? null : current, + ); + }, + onSettled: controller => { + setBackgroundHandoffPreparing(false); + setAbortController(current => + current === controller ? null : current, + ); + }, + }); + // The task is intentionally published only after the foreground settles, + // but its controller must be reachable during preparation so Escape can + // cancel the handoff before it dispatches a provider request. + setAbortController(backgroundSession.abortController); + abortController?.abort('background'); + }, [abortController, mainLoopModel, toolPermissionContext, mainThreadAgentDefinition, getToolUseContext, customSystemPrompt, appendSystemPrompt, canUseTool, setAppState, getAutoCompactTrackingForSession, setAutoCompactTrackingForSession, fallbackModel, setAbortController, addNotification, terminalTitle]); const { handleBackgroundSession } = useSessionBackgrounding({ @@ -2959,7 +3086,7 @@ export function REPL({ void removeTranscriptMessage(tombstonedMessage.uuid); }, setStreamingThinking, undefined, onStreamingText); }, [setMessages, setResponseLength, setStreamMode, setStreamingToolUses, setStreamingThinking, onStreamingText]); - const onQueryImpl = useCallback(async (messagesIncludingNewMessages: MessageType[], newMessages: MessageType[], abortController: AbortController, shouldQuery: boolean, additionalAllowedTools: string[], mainLoopModelParam: string, queryGeneration: number, effort?: EffortValue, queryLifecycle?: QueryLifecycleOperationTracker, requestOnlyMessages: MessageType[] = [], interruptionCorrectionQueryId?: string, onModelRequestStart?: () => void) => { + const onQueryImpl = useCallback(async (messagesIncludingNewMessages: MessageType[], newMessages: MessageType[], abortController: AbortController, shouldQuery: boolean, additionalAllowedTools: string[], mainLoopModelParam: string, queryGeneration: number, turnBudget: QueryTurnBudget, effort?: EffortValue, queryLifecycle?: QueryLifecycleOperationTracker, requestOnlyMessages: MessageType[] = [], interruptionCorrectionQueryId?: string, onModelRequestStart?: () => void): Promise => { // Prepare IDE integration for new prompt. Read mcpClients fresh from // store — useManageMCPConnections may have populated it since the // render that captured this closure (same pattern as computeTools). @@ -3096,7 +3223,7 @@ export function REPL({ resetTurnToolDuration(); resetTurnClassifierDuration(); let expectedAutoCompactTracking = queryAutoCompactTracking; - for await (const event of query({ + const queryGenerator = query({ messages: messagesIncludingNewMessages, requestOnlyMessages, onModelRequestStart: interruptionCorrectionQueryId @@ -3116,16 +3243,30 @@ export function REPL({ toolUseContext, querySource: getQuerySourceForREPL(), fallbackModel, - maxTurns, + turnBudget, autoCompactTracking: queryAutoCompactTracking, onAutoCompactTrackingChange: tracking => { if (setAutoCompactTrackingForSessionIfUnchanged(querySessionId, expectedAutoCompactTracking, tracking)) { expectedAutoCompactTracking = tracking; } } - })) { - queryGuard.registerActivity(`query_event:${event.type}`, queryGeneration); - onQueryEvent(event); + }); + let queryTerminal: QueryTerminal; + let generatorDone = false; + try { + while (true) { + const next = await queryGenerator.next(); + if (next.done) { + generatorDone = true; + queryTerminal = next.value; + break; + } + const event = next.value; + queryGuard.registerActivity(`query_event:${event.type}`, queryGeneration); + onQueryEvent(event); + } + } finally { + if (!generatorDone) await queryGenerator.return(undefined as never); } if (isBuddyEnabled()) { void fireCompanionObserver(messagesRef.current, reaction => setAppState(prev => prev.companionReaction === reaction ? prev : { @@ -3142,7 +3283,8 @@ export function REPL({ // Signal that a query turn has completed successfully await onTurnComplete?.(messagesRef.current); - }, [initialMcpClients, resetLoadingState, getToolUseContext, toolPermissionContext, setAppState, customSystemPrompt, onTurnComplete, appendSystemPrompt, canUseTool, mainThreadAgentDefinition, onQueryEvent, sessionTitle, titleDisabled, maxTurns, getAutoCompactTrackingForSession, setAutoCompactTrackingForSession, setAutoCompactTrackingForSessionIfUnchanged, queryGuard, interruptionCorrectionTracker]); + return queryTerminal; + }, [initialMcpClients, resetLoadingState, getToolUseContext, toolPermissionContext, setAppState, customSystemPrompt, onTurnComplete, appendSystemPrompt, canUseTool, mainThreadAgentDefinition, onQueryEvent, sessionTitle, titleDisabled, getAutoCompactTrackingForSession, setAutoCompactTrackingForSession, setAutoCompactTrackingForSessionIfUnchanged, queryGuard, interruptionCorrectionTracker]); const onQuery = useCallback(async (newMessages: MessageType[], abortController: AbortController, shouldQuery: boolean, additionalAllowedTools: string[], mainLoopModelParam: string, onBeforeQueryCallback?: (input: string, newMessages: MessageType[]) => Promise, input?: string, effort?: EffortValue, isInterruptionCorrectionEligible = false, onModelRequestStart?: () => void): Promise => { // If this is a teammate, mark them as active when starting a turn if (isAgentSwarmsEnabled()) { @@ -3159,6 +3301,16 @@ export function REPL({ // Returns null if already running — no separate check-then-set. const lifecycleTracker = queryLifecycleTrackerRef.current; const querySource = getQuerySourceForREPL(); + if (backgroundHandoffPreparing) { + logEvent('tengu_concurrent_onquery_detected', {}); + newMessages.filter((m): m is UserMessage => m.type === 'user' && !m.isMeta).map(_ => getContentText(_.message.content)).filter(_ => _ !== null).forEach((msg, i) => { + enqueue(buildConcurrentRequeuedPrompt(msg, isInterruptionCorrectionEligible)); + if (i === 0) { + logEvent('tengu_concurrent_onquery_enqueued', {}); + } + }); + return false; + } const startResult = queryGuard.tryStart({ queryId: randomUUID(), querySource, @@ -3181,10 +3333,21 @@ export function REPL({ } lifecycleTracker.clear(); const thisGeneration = startResult.generation; + backgroundHandoffStartedRef.current = false; + const turnBudgetHandoff = createForegroundTurnBudgetHandoff( + resolveReplMaxTurnsForSession(maxTurnsProp, { + isRemoteSession, + directConnectConfig, + sshSession, + }), + ); + const turnBudget = turnBudgetHandoff.budget; + foregroundTurnBudgetRef.current = turnBudgetHandoff; const queryContext = startResult.context; logQueryLifecycle('start', queryContext); logQueryLifecycle('guard_start', queryContext); let didThrow = false; + let queryTerminal: QueryTerminal | undefined; let preflightVetoed = false; let modelTurnStarted = false; let hasInterruptionCorrectionRequestOnlyMessage = false; @@ -3240,7 +3403,7 @@ export function REPL({ } if (!preflightVetoed) { modelTurnStarted = true; - await onQueryImpl(latestMessages, persistentNewMessages, abortController, shouldQuery, additionalAllowedTools, mainLoopModelParam, thisGeneration, effort, lifecycleTracker, requestOnlyMessages, isInterruptionCorrectionEligible ? queryContext.queryId : undefined, onModelRequestStart); + queryTerminal = await onQueryImpl(latestMessages, persistentNewMessages, abortController, shouldQuery, additionalAllowedTools, mainLoopModelParam, thisGeneration, turnBudget, effort, lifecycleTracker, requestOnlyMessages, isInterruptionCorrectionEligible ? queryContext.queryId : undefined, onModelRequestStart); } if (preflightVetoed) { return false; @@ -3254,6 +3417,25 @@ export function REPL({ } throw error; } finally { + // The ref is only an ownership marker for the currently foregrounded + // prompt. A background handoff already captured the budget object, so + // clear the ref on every terminal path without disturbing a newer query. + const shouldContinueBackground = + shouldContinueBackgroundAfterForegroundQuery({ + didThrow, + preflightVetoed, + abortReason: abortController.signal.reason, + queryTerminal, + }); + if (shouldContinueBackground) { + const deferredCap = computeDeferredMaxTurnsCapForBackgroundHandoff(abortController.signal.reason, queryTerminal, turnBudget.maxTurns, turnBudget.turnsStarted); + if (deferredCap) { + turnBudgetHandoff.deferredMaxTurnsCap = deferredCap; + turnBudgetHandoff.settledTranscriptTailUuid = + messagesRef.current.at(-1)?.uuid ?? null; + } + } + releaseForegroundTurnBudget(foregroundTurnBudgetRef, backgroundHandoffStartedRef, turnBudgetHandoff, shouldContinueBackground); // A provider response can hand off to tools before the assistant turn // finishes. Keep correction ownership through that work (and retries), // then clear it only when this query reaches its terminal cleanup. @@ -3368,7 +3550,9 @@ export function REPL({ // reads false at the idle prompt. Without this, the stale non-aborted // controller makes ctrl+c fire onCancel() (aborting nothing) instead of // propagating to the double-press exit flow. - setAbortController(null); + setAbortController(current => + current === abortController ? null : current, + ); } else { const guardCompletedContext = queryGuard.lastContext; if ((guardCompletedContext?.terminalReason === 'query-timeout' || guardCompletedContext?.terminalReason === 'hard-max-query-timeout') && guardCompletedContext.queryGeneration === thisGeneration) { @@ -3412,7 +3596,7 @@ export function REPL({ } } } - }, [onQueryImpl, setAppState, resetLoadingState, queryGuard, mrOnBeforeQuery, mrOnTurnComplete]); + }, [onQueryImpl, setAppState, resetLoadingState, queryGuard, mrOnBeforeQuery, mrOnTurnComplete, maxTurnsProp, isRemoteSession, directConnectConfig, sshSession, backgroundHandoffPreparing]); // Handle initial message (from CLI args or plan mode exit with context clear) // This effect runs when isLoading becomes false and there's a pending message diff --git a/src/screens/replMaxTurns.ts b/src/screens/replMaxTurns.ts index c60c8ab86a..b282847e6d 100644 --- a/src/screens/replMaxTurns.ts +++ b/src/screens/replMaxTurns.ts @@ -1,5 +1,189 @@ -export const DEFAULT_REPL_MAX_TURNS = 50 +import { + createQueryTurnBudget, + type QueryTurnBudget, +} from '../query.js' +import type { Terminal as QueryTerminal } from '../query/transitions.js' +import { normalizeAbortReason } from '../utils/abortReasons.js' +import { getReplMaxTurnsWarning, resolveReplMaxTurns } from '../utils/replMaxTurns.js' -export function resolveReplMaxTurns(maxTurns?: number): number { - return maxTurns ?? DEFAULT_REPL_MAX_TURNS +export { + DEFAULT_REPL_MAX_TURNS, + getReplMaxTurnsWarning, + MAX_TURNS_CLI_DESCRIPTION, + REPL_MAX_TURNS_OPTIONS, + normalizeReplMaxTurns, + resolveReplMaxTurns, +} from '../utils/replMaxTurns.js' + +export function isLocalInteractiveMaxTurnsSession(session: { + isRemoteSession: boolean + directConnectConfig: unknown + sshSession: unknown +}): boolean { + return ( + !session.isRemoteSession && + !session.directConnectConfig && + !session.sshSession + ) +} + +export function shouldShowReplMaxTurnsUnlimitedWarning( + maxTurns: number | undefined, + session: { + isRemoteSession: boolean + directConnectConfig: unknown + sshSession: unknown + }, +): boolean { + const warning = getReplMaxTurnsWarning(maxTurns) + return warning !== undefined && isLocalInteractiveMaxTurnsSession(session) +} + +export function shouldContinueBackgroundAfterForegroundQuery({ + didThrow, + preflightVetoed, + abortReason, + queryTerminal, +}: { + didThrow: boolean + preflightVetoed: boolean + abortReason: unknown + queryTerminal: QueryTerminal | undefined +}): boolean { + return ( + !didThrow && + !preflightVetoed && + abortReason === 'background' && + (queryTerminal?.reason === 'aborted_streaming' || + queryTerminal?.reason === 'aborted_tools') + ) +} + +type MutableRef = { current: T } + +export type DeferredMaxTurnsCap = { + maxTurns: number + turnCount: number +} + +export type ForegroundTurnBudgetHandoff = { + budget: QueryTurnBudget + settled: Promise + settle: (shouldContinue: boolean) => void + /** Cap suppressed on foreground `background` abort; restore if handoff is cancelled. */ + deferredMaxTurnsCap?: DeferredMaxTurnsCap + /** + * Last settled transcript message before a cancelled handoff may restore a + * deferred cap. When the live tail uuid differs, a newer prompt owns the view. + */ + settledTranscriptTailUuid?: string | null +} + +export function canRestoreDeferredMaxTurnsCap( + handoff: ForegroundTurnBudgetHandoff, + currentMessages: readonly { uuid: string }[], +): boolean { + if (!handoff.deferredMaxTurnsCap) return false + const anchor = handoff.settledTranscriptTailUuid + if (anchor === undefined) return true + if (anchor === null) return currentMessages.length === 0 + return currentMessages.at(-1)?.uuid === anchor +} + +export function resolveReplMaxTurnsForSession( + maxTurns: number | undefined, + session: { + isRemoteSession: boolean + directConnectConfig: unknown + sshSession: unknown + }, +): number | undefined { + if (!isLocalInteractiveMaxTurnsSession(session)) { + return undefined + } + return resolveReplMaxTurns(maxTurns) +} + +export function computeDeferredMaxTurnsCapForBackgroundHandoff( + abortReason: unknown, + queryTerminal: QueryTerminal | undefined, + maxTurns: number | undefined, + turnsStarted: number, +): DeferredMaxTurnsCap | undefined { + if ( + normalizeAbortReason(abortReason) !== 'background' || + queryTerminal?.reason !== 'aborted_tools' || + maxTurns === undefined + ) { + return undefined + } + const turnCount = turnsStarted + 1 + if (turnCount <= maxTurns) { + return undefined + } + return { maxTurns, turnCount } +} + +export function createForegroundTurnBudgetHandoff( + maxTurns?: number, +): ForegroundTurnBudgetHandoff { + let resolveSettled!: (shouldContinue: boolean) => void + let isSettled = false + const settled = new Promise(resolve => { + resolveSettled = resolve + }) + return { + budget: createQueryTurnBudget(maxTurns), + settled, + settle: shouldContinue => { + if (isSettled) return + isSettled = true + resolveSettled(shouldContinue) + }, + } +} + +export async function waitForForegroundTurnBudgetSettlement( + handoff: ForegroundTurnBudgetHandoff, + signal: AbortSignal, +): Promise { + if (signal.aborted) return null + + let resolveAborted!: () => void + const aborted = new Promise(resolve => { + resolveAborted = resolve + }) + const onAbort = () => resolveAborted() + signal.addEventListener('abort', onAbort, { once: true }) + try { + return await Promise.race([ + handoff.settled, + aborted.then(() => null), + ]) + } finally { + signal.removeEventListener('abort', onAbort) + } +} + +export function claimBackgroundTurnBudget( + budgetRef: MutableRef, + handoffStartedRef: MutableRef, +): ForegroundTurnBudgetHandoff | null { + if (!budgetRef.current || handoffStartedRef.current) return null + handoffStartedRef.current = true + return budgetRef.current +} + +export function releaseForegroundTurnBudget( + budgetRef: MutableRef, + handoffStartedRef: MutableRef, + ownedHandoff: ForegroundTurnBudgetHandoff, + shouldContinue: boolean, +): void { + // Always release waiters for this prompt, even if a newer prompt replaced + // the foreground ref before the stale finally ran. + ownedHandoff.settle(shouldContinue) + if (budgetRef.current !== ownedHandoff) return + budgetRef.current = null + handoffStartedRef.current = false } diff --git a/src/screens/replMaxTurnsProp.test.ts b/src/screens/replMaxTurnsProp.test.ts index c125ba721d..5ae0baea8f 100644 --- a/src/screens/replMaxTurnsProp.test.ts +++ b/src/screens/replMaxTurnsProp.test.ts @@ -1,52 +1,481 @@ -import { describe, expect, test } from 'bun:test' -import { readFileSync } from 'node:fs' -import { join } from 'node:path' -import { DEFAULT_REPL_MAX_TURNS, resolveReplMaxTurns } from './replMaxTurns.js' +import { afterEach, beforeEach, describe, expect, spyOn, test } from 'bun:test' +import { + Command as CommanderCommand, + InvalidArgumentError, + Option, +} from '@commander-js/extra-typings' +import { + canRestoreDeferredMaxTurnsCap, + claimBackgroundTurnBudget, + computeDeferredMaxTurnsCapForBackgroundHandoff, + createForegroundTurnBudgetHandoff, + isLocalInteractiveMaxTurnsSession, + releaseForegroundTurnBudget, + resolveReplMaxTurnsForSession, + shouldContinueBackgroundAfterForegroundQuery, + shouldShowReplMaxTurnsUnlimitedWarning, + waitForForegroundTurnBudgetSettlement, +} from './replMaxTurns.js' +import { createQueryTurnBudget } from '../query.js' +import { + DEFAULT_GLOBAL_CONFIG, + GLOBAL_CONFIG_KEYS, + getGlobalConfig, + isGlobalConfigKey, + saveGlobalConfig, +} from '../utils/config.js' +import { + DEFAULT_REPL_MAX_TURNS, + getReplMaxTurnsWarning, + MAX_TURNS_CLI_DESCRIPTION, + MAX_TURNS_UNLIMITED_WARNING, + normalizeReplMaxTurns, + parseMaxTurnsCommanderArgument, + parseMaxTurnsCli, + REPL_MAX_TURNS_OPTIONS, + resolveReplMaxTurns, +} from '../utils/replMaxTurns.js' +import * as debug from '../utils/debug.js' -const screenDir = import.meta.dirname +const ENV_KEYS = ['OPENCLAUDE_MAX_TURNS', 'CLAUDE_CODE_MAX_TURNS'] as const +const savedEnv: Partial> = + {} +const savedReplMaxTurns = getGlobalConfig().replMaxTurns -function readScreen(name: string): string { - return readFileSync(join(screenDir, name), 'utf8') +function clearTurnEnv(): void { + for (const key of ENV_KEYS) { + delete process.env[key] + } +} + +function setReplMaxTurnsConfig(value: number | undefined): void { + saveGlobalConfig(current => ({ + ...current, + replMaxTurns: value, + })) +} + +function createMaxTurnsCliProgram(): CommanderCommand { + const program = new CommanderCommand() + program + .name('openclaude') + .exitOverride() + .addOption( + new Option('--max-turns ', MAX_TURNS_CLI_DESCRIPTION).argParser( + value => { + return parseMaxTurnsCommanderArgument(value) + }, + ), + ) + .action(() => {}) + return program } -function objectBody(source: string, marker: RegExp): string { - const match = source.match(marker) - expect(match).not.toBeNull() - const start = match!.index! + match![0].length - 1 - let depth = 0 - for (let index = start; index < source.length; index++) { - if (source[index] === '{') depth++ - if (source[index] === '}') { - depth-- - if (depth === 0) return source.slice(start, index + 1) +beforeEach(() => { + clearTurnEnv() +}) + +afterEach(() => { + for (const key of ENV_KEYS) { + const previous = savedEnv[key] + if (previous === undefined) { + delete process.env[key] + } else { + process.env[key] = previous } } - throw new Error(`Unclosed object after ${marker}`) + setReplMaxTurnsConfig(savedReplMaxTurns) +}) + +for (const key of ENV_KEYS) { + savedEnv[key] = process.env[key] } describe('interactive REPL max-turn cap', () => { test('supplies the local interactive default at runtime', () => { + clearTurnEnv() + setReplMaxTurnsConfig(undefined) expect(DEFAULT_REPL_MAX_TURNS).toBe(50) expect(resolveReplMaxTurns()).toBe(50) }) test('preserves an explicit interactive cap at runtime', () => { + clearTurnEnv() expect(resolveReplMaxTurns(7)).toBe(7) }) - test('passes the resolved cap to foreground and background queries', () => { - const source = readScreen('REPL.tsx') - const foreground = objectBody(source, /for await \(const event of query\(\{/) - const background = objectBody(source, /queryParams:\s*\{/) + test('honors OPENCLAUDE_MAX_TURNS when no explicit cap is passed', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = '200' + expect(resolveReplMaxTurns()).toBe(200) + }) + + test('falls back to CLAUDE_CODE_MAX_TURNS when OPENCLAUDE_MAX_TURNS is unset', () => { + clearTurnEnv() + process.env.CLAUDE_CODE_MAX_TURNS = '125' + expect(resolveReplMaxTurns()).toBe(125) + }) + + test('prefers OPENCLAUDE_MAX_TURNS over CLAUDE_CODE_MAX_TURNS', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = '90' + process.env.CLAUDE_CODE_MAX_TURNS = '30' + expect(resolveReplMaxTurns()).toBe(90) + }) + + test('invalid OPENCLAUDE_MAX_TURNS does not fall through to legacy', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = 'nope' + process.env.CLAUDE_CODE_MAX_TURNS = '125' + expect(resolveReplMaxTurns()).toBe(DEFAULT_REPL_MAX_TURNS) + }) + + test('does not include an invalid environment value in debug logs', () => { + clearTurnEnv() + const invalidValue = 'private-value-that-must-not-be-logged' + process.env.OPENCLAUDE_MAX_TURNS = invalidValue + const logSpy = spyOn(debug, 'logForDebugging').mockImplementation(() => {}) + + try { + expect(resolveReplMaxTurns()).toBe(DEFAULT_REPL_MAX_TURNS) + expect(logSpy).toHaveBeenCalledTimes(1) + const logged = logSpy.mock.calls.flat().join(' ') + expect(logged).toContain('OPENCLAUDE_MAX_TURNS has an invalid value') + expect(logged).not.toContain(invalidValue) + } finally { + logSpy.mockRestore() + } + }) + + test('explicit CLI cap wins over environment overrides', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = '200' + expect(resolveReplMaxTurns(80)).toBe(80) + }) + + test('honors /config replMaxTurns when CLI and env are unset', () => { + clearTurnEnv() + setReplMaxTurnsConfig(200) + expect(resolveReplMaxTurns()).toBe(200) + }) + + test('env wins over /config replMaxTurns', () => { + clearTurnEnv() + setReplMaxTurnsConfig(200) + process.env.OPENCLAUDE_MAX_TURNS = '80' + expect(resolveReplMaxTurns()).toBe(80) + }) + + test('CLI wins over /config replMaxTurns', () => { + clearTurnEnv() + setReplMaxTurnsConfig(200) + expect(resolveReplMaxTurns(90)).toBe(90) + }) + + test('treats an explicit CLI zero as unlimited and invalid values as the default', () => { + clearTurnEnv() + setReplMaxTurnsConfig(undefined) + process.env.OPENCLAUDE_MAX_TURNS = 'nope' + expect(resolveReplMaxTurns()).toBe(DEFAULT_REPL_MAX_TURNS) + clearTurnEnv() + expect(resolveReplMaxTurns(0)).toBeUndefined() + expect(resolveReplMaxTurns(-3)).toBe(DEFAULT_REPL_MAX_TURNS) + expect(resolveReplMaxTurns(Number.NaN)).toBe(DEFAULT_REPL_MAX_TURNS) + expect(resolveReplMaxTurns(2.5)).toBe(DEFAULT_REPL_MAX_TURNS) + }) + + test('warns when the CLI explicitly disables the turn limit', () => { + expect(getReplMaxTurnsWarning(0)).toBe(MAX_TURNS_UNLIMITED_WARNING) + expect(getReplMaxTurnsWarning(50)).toBeUndefined() + expect(getReplMaxTurnsWarning()).toBeUndefined() + }) + + test('emits the unlimited warning only from a local REPL', () => { + expect( + shouldShowReplMaxTurnsUnlimitedWarning(0, { + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: undefined, + }), + ).toBe(true) + expect( + shouldShowReplMaxTurnsUnlimitedWarning(0, { + isRemoteSession: true, + directConnectConfig: undefined, + sshSession: undefined, + }), + ).toBe(false) + expect( + shouldShowReplMaxTurnsUnlimitedWarning(0, { + isRemoteSession: false, + directConnectConfig: {}, + sshSession: undefined, + }), + ).toBe(false) + expect( + shouldShowReplMaxTurnsUnlimitedWarning(0, { + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: {}, + }), + ).toBe(false) + expect( + shouldShowReplMaxTurnsUnlimitedWarning(50, { + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: undefined, + }), + ).toBe(false) + expect(isLocalInteractiveMaxTurnsSession({ + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: undefined, + })).toBe(true) + }) + + test('normalizeReplMaxTurns matches /config picker persistence', () => { + expect(normalizeReplMaxTurns(200)).toBe(200) + expect(normalizeReplMaxTurns('500')).toBe(500) + expect(normalizeReplMaxTurns(0)).toBe(DEFAULT_REPL_MAX_TURNS) + expect(normalizeReplMaxTurns('nope')).toBe(DEFAULT_REPL_MAX_TURNS) + }) + + test('replMaxTurns is registered for /config', () => { + expect(GLOBAL_CONFIG_KEYS).toContain('replMaxTurns') + expect(isGlobalConfigKey('replMaxTurns')).toBe(true) + expect(DEFAULT_GLOBAL_CONFIG.replMaxTurns).toBeUndefined() + expect(REPL_MAX_TURNS_OPTIONS).toEqual([50, 100, 200, 500]) + }) + + test('headless --max-turns 0 stays distinct from interactive unlimited resolution', () => { + expect(parseMaxTurnsCli('0')).toBe(0) + expect(resolveReplMaxTurns(0)).toBeUndefined() + const headlessTurnBudget = createQueryTurnBudget(parseMaxTurnsCli('0')) + expect(headlessTurnBudget.maxTurns).toBe(0) + // Headless passes the CLI value through; falsy maxTurns disables turn-cap guards. + expect(Boolean(headlessTurnBudget.maxTurns)).toBe(false) + const help = MAX_TURNS_CLI_DESCRIPTION.toLowerCase() + expect(help).toContain('local interactive mode') + expect(help).toContain('--print mode') + }) + + test('background budget handoff preserves identity, settlement, and one-shot ownership', async () => { + const handoff = createForegroundTurnBudgetHandoff(50) + const budgetRef = { current: handoff } + const handoffStartedRef = { current: false } + + expect( + claimBackgroundTurnBudget(budgetRef, handoffStartedRef), + ).toBe(handoff) + expect(handoff.budget.maxTurns).toBe(50) + expect( + claimBackgroundTurnBudget(budgetRef, handoffStartedRef), + ).toBeNull() + + const newerHandoff = createForegroundTurnBudgetHandoff(100) + budgetRef.current = newerHandoff + handoffStartedRef.current = false + releaseForegroundTurnBudget(budgetRef, handoffStartedRef, handoff, true) + await expect(handoff.settled).resolves.toBe(true) + expect(budgetRef.current).toBe(newerHandoff) + + releaseForegroundTurnBudget( + budgetRef, + handoffStartedRef, + newerHandoff, + false, + ) + await expect(newerHandoff.settled).resolves.toBe(false) + expect(budgetRef.current).toBeNull() + expect(handoffStartedRef.current).toBe(false) + }) + + test('a stopped background task does not remain blocked on foreground settlement', async () => { + const handoff = createForegroundTurnBudgetHandoff(50) + const abortController = new AbortController() + const wait = waitForForegroundTurnBudgetSettlement( + handoff, + abortController.signal, + ) + + abortController.abort('task stopped') - expect(foreground).toContain('maxTurns,') - expect(background).toContain('maxTurns,') + await expect(wait).resolves.toBeNull() }) - test('passes the cap from the resume selector into REPL', () => { - const source = readScreen('ResumeConversation.tsx') - const repl = source.slice(source.indexOf('', source.indexOf(' { + expect( + computeDeferredMaxTurnsCapForBackgroundHandoff( + 'background', + { reason: 'aborted_tools' }, + 1, + 1, + ), + ).toEqual({ maxTurns: 1, turnCount: 2 }) + expect( + computeDeferredMaxTurnsCapForBackgroundHandoff( + 'background', + { reason: 'aborted_streaming' }, + 1, + 1, + ), + ).toBeUndefined() + expect( + computeDeferredMaxTurnsCapForBackgroundHandoff( + 'user', + { reason: 'aborted_tools' }, + 1, + 1, + ), + ).toBeUndefined() + }) + + test('skips deferred max-turn restoration when a newer prompt owns the transcript', () => { + const handoff = createForegroundTurnBudgetHandoff(1) + handoff.deferredMaxTurnsCap = { maxTurns: 1, turnCount: 2 } + handoff.settledTranscriptTailUuid = 'prior-tail' + + expect( + canRestoreDeferredMaxTurnsCap(handoff, [ + { uuid: 'prior-tail' }, + ]), + ).toBe(true) + expect( + canRestoreDeferredMaxTurnsCap(handoff, [ + { uuid: 'prior-tail' }, + { uuid: 'new-user-turn' }, + ]), + ).toBe(false) + }) + + test('does not apply the local interactive cap in remote-backed sessions', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = '80' + expect( + resolveReplMaxTurnsForSession(undefined, { + isRemoteSession: true, + directConnectConfig: undefined, + sshSession: undefined, + }), + ).toBeUndefined() + expect( + resolveReplMaxTurnsForSession(90, { + isRemoteSession: false, + directConnectConfig: {}, + sshSession: undefined, + }), + ).toBeUndefined() + expect( + resolveReplMaxTurnsForSession(90, { + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: {}, + }), + ).toBeUndefined() + expect( + resolveReplMaxTurnsForSession(90, { + isRemoteSession: false, + directConnectConfig: undefined, + sshSession: undefined, + }), + ).toBe(90) + }) + + test('does not continue a background handoff after the foreground query throws', () => { + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: true, + preflightVetoed: false, + abortReason: 'background', + queryTerminal: { reason: 'aborted_streaming' }, + }), + ).toBe(false) + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: false, + preflightVetoed: true, + abortReason: 'background', + queryTerminal: undefined, + }), + ).toBe(false) + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: false, + preflightVetoed: false, + abortReason: 'background', + queryTerminal: { reason: 'aborted_streaming' }, + }), + ).toBe(true) + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: false, + preflightVetoed: false, + abortReason: 'background', + queryTerminal: { reason: 'aborted_tools' }, + }), + ).toBe(true) + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: false, + preflightVetoed: false, + abortReason: 'background', + queryTerminal: undefined, + }), + ).toBe(false) + expect( + shouldContinueBackgroundAfterForegroundQuery({ + didThrow: false, + preflightVetoed: false, + abortReason: 'user', + queryTerminal: { reason: 'aborted_streaming' }, + }), + ).toBe(false) + }) + + test('Commander --max-turns help scopes the interactive cap to local query loops', () => { + // Commander wraps long option help across lines; collapse whitespace. + const help = createMaxTurnsCliProgram() + .helpInformation() + .toLowerCase() + .replace(/\s+/g, ' ') + expect(help).toContain('--max-turns') + expect(help).toContain('local interactive') + expect(help).toContain('remote-backed') + expect(help).not.toContain('only works with --print') + }) + + test('Commander --max-turns parses into the value the local REPL resolves', async () => { + clearTurnEnv() + setReplMaxTurnsConfig(50) + const program = createMaxTurnsCliProgram() + await program.parseAsync(['node', 'openclaude', '--max-turns', '200'], { + from: 'node', + }) + const parsed = program.getOptionValue('maxTurns') as number + expect(parsed).toBe(200) + // Same handoff the interactive session uses: CLI option → resolveReplMaxTurns. + expect(resolveReplMaxTurns(parsed)).toBe(200) + }) + + test('Commander --max-turns rejects values that could disable the cap accidentally', () => { + expect(parseMaxTurnsCli('0')).toBe(0) + expect(parseMaxTurnsCli('200')).toBe(200) + for (const invalid of ['', ' ', 'nope', '-3', '2.5', 'Infinity']) { + expect(() => parseMaxTurnsCli(invalid)).toThrow( + '--max-turns must be a non-negative integer', + ) + } + }) - expect(repl).toContain('maxTurns={maxTurns}') + test('Commander formats invalid --max-turns as an option error', async () => { + const program = createMaxTurnsCliProgram() + await expect( + program.parseAsync(['node', 'openclaude', '--max-turns', 'nope'], { + from: 'node', + }), + ).rejects.toMatchObject({ + code: 'commander.invalidArgument', + exitCode: 1, + }) }) }) diff --git a/src/services/api/claude.lifecycle.test.ts b/src/services/api/claude.lifecycle.test.ts index 8f63f50d73..a975dd445a 100644 --- a/src/services/api/claude.lifecycle.test.ts +++ b/src/services/api/claude.lifecycle.test.ts @@ -337,6 +337,49 @@ afterEach(() => { }) describe('Claude API lifecycle tracking', () => { + test('checks provider-request ownership immediately before dispatch', async () => { + setClientTestEnv() + process.env.OPENCLAUDE_MAX_RETRIES = '0' + const queryLifecycle = new QueryLifecycleOperationTracker() + const events: string[] = [] + let permissionContextReads = 0 + const fetchOverride: FetchOverride = async () => { + events.push('fetch') + return makeJsonResponse(makeBetaMessage()) + } + + const generator = queryModelWithStreaming({ + messages: [ + { + type: 'user', + uuid: '00000000-0000-0000-0000-000000000001', + timestamp: '2026-06-17T00:00:00.000Z', + message: { role: 'user', content: 'hello' }, + } as Message, + ], + systemPrompt: asSystemPrompt([]), + thinkingConfig: { type: 'disabled' }, + tools: [], + signal: new AbortController().signal, + options: { + ...makeOptions(queryLifecycle), + fetchOverride, + getToolPermissionContext: async () => { + permissionContextReads++ + return getEmptyToolPermissionContext() + }, + onProviderRequestStart: () => { + events.push('ownership-check') + return false + }, + }, + }) + + expect(await drainGenerator(generator)).toBeUndefined() + expect(events).toEqual(['ownership-check']) + expect(permissionContextReads).toBe(0) + }) + test('ends a failed streaming dispatch before retry backoff is reported', async () => { setClientTestEnv() process.env.OPENCLAUDE_MAX_RETRIES = '1' @@ -652,6 +695,7 @@ describe('Claude API lifecycle tracking', () => { ), ) + if (result === null) throw new Error('expected non-streaming response') expect(result.id).toBe('msg-lifecycle-test') expect(requestSnapshots).toHaveLength(1) expect(requestSnapshots[0]?.apiCalls).toHaveLength(1) @@ -662,6 +706,38 @@ describe('Claude API lifecycle tracking', () => { expect(queryLifecycle.snapshot().apiCalls).toEqual([]) }) + test('non-streaming fallback checks ownership before dispatch', async () => { + setClientTestEnv() + const queryLifecycle = new QueryLifecycleOperationTracker() + let fetchCalls = 0 + const fetchOverride: FetchOverride = async () => { + fetchCalls++ + return makeJsonResponse(makeBetaMessage()) + } + + const result = await drainGenerator( + executeNonStreamingRequest( + { model: 'claude-lifecycle-test', source: 'sdk', fetchOverride }, + { + model: 'claude-lifecycle-test', + thinkingConfig: { type: 'disabled' }, + signal: new AbortController().signal, + querySource: 'sdk', + }, + makeParams, + () => {}, + () => {}, + null, + queryLifecycle, + () => false, + ), + ) + + expect(result).toBeNull() + expect(fetchCalls).toBe(0) + expect(queryLifecycle.snapshot().apiCalls).toEqual([]) + }) + test('clears non-streaming fallback lifecycle entries after request errors', async () => { setClientTestEnv() process.env.OPENCLAUDE_MAX_RETRIES = '0' diff --git a/src/services/api/claude.ts b/src/services/api/claude.ts index 3a778cb65b..91d2dd7646 100644 --- a/src/services/api/claude.ts +++ b/src/services/api/claude.ts @@ -737,6 +737,11 @@ export type Options = { providerOverride?: { model: string; baseURL: string; apiKey: string } queryLifecycle?: QueryLifecycleOperationTracker messageNormalizationTools?: Tools + /** + * Synchronous ownership check invoked immediately before an outbound + * provider request. Returning false skips the request without retrying. + */ + onProviderRequestStart?: () => boolean } export async function queryModelWithoutStreaming({ @@ -913,7 +918,8 @@ export async function* executeNonStreamingRequest( */ originatingRequestId?: string | null, queryLifecycle?: QueryLifecycleOperationTracker, -): AsyncGenerator { + onProviderRequestStart?: () => boolean, +): AsyncGenerator { const fallbackTimeoutMs = getNonstreamingFallbackTimeoutMs() const generator = withRetry( () => @@ -927,6 +933,7 @@ export async function* executeNonStreamingRequest( }), async (anthropic, attempt, context) => { const start = Date.now() + if (onProviderRequestStart?.() === false) return null const retryParams = paramsFromContext(context) captureRequest(retryParams) onAttempt(attempt, start, retryParams.max_tokens) @@ -1000,7 +1007,7 @@ export async function* executeNonStreamingRequest( } } while (!e.done) - return e.value as BetaMessage + return e.value as BetaMessage | null } /** @@ -1913,36 +1920,6 @@ async function* queryModel( } } - // Compute log scalars synchronously so the fire-and-forget .then() closure - // captures only primitives instead of paramsFromContext's full closure scope - // (messagesForAPI, system, allTools, betas — the entire request-building - // context), which would otherwise be pinned until the promise resolves. - { - const queryParams = paramsFromContext({ - model: options.model, - thinkingConfig, - }) - const logMessagesLength = queryParams.messages.length - const logBetas = useBetas ? (queryParams.betas ?? []) : [] - const logThinkingType = queryParams.thinking?.type ?? 'disabled' - const logEffortValue = queryParams.output_config?.effort - void options.getToolPermissionContext().then(permissionContext => { - logAPIQuery({ - model: options.model, - messagesLength: logMessagesLength, - temperature: options.temperatureOverride ?? 1, - betas: logBetas, - permissionMode: permissionContext.mode, - querySource: options.querySource, - queryTracking: options.queryTracking, - thinkingType: logThinkingType, - effortValue: logEffortValue, - fastMode: isFastMode, - previousRequestId, - }) - }) - } - const newMessages: AssistantMessage[] = [] let ttftMs = 0 let partialMessage: BetaMessage | undefined = undefined @@ -1957,10 +1934,11 @@ async function* queryModel( let research: unknown = undefined let isFastModeRequest = isFastMode // Keep separate state as it may change if falling back let isAdvisorInProgress = false + let apiQueryLogged = false try { queryCheckpoint('query_client_creation_start') - const generator = withRetry( + const generator = withRetry | null>( () => getAnthropicClient({ maxRetries: 0, // Disabled auto-retry in favor of manual implementation @@ -1981,11 +1959,44 @@ async function* queryModel( // client_creation_start is meaningful on attempt 1. queryCheckpoint('query_client_creation_end') + // Keep this immediately adjacent to the SDK call below. query.ts uses + // it to atomically reserve a shared foreground/background turn only + // after every asynchronous provider-preparation step has completed. + if (options.onProviderRequestStart?.() === false) { + return null + } + + // Everything below is synchronous until the SDK request is created, + // so ownership cannot change between this request build and dispatch. const params = paramsFromContext(context) captureAPIRequest(params, options.querySource) // Capture for bug reports - maxOutputTokens = params.max_tokens + if (!apiQueryLogged) { + apiQueryLogged = true + // Capture primitives only: the fire-and-forget permission lookup + // must not retain the full request-building closure. + const logMessagesLength = params.messages.length + const logBetas = useBetas ? (params.betas ?? []) : [] + const logThinkingType = params.thinking?.type ?? 'disabled' + const logEffortValue = params.output_config?.effort + void options.getToolPermissionContext().then(permissionContext => { + logAPIQuery({ + model: options.model, + messagesLength: logMessagesLength, + temperature: options.temperatureOverride ?? 1, + betas: logBetas, + permissionMode: permissionContext.mode, + querySource: options.querySource, + queryTracking: options.queryTracking, + thinkingType: logThinkingType, + effortValue: logEffortValue, + fastMode: isFastMode, + previousRequestId, + }) + }) + } + // Fire immediately before the fetch is dispatched. .withResponse() below // awaits until response headers arrive, so this MUST be before the await // or the "Network TTFB" phase measurement is wrong. @@ -2053,10 +2064,11 @@ async function* queryModel( e = await generator.next() // yield API error messages (the stream has a 'controller' property, error messages don't) - if (!('controller' in e.value)) { + if (!e.done && !('controller' in e.value)) { yield e.value } } while (!e.done) + if (e.value === null) return stream = e.value as Stream // reset state @@ -2876,8 +2888,11 @@ async function* queryModel( params => captureAPIRequest(params, options.querySource), streamRequestId, options.queryLifecycle, + options.onProviderRequestStart, ) + if (result === null) return + const m: AssistantMessage = { message: { ...result, @@ -2996,8 +3011,11 @@ async function* queryModel( params => captureAPIRequest(params, options.querySource), failedRequestId, options.queryLifecycle, + options.onProviderRequestStart, ) + if (result === null) return + const m: AssistantMessage = { message: { ...result, diff --git a/src/tasks/LocalMainSessionTask.test.ts b/src/tasks/LocalMainSessionTask.test.ts new file mode 100644 index 0000000000..7558720975 --- /dev/null +++ b/src/tasks/LocalMainSessionTask.test.ts @@ -0,0 +1,408 @@ +import { afterEach, beforeEach, describe, expect, test } from 'bun:test' +import { mkdtemp, mkdir, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +import { getOriginalCwd, setOriginalCwd } from '../bootstrap/state.js' +import { + acquireSharedMutationLock, + releaseSharedMutationLock, +} from '../test/sharedMutationLock.js' +import type { QueryParams } from '../query.js' +import { createQueryTurnBudget } from '../query.js' +import type { Terminal } from '../query/transitions.js' +import { getDefaultAppState, type AppState } from '../state/AppStateStore.js' +import type { Message } from '../types/message.js' +import { createAttachmentMessage } from '../utils/attachments.js' +import { createAssistantMessage } from '../utils/messages.js' +import { + getClaudeConfigHomeDir, + getClaudeConfigHomeDirOverrideForTesting, + setClaudeConfigHomeDirForTesting, +} from '../utils/envUtils.js' +import { dequeueAll } from '../utils/messageQueueManager.js' +import { getClaudeTempDir } from '../utils/permissions/filesystem.js' +import { + flushSessionStorage, + getProjectDir, + resetProjectForTesting, +} from '../utils/sessionStorage.js' +import { + _clearOutputsForTest, + _resetTaskOutputDirForTest, +} from '../utils/task/diskOutput.js' +import { + isMainSessionTask, + startBackgroundSession, +} from './LocalMainSessionTask.js' + +async function waitFor( + condition: () => boolean, + description: string, +): Promise { + const deadline = Date.now() + 2_000 + while (!condition()) { + if (Date.now() >= deadline) { + throw new Error(`Timed out waiting for ${description}`) + } + await new Promise(resolve => setTimeout(resolve, 5)) + } +} + +let originalCwd: string +let originalConfigDir: string | undefined +let originalClaudeTmpDir: string | undefined +let testRoot: string + +beforeEach(async () => { + await acquireSharedMutationLock('tasks/LocalMainSessionTask.test.ts') + originalCwd = getOriginalCwd() + originalConfigDir = getClaudeConfigHomeDirOverrideForTesting() + originalClaudeTmpDir = process.env.CLAUDE_CODE_TMPDIR + testRoot = await mkdtemp(join(tmpdir(), 'openclaude-main-session-task-')) + const projectDir = join(testRoot, 'project') + await mkdir(projectDir) + + setClaudeConfigHomeDirForTesting(join(testRoot, 'config')) + getClaudeConfigHomeDir.cache?.clear?.() + getProjectDir.cache?.clear?.() + setOriginalCwd(projectDir) + process.env.CLAUDE_CODE_TMPDIR = join(testRoot, 'tmp') + getClaudeTempDir.cache?.clear?.() + _resetTaskOutputDirForTest() + resetProjectForTesting() +}) + +afterEach(async () => { + dequeueAll() + try { + await _clearOutputsForTest() + await flushSessionStorage() + } finally { + resetProjectForTesting() + _resetTaskOutputDirForTest() + setClaudeConfigHomeDirForTesting(originalConfigDir) + getClaudeConfigHomeDir.cache?.clear?.() + getProjectDir.cache?.clear?.() + setOriginalCwd(originalCwd) + if (originalClaudeTmpDir === undefined) { + delete process.env.CLAUDE_CODE_TMPDIR + } else { + process.env.CLAUDE_CODE_TMPDIR = originalClaudeTmpDir + } + getClaudeTempDir.cache?.clear?.() + await rm(testRoot, { recursive: true, force: true }) + releaseSharedMutationLock() + } +}) + +describe('LocalMainSessionTask', () => { + test('does not register a task when foreground settlement needs no continuation', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + + const { taskId } = startBackgroundSession({ + description: 'settled foreground', + setAppState, + prepare: async () => null, + }) + + await new Promise(resolve => setTimeout(resolve, 0)) + + expect(state.tasks[taskId]).toBeUndefined() + }) + + test('restores claimed notifications when preparation returns after abort', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + let restoreCalls = 0 + + const { taskId } = startBackgroundSession({ + description: 'aborted preparation', + setAppState, + prepare: async abortController => { + abortController.abort('stopped') + return { + messages: [], + restoreNotificationsIfUnsent: () => { + restoreCalls++ + }, + queryParams: {} as Omit, + } + }, + }) + + await new Promise(resolve => setTimeout(resolve, 0)) + + expect(restoreCalls).toBe(1) + expect(state.tasks[taskId]).toBeUndefined() + }) + + test('restores claimed notifications when background aborts before provider dispatch', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + let restoreCalls = 0 + let releaseQuery!: () => void + const queryBlocked = new Promise(resolve => { + releaseQuery = resolve + }) + let taskRegistered = false + + const { taskId, abortController } = startBackgroundSession({ + description: 'pre-dispatch abort', + setAppState, + prepare: async () => ({ + messages: [], + restoreNotificationsIfUnsent: () => { + restoreCalls++ + }, + queryParams: {} as Omit, + }), + onRegistered: () => { + taskRegistered = true + }, + queryImpl: (async function* (): AsyncGenerator { + await queryBlocked + yield createAssistantMessage({ content: 'late' }) + return { reason: 'completed' } + }) as typeof import('../query.js').query, + }) + + await waitFor(() => taskRegistered, 'task registration') + abortController.abort('stopped') + releaseQuery() + await waitFor(() => restoreCalls === 1, 'notification restore') + expect(state.tasks[taskId]?.status).toBe('killed') + }) + + test('does not restore notifications after provider dispatch commits ownership', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + let restoreCalls = 0 + let ownershipActive = true + let releaseQuery!: () => void + const queryBlocked = new Promise(resolve => { + releaseQuery = resolve + }) + let taskRegistered = false + + const { taskId, abortController } = startBackgroundSession({ + description: 'post-dispatch abort', + setAppState, + prepare: async () => ({ + messages: [], + restoreNotificationsIfUnsent: () => { + if (!ownershipActive) return + ownershipActive = false + restoreCalls++ + }, + commitNotificationOwnership: () => { + ownershipActive = false + }, + queryParams: {} as Omit, + }), + onRegistered: () => { + taskRegistered = true + }, + queryImpl: (async function* ( + params, + ): AsyncGenerator { + params.onProviderDispatchAccepted?.() + await queryBlocked + yield createAssistantMessage({ content: 'late' }) + return { reason: 'completed' } + }) as typeof import('../query.js').query, + }) + + await waitFor(() => taskRegistered, 'task registration') + abortController.abort('stopped') + releaseQuery() + await waitFor(() => state.tasks[taskId]?.status === 'killed', 'task killed') + expect(restoreCalls).toBe(0) + }) + + test('does not restore merged notifications when background completes at turn cap without provider dispatch', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + let restoreCalls = 0 + let commitCalls = 0 + const notificationAttachment = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: + 'sabc1234done', + }) + + const { taskId } = startBackgroundSession({ + description: 'turn-cap completion', + setAppState, + prepare: async () => ({ + messages: [notificationAttachment], + restoreNotificationsIfUnsent: () => { + restoreCalls++ + }, + commitNotificationOwnership: () => { + commitCalls++ + }, + queryParams: { + turnBudget: createQueryTurnBudget(1), + } as Omit, + }), + queryImpl: (async function* (): AsyncGenerator { + return { reason: 'max_turns', turnCount: 2 } + }) as typeof import('../query.js').query, + }) + + await waitFor( + () => state.tasks[taskId]?.status === 'completed', + 'background task completion', + ) + + expect(restoreCalls).toBe(0) + expect(commitCalls).toBe(0) + }) + + test('retains a max-turn terminal attachment in task messages', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + const cap = createAttachmentMessage({ + type: 'max_turns_reached', + maxTurns: 1, + turnCount: 2, + }) + const queryImpl = async function* (): AsyncGenerator { + yield cap + return { reason: 'max_turns', turnCount: 2 } + } + + const { taskId } = startBackgroundSession({ + description: 'cap retention test', + setAppState, + queryImpl: queryImpl as typeof import('../query.js').query, + prepare: async () => ({ + messages: [], + queryParams: {} as Omit, + }), + }) + + await waitFor( + () => state.tasks[taskId]?.status === 'completed', + 'background task completion', + ) + + const task = state.tasks[taskId] + expect(isMainSessionTask(task)).toBe(true) + if (!isMainSessionTask(task)) throw new Error('expected main-session task') + expect(task.messages).toEqual([cap]) + }) + + test('reports a preparation failure before task registration', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + const failure = new Error('context failed') + const errors: unknown[] = [] + const cancelled: boolean[] = [] + + const { taskId } = startBackgroundSession({ + description: 'failing preparation', + setAppState, + prepare: async () => { + throw failure + }, + onPreparationError: error => errors.push(error), + onContinuationCancelled: () => cancelled.push(true), + }) + + await waitFor(() => errors.length === 1, 'preparation failure callback') + + expect(errors).toEqual([failure]) + expect(cancelled).toEqual([true]) + expect(state.tasks[taskId]).toBeUndefined() + }) + + test('settles the pending controller after preparation completes', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + const settledControllers: AbortController[] = [] + + const { taskId, abortController } = startBackgroundSession({ + description: 'settled preparation', + setAppState, + prepare: async () => null, + onSettled: controller => settledControllers.push(controller), + }) + + await waitFor(() => settledControllers.length === 1, 'handoff settlement') + + expect(settledControllers).toEqual([abortController]) + expect(state.tasks[taskId]).toBeUndefined() + }) + + test('reports when a pending handoff is registered as a task', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + const registeredControllers: AbortController[] = [] + + const { taskId, abortController } = startBackgroundSession({ + description: 'registered handoff', + setAppState, + prepare: async () => ({ + messages: [], + queryParams: {} as Omit, + }), + onRegistered: controller => registeredControllers.push(controller), + queryImpl: (async function* (): AsyncGenerator { + return { reason: 'completed' } + }) as typeof import('../query.js').query, + }) + + await waitFor(() => state.tasks[taskId]?.status === 'completed', 'task completion') + + expect(registeredControllers).toEqual([abortController]) + }) + + test('exposes a controller that cancels preparation before registration', async () => { + let state = getDefaultAppState() + const setAppState = (update: (previous: AppState) => AppState): void => { + state = update(state) + } + let releasePreparation!: () => void + const preparation = new Promise(resolve => { + releasePreparation = resolve + }) + + const { taskId, abortController } = startBackgroundSession({ + description: 'cancellable preparation', + setAppState, + prepare: async controller => { + await preparation + if (controller.signal.aborted) return null + throw new Error('preparation should have been cancelled') + }, + }) + + abortController.abort('user-cancel') + releasePreparation() + await new Promise(resolve => setTimeout(resolve, 0)) + + expect(state.tasks[taskId]).toBeUndefined() + }) +}) diff --git a/src/tasks/LocalMainSessionTask.ts b/src/tasks/LocalMainSessionTask.ts index 57763f3b3a..0c87ab76d2 100644 --- a/src/tasks/LocalMainSessionTask.ts +++ b/src/tasks/LocalMainSessionTask.ts @@ -48,7 +48,11 @@ import { getTaskOutputPath, initTaskOutputAsSymlink, } from '../utils/task/diskOutput.js' -import { registerTask, updateTaskState } from '../utils/task/framework.js' +import { + PANEL_GRACE_MS, + registerTask, + updateTaskState, +} from '../utils/task/framework.js' import type { LocalAgentTaskState } from './LocalAgentTask/LocalAgentTask.js' // Main session tasks use LocalAgentTaskState with agentType='main-session' @@ -88,15 +92,16 @@ function generateMainSessionTaskId(): string { * @param setAppState - State setter function * @param mainThreadAgentDefinition - Optional agent definition if running with --agent * @param existingAbortController - Optional abort controller to reuse (for backgrounding an active query) - * @returns Object with task ID and abort signal for stopping the background query + * @returns Object with task ID and the controller that owns background cancellation */ export function registerMainSessionTask( description: string, setAppState: SetAppState, mainThreadAgentDefinition?: AgentDefinition, existingAbortController?: AbortController, -): { taskId: string; abortSignal: AbortSignal } { - const taskId = generateMainSessionTaskId() + existingTaskId?: string, +): { taskId: string; abortController: AbortController } { + const taskId = existingTaskId ?? generateMainSessionTaskId() // Link output to an isolated per-task transcript file (same layout as // sub-agents). Do NOT use getTranscriptPath() — that's the main session's @@ -157,7 +162,7 @@ export function registerMainSessionTask( return prev }) - return { taskId, abortSignal: abortController.signal } + return { taskId, abortController } } /** @@ -328,37 +333,99 @@ type ToolActivity = { input: Record } +export type BackgroundSessionHandle = { + taskId: string + abortController: AbortController +} + /** * Start a fresh background session with the given messages. * - * Spawns an independent query() call with the current messages and registers it - * as a background task. The caller's foreground query continues running normally. + * Prepares the continuation first, then registers the background task and + * starts an independent query() call. A foreground that settles without a + * successor therefore never publishes a task. */ export function startBackgroundSession({ - messages, - queryParams, + prepare, description, setAppState, agentDefinition, + queryImpl = query, + onPreparationError, + onContinuationCancelled, + onRegistered, + onSettled, }: { - messages: Message[] - queryParams: Omit + prepare: (abortController: AbortController) => Promise< + | { + messages: Message[] + restoreNotificationsIfUnsent?: () => void + commitNotificationOwnership?: () => void + queryParams: Omit + } + | null + > description: string setAppState: SetAppState agentDefinition?: AgentDefinition -}): string { - const { taskId, abortSignal } = registerMainSessionTask( - description, - setAppState, - agentDefinition, - ) + queryImpl?: typeof query + onPreparationError?: (error: unknown) => void + onContinuationCancelled?: () => void + onRegistered?: (abortController: AbortController) => void + onSettled?: (abortController: AbortController) => void +}): BackgroundSessionHandle { + // Keep a controller while preparation waits for the foreground settlement, + // but do not publish a task until a successor is actually required. A + // settled foreground can complete normally after Ctrl+B is pressed. + const taskId = generateMainSessionTaskId() + const abortController = createAbortController() + const abortSignal = abortController.signal + let taskRegistered = false + let providerRequestStarted = false + let restoreNotificationsIfUnsent: (() => void) | undefined + let commitNotificationOwnership: (() => void) | undefined + + const restorePreDispatchNotifications = (): void => { + if (providerRequestStarted) return + const restore = restoreNotificationsIfUnsent + restoreNotificationsIfUnsent = undefined + restore?.() + } - // Persist the pre-backgrounding conversation to the task's isolated - // transcript so TaskOutput shows context immediately. Subsequent messages - // are written incrementally below. - void recordSidechainTranscript(messages, taskId).catch(err => - logForDebugging(`bg-session initial transcript write failed: ${err}`), - ) + const finishAbortedTask = (): void => { + restorePreDispatchNotifications() + onContinuationCancelled?.() + // chat:killAgents already marks the task notified and emits this event. + // stopTask kills it without doing either for local-agent tasks, so close + // the SDK lifecycle exactly once no matter which phase observed abort. + let alreadyNotified = false + let shouldEvictOutput = false + updateTaskState(taskId, setAppState, task => { + alreadyNotified = task.notified === true + if (task.status === 'running') { + shouldEvictOutput = true + task.unregisterCleanup?.() + const endTime = Date.now() + return { + ...task, + status: 'killed', + endTime, + evictAfter: task.retain ? undefined : endTime + PANEL_GRACE_MS, + abortController: undefined, + unregisterCleanup: undefined, + selectedAgent: undefined, + notified: true, + } + } + return alreadyNotified ? task : { ...task, notified: true } + }) + if (shouldEvictOutput) void evictTaskOutput(taskId) + if (!alreadyNotified) { + emitTaskTerminatedSdk(taskId, 'stopped', { + summary: description, + }) + } + } // Wrap in agent context so skill invocations scope to this task's agentId // (not null). This lets clearInvokedSkills(preservedAgentIds) selectively @@ -373,36 +440,73 @@ export function startBackgroundSession({ void runWithAgentContext(agentContext, async () => { try { + // Wait for the foreground settlement before publishing a task. This + // keeps a completion race from producing a phantom background task. + const prepared = await prepare(abortController) + if (prepared === null) { + return + } + restoreNotificationsIfUnsent = prepared.restoreNotificationsIfUnsent + commitNotificationOwnership = prepared.commitNotificationOwnership + if (abortSignal.aborted) { + restorePreDispatchNotifications() + onContinuationCancelled?.() + return + } + try { + registerMainSessionTask( + description, + setAppState, + agentDefinition, + abortController, + taskId, + ) + } catch (error) { + restorePreDispatchNotifications() + if (abortSignal.aborted) { + onContinuationCancelled?.() + return + } + throw error + } + taskRegistered = true + onRegistered?.(abortController) + const { messages, queryParams } = prepared + + // Persist the pre-backgrounding conversation to the task's isolated + // transcript so TaskOutput shows context immediately. Subsequent + // messages are written incrementally below. + void recordSidechainTranscript(messages, taskId).catch(err => + logForDebugging(`bg-session initial transcript write failed: ${err}`), + ) + const bgMessages: Message[] = [...messages] const recentActivities: ToolActivity[] = [] let toolCount = 0 let tokenCount = 0 let lastRecordedUuid: UUID | null = messages.at(-1)?.uuid ?? null - for await (const event of query({ + for await (const event of queryImpl({ messages: bgMessages, ...queryParams, + onProviderDispatchAccepted: () => { + providerRequestStarted = true + commitNotificationOwnership?.() + }, })) { if (abortSignal.aborted) { - // Aborted mid-stream — completeMainSessionTask won't be reached. - // chat:killAgents path already marked notified + emitted; stopTask path did not. - let alreadyNotified = false - updateTaskState(taskId, setAppState, task => { - alreadyNotified = task.notified === true - return alreadyNotified ? task : { ...task, notified: true } - }) - if (!alreadyNotified) { - emitTaskTerminatedSdk(taskId, 'stopped', { - summary: description, - }) - } + finishAbortedTask() return } if ( event.type !== 'user' && event.type !== 'assistant' && - event.type !== 'system' + event.type !== 'system' && + !( + event.type === 'attachment' && + event.attachment.type === 'max_turns_reached' + ) ) { continue } @@ -467,12 +571,39 @@ export function startBackgroundSession({ }) } + if (abortSignal.aborted) { + finishAbortedTask() + return + } + if (providerRequestStarted) { + commitNotificationOwnership?.() + } completeMainSessionTask(taskId, true, setAppState) } catch (error) { + if (abortSignal.aborted) { + if (taskRegistered) { + finishAbortedTask() + } else { + restorePreDispatchNotifications() + onContinuationCancelled?.() + } + return + } logError(error) - completeMainSessionTask(taskId, false, setAppState) + if (taskRegistered) { + if (providerRequestStarted) { + commitNotificationOwnership?.() + } + completeMainSessionTask(taskId, false, setAppState) + } else { + restorePreDispatchNotifications() + onPreparationError?.(error) + onContinuationCancelled?.() + } + } finally { + onSettled?.(abortController) } }) - return taskId + return { taskId, abortController } } diff --git a/src/utils/config.ts b/src/utils/config.ts index e922754771..3468d4adda 100644 --- a/src/utils/config.ts +++ b/src/utils/config.ts @@ -294,6 +294,11 @@ export type GlobalConfig = { contextCollapseEnabled: boolean // Opt-in: collapse old transcript spans into summaries (lossy; off by default) toolHistoryCompressionEnabled: boolean // Compress old tool_result content (shim providers; Anthropic-native only while prompt caching is inactive) compactTailTurns?: number // Recent messages preserved verbatim by auto-compact's relevance pruning (default: 3) + /** + * Per-prompt local interactive REPL turn cap (default: 50). + * Overridden by CLI `--max-turns` and OPENCLAUDE_MAX_TURNS / CLAUDE_CODE_MAX_TURNS. + */ + replMaxTurns?: number showTurnDuration: boolean // Controls whether to show turn duration message (e.g., "Cooked for 1m 6s") // Controls whether to show per-query cache hit/miss stats at the end of each turn. // 'off' — no display @@ -771,6 +776,7 @@ export const GLOBAL_CONFIG_KEYS = [ 'hasUsedBackslashReturn', 'autoCompactEnabled', 'compactTailTurns', + 'replMaxTurns', 'contextCollapseEnabled', 'toolHistoryCompressionEnabled', 'showTurnDuration', diff --git a/src/utils/context.test.ts b/src/utils/context.test.ts index cec4c1d110..edf3c13092 100644 --- a/src/utils/context.test.ts +++ b/src/utils/context.test.ts @@ -32,6 +32,8 @@ const originalEnv = { LONGCAT_API_KEY: process.env.LONGCAT_API_KEY, CLAUDE_CODE_MAX_CONTEXT_TOKENS: process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS, USER_TYPE: process.env.USER_TYPE, + OPENCLAUDE_MAX_TURNS: process.env.OPENCLAUDE_MAX_TURNS, + CLAUDE_CODE_MAX_TURNS: process.env.CLAUDE_CODE_MAX_TURNS, } beforeEach(async () => { @@ -53,6 +55,8 @@ beforeEach(async () => { delete process.env.LONGCAT_API_KEY delete process.env.CLAUDE_CODE_MAX_CONTEXT_TOKENS delete process.env.USER_TYPE + delete process.env.OPENCLAUDE_MAX_TURNS + delete process.env.CLAUDE_CODE_MAX_TURNS }) afterEach(() => { @@ -142,6 +146,16 @@ afterEach(() => { } else { process.env.USER_TYPE = originalEnv.USER_TYPE } + if (originalEnv.OPENCLAUDE_MAX_TURNS === undefined) { + delete process.env.OPENCLAUDE_MAX_TURNS + } else { + process.env.OPENCLAUDE_MAX_TURNS = originalEnv.OPENCLAUDE_MAX_TURNS + } + if (originalEnv.CLAUDE_CODE_MAX_TURNS === undefined) { + delete process.env.CLAUDE_CODE_MAX_TURNS + } else { + process.env.CLAUDE_CODE_MAX_TURNS = originalEnv.CLAUDE_CODE_MAX_TURNS + } } finally { clearSessionContextWindowOverride() releaseSharedMutationLock() @@ -559,8 +573,13 @@ test('unknown openai-compatible model fallback logs one debug warning and no con contextModule.getContextWindowForModel('another-unknown-3p-model'), ).toBe(128_000) expect(consoleError).not.toHaveBeenCalled() - expect(logForDebugging).toHaveBeenCalledTimes(1) - expect(logForDebugging.mock.calls[0]?.[1]).toEqual({ level: 'warn' }) + const contextWarnings = logForDebugging.mock.calls.filter( + ([message, options]) => + typeof message === 'string' && + message.startsWith('[context] Warning:') && + options?.level === 'warn', + ) + expect(contextWarnings).toHaveLength(1) } finally { console.error = originalConsoleError mock.restore() diff --git a/src/utils/messageQueueManager.test.ts b/src/utils/messageQueueManager.test.ts new file mode 100644 index 0000000000..63eaaaf292 --- /dev/null +++ b/src/utils/messageQueueManager.test.ts @@ -0,0 +1,54 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test' +import { + dequeue, + dequeueAll, + enqueue, + getCommandQueue, + prepend, + resetCommandQueue, + subscribeToCommandQueue, +} from './messageQueueManager.js' + +beforeEach(() => resetCommandQueue()) +afterEach(() => resetCommandQueue()) + +test('prepend ignores an empty array', () => { + const notifications: number[] = [] + const unsubscribe = subscribeToCommandQueue(() => notifications.push(1)) + try { + prepend([]) + expect(getCommandQueue()).toEqual([]) + expect(notifications).toHaveLength(0) + } finally { + unsubscribe() + } +}) + +test('prepend restores commands ahead of later enqueues in FIFO order', () => { + const notifications: number[] = [] + const unsubscribe = subscribeToCommandQueue(() => notifications.push(1)) + try { + const restored = [ + { value: 'first restored', mode: 'prompt' as const, priority: 'later' as const }, + { value: 'second restored', mode: 'prompt' as const, priority: 'now' as const }, + { value: 'third restored', mode: 'prompt' as const, priority: 'now' as const }, + ] + prepend(restored) + expect(notifications).toHaveLength(1) + enqueue({ value: 'later enqueue', mode: 'prompt' }) + + expect(getCommandQueue()).toEqual([ + ...restored, + { value: 'later enqueue', mode: 'prompt', priority: 'next' }, + ]) + expect(notifications).toHaveLength(2) + expect(dequeue()).toMatchObject(restored[1]!) + expect(dequeue()).toMatchObject(restored[2]!) + expect(dequeueAll()).toEqual([ + restored[0], + { value: 'later enqueue', mode: 'prompt', priority: 'next' }, + ]) + } finally { + unsubscribe() + } +}) diff --git a/src/utils/messageQueueManager.ts b/src/utils/messageQueueManager.ts index 1db54b53d8..587901af99 100644 --- a/src/utils/messageQueueManager.ts +++ b/src/utils/messageQueueManager.ts @@ -133,6 +133,22 @@ export function enqueue(command: QueuedCommand): void { ) } +/** + * Restore commands ahead of commands added after they were claimed, without + * changing their priority. Commands in the supplied array retain FIFO order. + */ +export function prepend(commands: readonly QueuedCommand[]): void { + if (commands.length === 0) return + commandQueue.unshift(...commands) + notifySubscribers() + for (const command of commands) { + logOperation( + 'enqueue', + typeof command.value === 'string' ? command.value : undefined, + ) + } +} + /** * Add a task notification to the queue. * Convenience wrapper that defaults priority to 'later' so user input diff --git a/src/utils/replMaxTurns.ts b/src/utils/replMaxTurns.ts new file mode 100644 index 0000000000..e5cd97e741 --- /dev/null +++ b/src/utils/replMaxTurns.ts @@ -0,0 +1,131 @@ +import { getGlobalConfig } from './config.js' +import { logForDebugging } from './debug.js' +import { InvalidArgumentError } from '@commander-js/extra-typings' + +export const DEFAULT_REPL_MAX_TURNS = 50 +export const MAX_TURNS_UNLIMITED_WARNING = + 'Warning: --max-turns 0 disables the local turn limit when applicable. Use with caution.' + +/** Preset values offered in `/config` (plus any current custom value). */ +export const REPL_MAX_TURNS_OPTIONS = [50, 100, 200, 500] as const + +/** + * Shared `--max-turns` help text for Commander registration and behavior tests. + * Keep remote-backed sessions explicitly out of scope in this string. + */ +export const MAX_TURNS_CLI_DESCRIPTION = + 'Maximum number of agentic turns per prompt. In local interactive mode, set to 0 for unlimited turns (use with caution). This overrides the default 50-turn REPL query cap and is also configurable via OPENCLAUDE_MAX_TURNS or /config. Does not apply to remote-backed sessions (connect/ssh/--remote). In --print mode this early-exits after the specified number of turns.' + +export function parseMaxTurnsCli(value: string): number { + const parsed = value.trim() === '' ? Number.NaN : Number(value) + if (!Number.isSafeInteger(parsed) || parsed < 0) { + throw new Error('--max-turns must be a non-negative integer') + } + return parsed +} + +/** Commander-compatible parser used by the production --max-turns option. */ +export function parseMaxTurnsCommanderArgument(value: string): number { + try { + return parseMaxTurnsCli(value) + } catch (error) { + throw new InvalidArgumentError( + error instanceof Error ? error.message : String(error), + ) + } +} + +/** + * Prefer OPENCLAUDE_MAX_TURNS; honor legacy CLAUDE_CODE_MAX_TURNS only when + * the new variable is unset/empty. Invalid, zero, negative, non-integer, or + * unsafe values are ignored (treated as absent for the chosen variable). + */ +function parsePositiveTurnEnv(name: string): number | undefined { + const raw = process.env[name] + if (!raw?.trim()) return undefined + const parsed = Number(raw.trim()) + if (Number.isSafeInteger(parsed) && parsed > 0) { + return parsed + } + return undefined +} + +/** + * Normalize a persisted or UI-selected interactive turn cap. + * Invalid values fall back to DEFAULT_REPL_MAX_TURNS. + */ +export function normalizeReplMaxTurns(value: unknown): number { + if (typeof value === 'number') { + if (Number.isSafeInteger(value) && value > 0) return value + return DEFAULT_REPL_MAX_TURNS + } + if (typeof value === 'string' && value.trim()) { + const parsed = Number(value.trim()) + if (Number.isSafeInteger(parsed) && parsed > 0) return parsed + } + return DEFAULT_REPL_MAX_TURNS +} + +export function getReplMaxTurnsWarning(maxTurns?: number): string | undefined { + return maxTurns === 0 ? MAX_TURNS_UNLIMITED_WARNING : undefined +} + +function resolveConfiguredReplMaxTurns(): number { + const configured = getGlobalConfig().replMaxTurns + if (configured === undefined) { + return DEFAULT_REPL_MAX_TURNS + } + return normalizeReplMaxTurns(configured) +} + +/** + * Resolve the per-prompt local interactive REPL turn cap. + * + * Precedence: explicit prop (CLI `--max-turns`) → OPENCLAUDE_MAX_TURNS → + * CLAUDE_CODE_MAX_TURNS (only if OPENCLAUDE_MAX_TURNS unset) → + * `/config` `replMaxTurns` → DEFAULT_REPL_MAX_TURNS (50). + * + * Applies to local interactive query loops only. Remote-backed sessions + * (connect/ssh/--remote) send prompts to a remote executor and are not + * capped here. + * + * An explicit CLI value of zero disables the cap. Other invalid explicit + * values fall through so a bad CLI parse cannot disable the interactive + * safety cap (unlike headless, where omitted maxTurns means no cap). + * If OPENCLAUDE_MAX_TURNS is set but invalid, DEFAULT_REPL_MAX_TURNS is + * used and lower layers (legacy env, /config) are not consulted — matching + * OPENCLAUDE_MAX_RETRIES precedence. + */ +export function resolveReplMaxTurns(maxTurns?: number): number | undefined { + if (maxTurns === 0) { + return undefined + } + if ( + typeof maxTurns === 'number' && + Number.isSafeInteger(maxTurns) && + maxTurns > 0 + ) { + return maxTurns + } + + const openClaudeRaw = process.env.OPENCLAUDE_MAX_TURNS + if (openClaudeRaw?.trim()) { + const parsed = parsePositiveTurnEnv('OPENCLAUDE_MAX_TURNS') + if (parsed !== undefined) { + return parsed + } + // Match OPENCLAUDE_MAX_RETRIES: set-but-invalid uses the default and does + // not fall through to legacy env or /config; surface a debug diagnostic. + logForDebugging( + `OPENCLAUDE_MAX_TURNS has an invalid value (using default: ${DEFAULT_REPL_MAX_TURNS})`, + ) + return DEFAULT_REPL_MAX_TURNS + } + + const legacy = parsePositiveTurnEnv('CLAUDE_CODE_MAX_TURNS') + if (legacy !== undefined) { + return legacy + } + + return resolveConfiguredReplMaxTurns() +} diff --git a/src/utils/taskNotificationIdentity.test.ts b/src/utils/taskNotificationIdentity.test.ts new file mode 100644 index 0000000000..f190559f1d --- /dev/null +++ b/src/utils/taskNotificationIdentity.test.ts @@ -0,0 +1,134 @@ +import { describe, expect, test } from 'bun:test' +import { + dedupeQueuedTaskNotifications, + filterClaimedTaskNotificationsForRestore, + getTaskNotificationDedupKey, + parseTaskNotificationTaskId, + pendingCommandsForEmbeddedNotifications, +} from './taskNotificationIdentity.js' +import { createAttachmentMessage } from './attachments.js' + +function taskNotification(taskId: string, summary: string): string { + return ` +${taskId} +/tmp/${taskId}.jsonl +completed +${summary} +` +} + +describe('task notification identity', () => { + test('parses the embedded task id from notification payloads', () => { + expect(parseTaskNotificationTaskId(taskNotification('sabc1234', 'done'))).toBe( + 'sabc1234', + ) + expect(parseTaskNotificationTaskId('no task id here')).toBeUndefined() + }) + + test('dedup keys stay distinct when summary text matches', () => { + const summary = 'Background session "work" completed' + const first = taskNotification('s1111111', summary) + const second = taskNotification('s2222222', summary) + + expect(getTaskNotificationDedupKey(first)).toBe('s1111111') + expect(getTaskNotificationDedupKey(second)).toBe('s2222222') + expect(getTaskNotificationDedupKey(first)).not.toBe( + getTaskNotificationDedupKey(second), + ) + }) + + test('falls back to full content when no task id is present', () => { + const payload = 'hook' + expect(getTaskNotificationDedupKey(payload)).toBe(payload) + }) + + test('dedupes duplicate task ids within a claimed notification batch', () => { + const summary = 'Background session "work" completed' + const first = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: taskNotification('s1111111', summary), + }) + const duplicate = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: taskNotification('s1111111', summary), + }) + const distinct = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: taskNotification('s2222222', summary), + }) + + expect( + dedupeQueuedTaskNotifications([], [first, duplicate, distinct]), + ).toEqual([first, distinct]) + }) + + test('restore excludes notifications already present in the settled transcript', () => { + const summary = 'Background session "work" completed' + const settledPrompt = taskNotification('s1111111', summary) + const settled = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: settledPrompt, + }) + const pending = [ + { + value: settledPrompt, + mode: 'task-notification' as const, + priority: 'later' as const, + }, + { + value: taskNotification('s2222222', summary), + mode: 'task-notification' as const, + priority: 'later' as const, + }, + ] + + expect( + filterClaimedTaskNotificationsForRestore(pending, [settled]), + ).toEqual([pending[1]]) + }) + + test('restore maps embedded successor notifications back to claimed commands', () => { + const summary = 'Background session "work" completed' + const first = { + value: taskNotification('s1111111', summary), + mode: 'task-notification' as const, + priority: 'later' as const, + } + const second = { + value: taskNotification('s2222222', summary), + mode: 'task-notification' as const, + priority: 'later' as const, + } + const embedded = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt: first.value, + }) + + expect( + pendingCommandsForEmbeddedNotifications([first, second], [embedded]), + ).toEqual([first]) + }) + + test('does not dedupe non-task-notification queued commands in the batch', () => { + const prompt = taskNotification('s1111111', 'done') + const taskAttachment = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'task-notification', + prompt, + }) + const promptAttachment = createAttachmentMessage({ + type: 'queued_command', + commandMode: 'prompt', + prompt, + }) + + expect(dedupeQueuedTaskNotifications([], [taskAttachment, promptAttachment])).toEqual( + [taskAttachment, promptAttachment], + ) + }) +}) diff --git a/src/utils/taskNotificationIdentity.ts b/src/utils/taskNotificationIdentity.ts new file mode 100644 index 0000000000..a6bd00e478 --- /dev/null +++ b/src/utils/taskNotificationIdentity.ts @@ -0,0 +1,114 @@ +import { TASK_ID_TAG } from '../constants/xml.js' +import type { Message } from '../types/message.js' +import type { QueuedCommand } from '../types/textInputTypes.js' + +const TASK_ID_PATTERN = new RegExp( + `<${TASK_ID_TAG}>([^<]+)`, +) + +/** + * Parse the stable task id embedded in a task-notification payload. + */ +export function parseTaskNotificationTaskId( + content: string, +): string | undefined { + return content.match(TASK_ID_PATTERN)?.[1] +} + +/** + * Dedup key for task notifications. Prefer the embedded task id so distinct + * tasks with identical summary text are not collapsed. + */ +export function getTaskNotificationDedupKey(content: string): string { + return parseTaskNotificationTaskId(content) ?? content +} + +function getQueuedCommandNotificationKey( + command: QueuedCommand, +): string | undefined { + if (typeof command.value !== 'string') return undefined + return getTaskNotificationDedupKey(command.value) +} + +function collectSettledNotificationKeys( + settledMessages: readonly Message[], +): Set { + const existingNotificationKeys = new Set() + for (const message of settledMessages) { + if ( + message.type === 'attachment' && + message.attachment.type === 'queued_command' && + message.attachment.commandMode === 'task-notification' && + typeof message.attachment.prompt === 'string' + ) { + existingNotificationKeys.add( + getTaskNotificationDedupKey(message.attachment.prompt), + ) + } + } + return existingNotificationKeys +} + +/** + * Claimed notifications already present in the settled foreground transcript + * must not be restored to the queue on handoff abort. + */ +export function filterClaimedTaskNotificationsForRestore( + pendingNotifications: readonly QueuedCommand[], + settledMessages: readonly Message[], +): QueuedCommand[] { + const existingNotificationKeys = collectSettledNotificationKeys(settledMessages) + return pendingNotifications.filter(command => { + const key = getQueuedCommandNotificationKey(command) + return key === undefined || !existingNotificationKeys.has(key) + }) +} + +/** + * Map embedded handoff notification attachments back to their claimed queue + * commands so restore only returns notifications actually woven into the + * successor batch. + */ +export function pendingCommandsForEmbeddedNotifications( + pendingNotifications: readonly QueuedCommand[], + embeddedNotifications: readonly Message[], +): QueuedCommand[] { + const embeddedKeys = new Set() + for (const message of embeddedNotifications) { + if ( + message.type === 'attachment' && + message.attachment.type === 'queued_command' && + message.attachment.commandMode === 'task-notification' && + typeof message.attachment.prompt === 'string' + ) { + embeddedKeys.add(getTaskNotificationDedupKey(message.attachment.prompt)) + } + } + return pendingNotifications.filter(command => { + const key = getQueuedCommandNotificationKey(command) + return key !== undefined && embeddedKeys.has(key) + }) +} + +export function dedupeQueuedTaskNotifications( + settledMessages: readonly Message[], + notificationMessages: readonly Message[], +): Message[] { + const existingNotificationKeys = collectSettledNotificationKeys(settledMessages) + return notificationMessages.filter(message => { + if ( + message.type !== 'attachment' || + message.attachment.type !== 'queued_command' || + message.attachment.commandMode !== 'task-notification' || + typeof message.attachment.prompt !== 'string' + ) { + return true + } + const key = getTaskNotificationDedupKey(message.attachment.prompt) + if (existingNotificationKeys.has(key)) { + return false + } + existingNotificationKeys.add(key) + return true + }) +}