From 5d9dc0c4d52da1f5c606aefd689d397c449f5bcb Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 30 Apr 2026 15:23:33 +0700 Subject: [PATCH] fix(claude): preserve session across query restarts Fixes the Claude query restart path that could drop conversation context after thinking effort changes or rewind-driven restarts. Addresses #439. Related to #156 and #200. --- ...agent.interrupt-restart-regression.test.ts | 46 +++++++ .../providers/claude-agent.redesign.test.ts | 10 +- .../server/agent/providers/claude-agent.ts | 76 ++++++++---- .../claude-thinking-memory.real.e2e.test.ts | 115 ++++++++++++++++++ 4 files changed, 222 insertions(+), 25 deletions(-) create mode 100644 packages/server/src/server/daemon-e2e/claude-thinking-memory.real.e2e.test.ts diff --git a/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts b/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts index 48a727f4c..4ad86767a 100644 --- a/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts +++ b/packages/server/src/server/agent/providers/claude-agent.interrupt-restart-regression.test.ts @@ -338,6 +338,52 @@ test("reuses the existing query after interrupt before starting the next prompt" await session.close(); }); +test("emits an assistant system notice when Claude changes session id mid-turn", async () => { + const logger = createTestLogger(); + let queryRef: ScriptedQuery | null = null; + + sdkMocks.query.mockImplementation(({ prompt }: { prompt: AsyncIterable }) => { + queryRef = createScriptedQuery({ + prompt, + sessionId: "claude-original-session", + async handlePrompt({ promptRecord, query }) { + if (promptRecord.text !== "trigger session switch") { + return; + } + query.emit({ + type: "assistant", + message: { content: "Claude kept working." }, + session_id: "claude-provider-switched-session", + }); + query.emit(buildSuccessResult("claude-provider-switched-session")); + }, + }); + return queryRef; + }); + + const client = new ClaudeAgentClient({ logger }); + const session = await client.createSession({ + provider: "claude", + cwd: process.cwd(), + }); + + const events = await collectUntilTerminal(streamSession(session, "trigger session switch")); + const assistantMessages = events.flatMap((event) => { + if (event.type !== "timeline" || event.item.type !== "assistant_message") { + return []; + } + return [event.item.text]; + }); + + expect(session.id).toBe("claude-provider-switched-session"); + expect(assistantMessages).toContain("Claude kept working."); + expect(assistantMessages).toContain( + "Claude switched to a new session: claude-original-session -> claude-provider-switched-session", + ); + + await session.close(); +}); + test("recovers when the query pump sees a single interrupt abort before the next prompt", async () => { const logger = createTestLogger(); const output = createAsyncQueue>(); diff --git a/packages/server/src/server/agent/providers/claude-agent.redesign.test.ts b/packages/server/src/server/agent/providers/claude-agent.redesign.test.ts index 432081f51..677f339e0 100644 --- a/packages/server/src/server/agent/providers/claude-agent.redesign.test.ts +++ b/packages/server/src/server/agent/providers/claude-agent.redesign.test.ts @@ -412,7 +412,10 @@ test("captures session IDs from fixture-driven init message variants", async () fixtures.map(async (fixture) => { const session = await createSession(); const internal = session as unknown as { - handleSystemMessage: (message: Record) => string | null; + handleSystemMessage: (message: Record) => { + threadStartedSessionId: string | null; + notice: AgentTimelineItem | null; + }; }; try { const started = internal.handleSystemMessage({ @@ -422,7 +425,10 @@ test("captures session IDs from fixture-driven init message variants", async () model: "opus", ...fixture.payload, }); - expect(started).toBe(fixture.expected); + expect(started).toEqual({ + threadStartedSessionId: fixture.expected, + notice: null, + }); expect(session.describePersistence()?.sessionId).toBe(fixture.expected); } finally { await session.close(); diff --git a/packages/server/src/server/agent/providers/claude-agent.ts b/packages/server/src/server/agent/providers/claude-agent.ts index adb000054..833838d82 100644 --- a/packages/server/src/server/agent/providers/claude-agent.ts +++ b/packages/server/src/server/agent/providers/claude-agent.ts @@ -2152,9 +2152,6 @@ class ClaudeAgentSession implements AgentSession { this.input = null; this.queryPumpPromise = null; this.queryRestartNeeded = false; - // Reset session identity for explicit restarts so the new query starts - // a fresh session rather than resuming the previous one. - this.claudeSessionId = null; oldInput?.end(); oldQuery.close?.(); try { @@ -2164,10 +2161,8 @@ class ClaudeAgentSession implements AgentSession { } } - // When the pump died unexpectedly (query became null, e.g. after a session - // ID overwrite error), preserve claudeSessionId so buildOptions() passes - // resume: sessionId and the new query auto-resumes the previous session. - // For explicit restarts above, claudeSessionId was already cleared. + // Preserve claudeSessionId across query recreation so buildOptions() passes + // resume: sessionId and the new query continues the existing conversation. this.persistence = null; const input = createAsyncMessageInput(); @@ -2840,7 +2835,6 @@ class ClaudeAgentSession implements AgentSession { this.query = null; this.input = null; } - this.claudeSessionId = null; this.persistence = null; this.persistedHistory = []; this.historyPending = false; @@ -2885,12 +2879,19 @@ class ClaudeAgentSession implements AgentSession { } const events: AgentStreamEvent[] = []; - const fallbackThreadSessionId = this.captureSessionIdFromMessage(message); - if (fallbackThreadSessionId) { + const sessionCapture = this.captureSessionIdFromMessage(message); + if (sessionCapture.notice) { + events.push({ + type: "timeline", + provider: "claude", + item: sessionCapture.notice, + }); + } + if (sessionCapture.threadStartedSessionId) { events.push({ type: "thread_started", provider: "claude", - sessionId: fallbackThreadSessionId, + sessionId: sessionCapture.threadStartedSessionId, }); } @@ -2929,12 +2930,19 @@ class ClaudeAgentSession implements AgentSession { events: AgentStreamEvent[], ): void { if (message.subtype === "init") { - const threadSessionId = this.handleSystemMessage(message); - if (threadSessionId) { + const sessionUpdate = this.handleSystemMessage(message); + if (sessionUpdate.notice) { + events.push({ + type: "timeline", + provider: "claude", + item: sessionUpdate.notice, + }); + } + if (sessionUpdate.threadStartedSessionId) { events.push({ type: "thread_started", provider: "claude", - sessionId: threadSessionId, + sessionId: sessionUpdate.threadStartedSessionId, }); } return; @@ -3091,7 +3099,20 @@ class ClaudeAgentSession implements AgentSession { events.push(this.buildTurnFailedEvent(errorMessage)); } - private captureSessionIdFromMessage(message: SDKMessage): string | null { + private createClaudeSessionChangedNotice( + oldSessionId: string, + newSessionId: string, + ): AgentTimelineItem { + return { + type: "assistant_message", + text: `Claude switched to a new session: ${oldSessionId} -> ${newSessionId}`, + }; + } + + private captureSessionIdFromMessage(message: SDKMessage): { + threadStartedSessionId: string | null; + notice: AgentTimelineItem | null; + } { const msg = message as unknown as { session_id?: unknown; sessionId?: unknown; @@ -3099,16 +3120,17 @@ class ClaudeAgentSession implements AgentSession { }; const sessionId = extractSessionIdRaw(msg).trim(); if (!sessionId) { - return null; + return { threadStartedSessionId: null, notice: null }; } if (this.claudeSessionId === null) { this.claudeSessionId = sessionId; this.persistence = null; - return sessionId; + return { threadStartedSessionId: sessionId, notice: null }; } if (this.claudeSessionId === sessionId) { - return null; + return { threadStartedSessionId: null, notice: null }; } + const oldSessionId = this.claudeSessionId; // Session ID changed mid-stream (e.g. a hook caused Claude to restart // with a new session). Accept the new ID and continue — the turn should // not be failed just because the underlying subprocess cycled. @@ -3118,12 +3140,18 @@ class ClaudeAgentSession implements AgentSession { ); this.claudeSessionId = sessionId; this.persistence = null; - return sessionId; + return { + threadStartedSessionId: sessionId, + notice: this.createClaudeSessionChangedNotice(oldSessionId, sessionId), + }; } - private handleSystemMessage(message: SDKSystemMessage): string | null { + private handleSystemMessage(message: SDKSystemMessage): { + threadStartedSessionId: string | null; + notice: AgentTimelineItem | null; + } { if (message.subtype !== "init") { - return null; + return { threadStartedSessionId: null, notice: null }; } const msg = message as unknown as { @@ -3133,10 +3161,11 @@ class ClaudeAgentSession implements AgentSession { }; const newSessionId = extractSessionIdRaw(msg).trim(); if (!newSessionId) { - return null; + return { threadStartedSessionId: null, notice: null }; } const existingSessionId = this.claudeSessionId; let threadStartedSessionId: string | null = null; + let notice: AgentTimelineItem | null = null; if (existingSessionId === null) { this.claudeSessionId = newSessionId; @@ -3153,6 +3182,7 @@ class ClaudeAgentSession implements AgentSession { ); this.claudeSessionId = newSessionId; threadStartedSessionId = newSessionId; + notice = this.createClaudeSessionChangedNotice(existingSessionId, newSessionId); } this.availableModes = DEFAULT_MODES; this.currentMode = message.permissionMode; @@ -3174,7 +3204,7 @@ class ClaudeAgentSession implements AgentSession { this.lastRuntimeModel = message.model; this.cachedRuntimeInfo = null; } - return threadStartedSessionId; + return { threadStartedSessionId, notice }; } private readMissingResumedConversationError(message: SDKMessage): string | null { diff --git a/packages/server/src/server/daemon-e2e/claude-thinking-memory.real.e2e.test.ts b/packages/server/src/server/daemon-e2e/claude-thinking-memory.real.e2e.test.ts new file mode 100644 index 000000000..fa908b68c --- /dev/null +++ b/packages/server/src/server/daemon-e2e/claude-thinking-memory.real.e2e.test.ts @@ -0,0 +1,115 @@ +import { beforeAll, beforeEach, describe, expect, test } from "vitest"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import pino from "pino"; + +import type { AgentTimelineItem } from "../agent/agent-sdk-types.js"; +import { ClaudeAgentClient } from "../agent/providers/claude-agent.js"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; +import { getFullAccessConfig, isProviderAvailable } from "./agent-configs.js"; + +function tmpCwd(): string { + return mkdtempSync(path.join(tmpdir(), "daemon-real-claude-thinking-memory-")); +} + +function compactText(value: string): string { + return value.replace(/\s+/g, "").toLowerCase(); +} + +function assistantMessages(items: AgentTimelineItem[]): string[] { + return items + .filter( + (item): item is Extract => + item.type === "assistant_message", + ) + .map((item) => item.text); +} + +function modelHasLowThinkingOption(model: { + id: string; + thinkingOptions?: Array<{ id: string }>; +}): boolean { + const thinkingOptionIds = new Set(model.thinkingOptions?.map((option) => option.id) ?? []); + return model.id.includes("sonnet") && thinkingOptionIds.has("low"); +} + +async function getAssistantText(client: DaemonClient, agentId: string): Promise { + const timeline = await client.fetchAgentTimeline(agentId, { + direction: "tail", + limit: 0, + projection: "canonical", + }); + return assistantMessages(timeline.entries.map((entry) => entry.item)).join("\n"); +} + +describe("daemon E2E (real claude) - thinking effort memory", () => { + let canRun = false; + + beforeAll(async () => { + canRun = await isProviderAvailable("claude"); + }); + + beforeEach((context) => { + if (!canRun) { + context.skip(); + } + }); + + test("changing thinking effort preserves the previous conversation", async () => { + const logger = pino({ level: "silent" }); + const cwd = tmpCwd(); + const daemon = await createTestPaseoDaemon({ + agentClients: { claude: new ClaudeAgentClient({ logger }) }, + logger, + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + + try { + await client.connect(); + await client.fetchAgents({ + subscribe: { subscriptionId: "claude-thinking-memory-real" }, + }); + + const modelList = await client.listProviderModels("claude", { cwd }); + const model = modelList.models.find(modelHasLowThinkingOption); + if (!model) { + throw new Error("No Claude Sonnet model with low thinking effort returned"); + } + + const agent = await client.createAgent({ + cwd, + title: "claude-thinking-memory-real", + ...getFullAccessConfig("claude"), + model: model.id, + }); + + await client.sendMessage( + agent.id, + "Remember the code phrase PASEO_MEMORY_56. Reply exactly: ACK_56", + ); + const firstFinish = await client.waitForFinish(agent.id, 180_000); + expect(firstFinish.status).toBe("idle"); + expect(firstFinish.final?.lastError).toBeUndefined(); + expect(compactText(await getAssistantText(client, agent.id))).toContain("ack_56"); + + await client.setAgentThinkingOption(agent.id, "low"); + + await client.sendMessage( + agent.id, + "What code phrase did I ask you to remember? Reply exactly with that code phrase and nothing else.", + ); + const secondFinish = await client.waitForFinish(agent.id, 180_000); + expect(secondFinish.status).toBe("idle"); + expect(secondFinish.final?.lastError).toBeUndefined(); + + const assistantText = await getAssistantText(client, agent.id); + expect(compactText(assistantText)).toContain("paseo_memory_56"); + } finally { + await client.close().catch(() => undefined); + await daemon.close().catch(() => undefined); + rmSync(cwd, { recursive: true, force: true }); + } + }, 420_000); +});