From cdbda3d2c916c857536b77ebcd7ac6eb4c7d905a Mon Sep 17 00:00:00 2001 From: Merge Sim Date: Wed, 2 Sep 2026 10:30:39 -0700 Subject: [PATCH] fix(claude): preserve SDK preset and close ownership --- .../claude-stream-json-connection.test.ts | 33 ++++++++++++++----- .../claude/claude-stream-json-connection.ts | 5 +-- ...laude-structured-launch-resolution.test.ts | 1 + .../claude-structured-launch-resolution.ts | 4 +++ ...tured-agent-session-adapter-router.test.ts | 33 +++++++++++++++++++ ...structured-agent-session-adapter-router.ts | 6 +++- 6 files changed, 71 insertions(+), 11 deletions(-) create mode 100644 src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.test.ts diff --git a/src/main/claude/claude-stream-json-connection.test.ts b/src/main/claude/claude-stream-json-connection.test.ts index 5e5125ac461..fae1b2ae05a 100644 --- a/src/main/claude/claude-stream-json-connection.test.ts +++ b/src/main/claude/claude-stream-json-connection.test.ts @@ -4,7 +4,7 @@ import { join } from 'node:path' import { afterEach, describe, expect, it, vi } from 'vitest' import { spawnProcess, type SpawnedProcess } from '../../shared/child-process/run-process' import type { ProcessSpec } from '../../shared/child-process/process-spec' -import type { CanUseTool } from '@anthropic-ai/claude-agent-sdk' +import { query, type CanUseTool, type Options } from '@anthropic-ai/claude-agent-sdk' import { openClaudeStreamJsonConnection, type ClaudeStreamJsonConnection, @@ -85,14 +85,20 @@ const spawnedChildren: SpawnedProcess[] = [] async function open( launch: ClaudeStreamJsonLaunch, - handlers: Parameters[1] = {} + handlers: Parameters[1] = {}, + queryImpl?: typeof query ): Promise { - const connection = await openClaudeStreamJsonConnection(launch, handlers, (spec) => { - spawned.push(spec) - const child = spawnProcess(spec) - spawnedChildren.push(child) - return child - }) + const connection = await openClaudeStreamJsonConnection( + launch, + handlers, + (spec) => { + spawned.push(spec) + const child = spawnProcess(spec) + spawnedChildren.push(child) + return child + }, + queryImpl + ) openConnections.push(connection) return connection } @@ -121,6 +127,17 @@ function readReportSafely(scenario: { readReport: () => ScriptedCliReport }) { } describe('Claude stream-json connection', () => { + it('passes the Claude Code system-prompt preset through to SDK query', async () => { + const scenario = scriptScenario([HOLD_OPEN]) + let captured: Options | undefined + await open(launchFor(scenario), {}, (params) => { + captured = params.options + return query(params) + }) + + expect(captured?.systemPrompt).toEqual({ type: 'preset', preset: 'claude_code' }) + }) + it('hands the child a derived environment, the resolved CLI path, and keeps the pid', async () => { vi.stubEnv('ANTHROPIC_API_KEY', 'sk-ant-SHELL-LEAK') vi.stubEnv('CLAUDE_CODE_CHILD_SESSION', '1') diff --git a/src/main/claude/claude-stream-json-connection.ts b/src/main/claude/claude-stream-json-connection.ts index f90853d7e19..ceb6c10933c 100644 --- a/src/main/claude/claude-stream-json-connection.ts +++ b/src/main/claude/claude-stream-json-connection.ts @@ -91,12 +91,13 @@ function exitError(stderrTail: string, status: ExitStatus | null, cause?: Error) export async function openClaudeStreamJsonConnection( launch: ClaudeStreamJsonLaunch, handlers: ClaudeStreamJsonConnectionHandlers = {}, - spawnImpl: typeof spawnProcess = spawnProcess + spawnImpl: typeof spawnProcess = spawnProcess, + queryImpl?: typeof ClaudeAgentSdk.query ): Promise { const { query } = await loadClaudeAgentSdk() const spawner = createClaudeCodeProcessSpawn(spawnImpl) const inbox = createClaudeUserMessageQueue() - const session = query({ + const session = (queryImpl ?? query)({ prompt: inbox.messages, options: { ...launch.options, diff --git a/src/main/claude/claude-structured-launch-resolution.test.ts b/src/main/claude/claude-structured-launch-resolution.test.ts index b8959d9e881..4e99a19dd31 100644 --- a/src/main/claude/claude-structured-launch-resolution.test.ts +++ b/src/main/claude/claude-structured-launch-resolution.test.ts @@ -69,6 +69,7 @@ describe('claude structured launch resolution', () => { includePartialMessages: true, settingSources: [...CLAUDE_DEFAULT_SETTING_SOURCES], supportedDialogKinds: [], + systemPrompt: { type: 'preset', preset: 'claude_code' }, extraArgs: { 'replay-user-messages': null }, sessionId: first.providerSessionId }) diff --git a/src/main/claude/claude-structured-launch-resolution.ts b/src/main/claude/claude-structured-launch-resolution.ts index fda0debe2f6..3476e472332 100644 --- a/src/main/claude/claude-structured-launch-resolution.ts +++ b/src/main/claude/claude-structured-launch-resolution.ts @@ -16,6 +16,7 @@ export type ClaudeStructuredSdkOptions = Pick< | 'settingSources' | 'supportedDialogKinds' | 'extraArgs' + | 'systemPrompt' | 'model' | 'effort' | 'sessionId' @@ -34,6 +35,9 @@ export const CLAUDE_STRUCTURED_BASE_OPTIONS: ClaudeStructuredSdkOptions = { includePartialMessages: true, settingSources: [...CLAUDE_DEFAULT_SETTING_SOURCES], supportedDialogKinds: [], + // Keep the SDK on Claude Code's own system-prompt contract. Without this + // preset, query() can run with a generic/empty prompt under SDK defaults. + systemPrompt: { type: 'preset', preset: 'claude_code' }, extraArgs: { 'replay-user-messages': null } } diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.test.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.test.ts new file mode 100644 index 00000000000..1356877bc32 --- /dev/null +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it, vi } from 'vitest' +import type { StructuredAgentSessionAdapter } from './structured-agent-session-adapter' +import { StructuredAgentSessionAdapterRouter } from './structured-agent-session-adapter-router' + +function adapterOf( + closeSession: StructuredAgentSessionAdapter['closeSession'] +): StructuredAgentSessionAdapter { + return { + acquire: vi.fn(async () => ({ process: { pid: 1 } }) as never), + closeSession, + dispatch: vi.fn(), + cancelTurn: vi.fn(), + answerPrompt: vi.fn(), + setOption: vi.fn() + } as unknown as StructuredAgentSessionAdapter +} + +describe('StructuredAgentSessionAdapterRouter.closeSession', () => { + it('retains the owner when provider exit is unproven so close can retry', async () => { + const closeSession = vi.fn().mockResolvedValueOnce(false).mockResolvedValueOnce(true) + const adapter = adapterOf(closeSession) + const router = new StructuredAgentSessionAdapterRouter( + { claude: adapter, codex: adapter }, + async () => {} + ) + const identity = { sessionId: 'session-1', agent: 'claude' } as never + await router.acquire({ identity, fence: 1, spawnToken: 'spawn-1' }) + + await expect(router.closeSession('session-1')).resolves.toBe(false) + await expect(router.closeSession('session-1')).resolves.toBe(true) + expect(closeSession).toHaveBeenCalledTimes(2) + }) +}) diff --git a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.ts b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.ts index d5968bbadef..d502b3d23b5 100644 --- a/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.ts +++ b/src/main/native-chat/agent-session-wire/structured-agent-session-adapter-router.ts @@ -70,7 +70,11 @@ export class StructuredAgentSessionAdapterRouter implements StructuredAgentSessi return false } const closed = await adapter.closeSession?.(sessionId) - this.owners.delete(sessionId) + // Keep routing ownership while the provider cannot prove its child exited; + // a retry must still reach the adapter that owns the live process. + if (closed === true) { + this.owners.delete(sessionId) + } return closed === true }