From b0cbfe11000eef4d85b96fc3173afbb8b0930c81 Mon Sep 17 00:00:00 2001 From: JATMN Date: Thu, 6 Aug 2026 05:24:34 -0700 Subject: [PATCH] fix(repl): make local interactive max-turns configurable (#2086) * fix(repl): make interactive max-turns configurable Wire --max-turns into interactive sessionConfig and honor OPENCLAUDE_MAX_TURNS / CLAUDE_CODE_MAX_TURNS so long autonomous REPL sessions can raise the default 50-turn per-prompt cap (fixes #2079). * fix(repl): forward --max-turns on connect/ssh/remote launches sessionConfig covered the normal interactive paths; connect, SSH, assistant, and --remote built REPL props without spreading it, so the CLI override was dropped despite help advertising interactive support. * fix(repl): scope interactive max-turns to local query loops Remote-backed sessions bypass local query(), so forwarding --max-turns into those REPL props over-claimed enforcement. Clarify help/docs and match OPENCLAUDE_MAX_RETRIES precedence when OPENCLAUDE_MAX_TURNS is set but invalid. * feat(config): add interactive max turns under /config Expose replMaxTurns in the Config panel (50/100/200/500) and resolve it after CLI/env so local interactive sessions can raise the per-prompt cap without restarting. Resolve at query time so mid-session /config changes apply on the next prompt. * docs(repl): clarify invalid OPENCLAUDE_MAX_TURNS precedence Match the OPENCLAUDE_MAX_RETRIES contract: a set-but-invalid primary env var uses the default and does not fall through to legacy or /config. * fix(repl): address PR review on max-turns help and web version gate Share the --max-turns Commander description via an imported constant so help and tests stay in sync without breaking the CLI bundle, replace source-only help assertions with Commander behavior coverage, and add the published 0.27.0 entry so web verify-dist passes. * fix(repl): typecheck Commander maxTurns opts and warn on invalid env Avoid TS2339 on untyped Commander opts, log invalid OPENCLAUDE_MAX_TURNS like MAX_RETRIES, and clarify that /config shows the persisted preference. * fix(repl): warn when max turns is unlimited * fix(repl): scope unlimited-turn warning locally * fix(repl): preserve interactive turn caps across backgrounding * fix(repl): preserve turn caps when backgrounding * fix(repl): share turn budget across background handoff * fix(repl): reserve turns at provider dispatch * fix(repl): snapshot background handoff transcript * fix(tasks): avoid phantom background session task * fix(repl): preserve handoff lifecycle state * fix(repl): own pending background handoffs * test(tasks): isolate background session task storage * fix(repl): refresh background task title and test cleanup * test(repl): cover max-turn CLI dispatch paths * test(queue): cover prepend notification and priority * fix(repl): skip background handoff after foreground query throws Rebased onto main and gate Ctrl+B continuation on !didThrow so a faulted foreground turn cannot start a background session from partial state. * fix(repl): address PR review findings on notifications and handoff Dedupe background task notifications by embedded task id, gate background continuation on preflight veto, scope queue removal to main-thread notifications, and forward maxTurns through all launchRepl entry points. * fix(repl): resolve latest CodeRabbit inline review findings Dedupe claimed notification batches by task id, restore notifications on pre-registration abort, tighten test isolation, and replace remaining brittle source-text assertions with behavioral coverage. * fix(repl): keep notification restore active until provider dispatch Stop clearing notification ownership when preparation succeeds so pre-dispatch aborts can restore claimed queue items, and commit ownership once the provider starts. Add regression coverage for the abort path and headless max-turns zero. * test(repl): cover post-dispatch ownership and headless max-turns 0 Add regression tests for notification restore after provider dispatch commits ownership, and assert headless --max-turns 0 reaches query() without interactive resolution stripping the value. * fix(repl): restore only embeddable notifications on Ctrl+B handoff abort Track the deduped successor subset when restoring claimed main-thread task notifications so items already in the settled foreground transcript are not re-queued. Clarify that agent-scoped notifications intentionally stay on their owner drain path (issue #2079 scope is interactive turn caps only). * fix(repl): address review findings on background handoff Commit notification ownership when background sessions complete without provider dispatch, forward all task notifications on Ctrl+B again, and restore deferred max-turn cap attachments when continuation is cancelled. * fix(repl): guard deferred cap restore and skip remote turn limits Anchor deferred max-turn restoration to the handed-off transcript tail so a cancelled Ctrl+B handoff cannot attach the prior prompt's cap to a newer turn. Apply the interactive turn cap only in local sessions and align remote-session docs/help wording. * fix(repl): use messagesRef for deferred cap transcript anchor persistentMessages is block-scoped inside onQuery try; read the settled tail from messagesRef in finally so typecheck passes. * test: harden context fallback warning assertion after max-turns tests Scope the unknown-model context test to [context] warnings only so unrelated import-time debug logs do not fail CI, and clear turn env vars in both that test and replMaxTurnsProp setup to avoid cross-file pollution. * test: address PR review findings on headless max-turns boundary Add a runHeadless-to-ask regression that asserts maxTurns 0 is forwarded through the headless print path, and restore OPENCLAUDE_MAX_TURNS env vars in context.test.ts after the unknown-model fallback test mutates them. * test: tidy headless max-turns boundary test and env isolation Mock headless stdout so runHeadless completes cleanly without leaking output, restore spies in finally, and centralize turn-env cleanup in context.test beforeEach. * fix: address PR review findings for max-turns background handoff Separate model-request lifecycle from provider dispatch acceptance so interruption correction arms before async prep, notification ownership commits only after dispatch, deferred turn caps restore on every abort path, and foreground work stays blocked while handoff preparation runs. --- docs/advanced-setup.md | 1 + src/cli/printMaxTurns.test.ts | 223 ++++++++++ src/components/SessionBackgroundHint.tsx | 2 + src/components/Settings/Config.tsx | 27 ++ src/main.tsx | 19 +- src/query.ts | 237 ++++++++-- src/query/agentStepLimit.test.ts | 438 +++++++++++++++++- src/query/requestOnlyMessages.test.ts | 28 +- src/screens/REPL.tsx | 318 +++++++++++--- src/screens/replMaxTurns.ts | 190 +++++++- src/screens/replMaxTurnsProp.test.ts | 487 +++++++++++++++++++-- src/services/api/claude.lifecycle.test.ts | 76 ++++ src/services/api/claude.ts | 88 ++-- src/tasks/LocalMainSessionTask.test.ts | 408 +++++++++++++++++ src/tasks/LocalMainSessionTask.ts | 209 +++++++-- src/utils/config.ts | 6 + src/utils/context.test.ts | 23 +- src/utils/messageQueueManager.test.ts | 54 +++ src/utils/messageQueueManager.ts | 16 + src/utils/replMaxTurns.ts | 131 ++++++ src/utils/taskNotificationIdentity.test.ts | 134 ++++++ src/utils/taskNotificationIdentity.ts | 114 +++++ 22 files changed, 3000 insertions(+), 229 deletions(-) create mode 100644 src/cli/printMaxTurns.test.ts create mode 100644 src/tasks/LocalMainSessionTask.test.ts create mode 100644 src/utils/messageQueueManager.test.ts create mode 100644 src/utils/replMaxTurns.ts create mode 100644 src/utils/taskNotificationIdentity.test.ts create mode 100644 src/utils/taskNotificationIdentity.ts diff --git a/docs/advanced-setup.md b/docs/advanced-setup.md index e447b7eef..8f096f55c 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 000000000..215366191 --- /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 bd5e224c1..5be8328e5 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 815c58e09..89d246efb 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 a7832be9b..be9171f38 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 1e4581765..8dc317f78 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 ee9a8f639..6501d936a 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 37da35576..268c9a8a5 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 decd2637f..f40a77f84 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). @@ -1043,7 +1068,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. @@ -2839,60 +2864,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, - toolUseContext, - fallbackModel, - maxTurns, - querySource: getQuerySourceForREPL(), - autoCompactTracking: getAutoCompactTrackingForSession(backgroundSessionId), - onAutoCompactTrackingChange: tracking => { - setAutoCompactTrackingForSession(backgroundSessionId, tracking); + }; + 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); } - }, - description: terminalTitle, - setAppState, - agentDefinition: mainThreadAgentDefinition - }); - })(); - }, [abortController, mainLoopModel, toolPermissionContext, mainThreadAgentDefinition, getToolUseContext, customSystemPrompt, appendSystemPrompt, canUseTool, setAppState, getAutoCompactTrackingForSession, setAutoCompactTrackingForSession, fallbackModel, maxTurns]); + }; + 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, + 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); + } + } + }; + } 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({ @@ -2967,7 +3094,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). @@ -3111,7 +3238,7 @@ export function REPL({ resetTurnToolDuration(); resetTurnClassifierDuration(); let expectedAutoCompactTracking = queryAutoCompactTracking; - for await (const event of query({ + const queryGenerator = query({ messages: messagesIncludingNewMessages, requestOnlyMessages, onModelRequestStart: interruptionCorrectionQueryId @@ -3131,16 +3258,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 : { @@ -3157,7 +3298,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()) { @@ -3174,6 +3316,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, @@ -3196,10 +3348,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; @@ -3255,7 +3418,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; @@ -3269,6 +3432,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. @@ -3383,7 +3565,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) { @@ -3427,7 +3611,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 c60c8ab86..b282847e6 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 c125ba721..5ae0baea8 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 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) +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 +} + +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*\{/) - - expect(foreground).toContain('maxTurns,') - expect(background).toContain('maxTurns,') + test('honors OPENCLAUDE_MAX_TURNS when no explicit cap is passed', () => { + clearTurnEnv() + process.env.OPENCLAUDE_MAX_TURNS = '200' + expect(resolveReplMaxTurns()).toBe(200) }) - test('passes the cap from the resume selector into REPL', () => { - const source = readScreen('ResumeConversation.tsx') - const repl = source.slice(source.indexOf('', source.indexOf(' { + clearTurnEnv() + process.env.CLAUDE_CODE_MAX_TURNS = '125' + expect(resolveReplMaxTurns()).toBe(125) + }) - expect(repl).toContain('maxTurns={maxTurns}') + 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') + + await expect(wait).resolves.toBeNull() + }) + + test('defers max-turn cap restoration when foreground suppresses it for background handoff', () => { + 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', + ) + } + }) + + 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 8f63f50d7..a975dd445 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 3a778cb65..91d2dd764 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 000000000..755872097 --- /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 57763f3b3..0c87ab76d 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 - // 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 restorePreDispatchNotifications = (): void => { + if (providerRequestStarted) return + const restore = restoreNotificationsIfUnsent + restoreNotificationsIfUnsent = undefined + restore?.() + } + + 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 e92275477..3468d4add 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 cec4c1d11..edf3c1309 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 000000000..63eaaaf29 --- /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 1db54b53d..587901af9 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 000000000..e5cd97e74 --- /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 000000000..f190559f1 --- /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 000000000..a6bd00e47 --- /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 + }) +}