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.
This commit is contained in:
Mohamed Boudra
2026-04-30 15:23:33 +07:00
parent 0b78379b77
commit 5d9dc0c4d5
4 changed files with 222 additions and 25 deletions

View File

@@ -338,6 +338,52 @@ test("reuses the existing query after interrupt before starting the next prompt"
await session.close(); 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<unknown> }) => {
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 () => { test("recovers when the query pump sees a single interrupt abort before the next prompt", async () => {
const logger = createTestLogger(); const logger = createTestLogger();
const output = createAsyncQueue<Record<string, unknown>>(); const output = createAsyncQueue<Record<string, unknown>>();

View File

@@ -412,7 +412,10 @@ test("captures session IDs from fixture-driven init message variants", async ()
fixtures.map(async (fixture) => { fixtures.map(async (fixture) => {
const session = await createSession(); const session = await createSession();
const internal = session as unknown as { const internal = session as unknown as {
handleSystemMessage: (message: Record<string, unknown>) => string | null; handleSystemMessage: (message: Record<string, unknown>) => {
threadStartedSessionId: string | null;
notice: AgentTimelineItem | null;
};
}; };
try { try {
const started = internal.handleSystemMessage({ const started = internal.handleSystemMessage({
@@ -422,7 +425,10 @@ test("captures session IDs from fixture-driven init message variants", async ()
model: "opus", model: "opus",
...fixture.payload, ...fixture.payload,
}); });
expect(started).toBe(fixture.expected); expect(started).toEqual({
threadStartedSessionId: fixture.expected,
notice: null,
});
expect(session.describePersistence()?.sessionId).toBe(fixture.expected); expect(session.describePersistence()?.sessionId).toBe(fixture.expected);
} finally { } finally {
await session.close(); await session.close();

View File

@@ -2152,9 +2152,6 @@ class ClaudeAgentSession implements AgentSession {
this.input = null; this.input = null;
this.queryPumpPromise = null; this.queryPumpPromise = null;
this.queryRestartNeeded = false; 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(); oldInput?.end();
oldQuery.close?.(); oldQuery.close?.();
try { try {
@@ -2164,10 +2161,8 @@ class ClaudeAgentSession implements AgentSession {
} }
} }
// When the pump died unexpectedly (query became null, e.g. after a session // Preserve claudeSessionId across query recreation so buildOptions() passes
// ID overwrite error), preserve claudeSessionId so buildOptions() passes // resume: sessionId and the new query continues the existing conversation.
// resume: sessionId and the new query auto-resumes the previous session.
// For explicit restarts above, claudeSessionId was already cleared.
this.persistence = null; this.persistence = null;
const input = createAsyncMessageInput<SDKUserMessage>(); const input = createAsyncMessageInput<SDKUserMessage>();
@@ -2840,7 +2835,6 @@ class ClaudeAgentSession implements AgentSession {
this.query = null; this.query = null;
this.input = null; this.input = null;
} }
this.claudeSessionId = null;
this.persistence = null; this.persistence = null;
this.persistedHistory = []; this.persistedHistory = [];
this.historyPending = false; this.historyPending = false;
@@ -2885,12 +2879,19 @@ class ClaudeAgentSession implements AgentSession {
} }
const events: AgentStreamEvent[] = []; const events: AgentStreamEvent[] = [];
const fallbackThreadSessionId = this.captureSessionIdFromMessage(message); const sessionCapture = this.captureSessionIdFromMessage(message);
if (fallbackThreadSessionId) { if (sessionCapture.notice) {
events.push({
type: "timeline",
provider: "claude",
item: sessionCapture.notice,
});
}
if (sessionCapture.threadStartedSessionId) {
events.push({ events.push({
type: "thread_started", type: "thread_started",
provider: "claude", provider: "claude",
sessionId: fallbackThreadSessionId, sessionId: sessionCapture.threadStartedSessionId,
}); });
} }
@@ -2929,12 +2930,19 @@ class ClaudeAgentSession implements AgentSession {
events: AgentStreamEvent[], events: AgentStreamEvent[],
): void { ): void {
if (message.subtype === "init") { if (message.subtype === "init") {
const threadSessionId = this.handleSystemMessage(message); const sessionUpdate = this.handleSystemMessage(message);
if (threadSessionId) { if (sessionUpdate.notice) {
events.push({
type: "timeline",
provider: "claude",
item: sessionUpdate.notice,
});
}
if (sessionUpdate.threadStartedSessionId) {
events.push({ events.push({
type: "thread_started", type: "thread_started",
provider: "claude", provider: "claude",
sessionId: threadSessionId, sessionId: sessionUpdate.threadStartedSessionId,
}); });
} }
return; return;
@@ -3091,7 +3099,20 @@ class ClaudeAgentSession implements AgentSession {
events.push(this.buildTurnFailedEvent(errorMessage)); 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 { const msg = message as unknown as {
session_id?: unknown; session_id?: unknown;
sessionId?: unknown; sessionId?: unknown;
@@ -3099,16 +3120,17 @@ class ClaudeAgentSession implements AgentSession {
}; };
const sessionId = extractSessionIdRaw(msg).trim(); const sessionId = extractSessionIdRaw(msg).trim();
if (!sessionId) { if (!sessionId) {
return null; return { threadStartedSessionId: null, notice: null };
} }
if (this.claudeSessionId === null) { if (this.claudeSessionId === null) {
this.claudeSessionId = sessionId; this.claudeSessionId = sessionId;
this.persistence = null; this.persistence = null;
return sessionId; return { threadStartedSessionId: sessionId, notice: null };
} }
if (this.claudeSessionId === sessionId) { 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 // 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 // with a new session). Accept the new ID and continue — the turn should
// not be failed just because the underlying subprocess cycled. // not be failed just because the underlying subprocess cycled.
@@ -3118,12 +3140,18 @@ class ClaudeAgentSession implements AgentSession {
); );
this.claudeSessionId = sessionId; this.claudeSessionId = sessionId;
this.persistence = null; 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") { if (message.subtype !== "init") {
return null; return { threadStartedSessionId: null, notice: null };
} }
const msg = message as unknown as { const msg = message as unknown as {
@@ -3133,10 +3161,11 @@ class ClaudeAgentSession implements AgentSession {
}; };
const newSessionId = extractSessionIdRaw(msg).trim(); const newSessionId = extractSessionIdRaw(msg).trim();
if (!newSessionId) { if (!newSessionId) {
return null; return { threadStartedSessionId: null, notice: null };
} }
const existingSessionId = this.claudeSessionId; const existingSessionId = this.claudeSessionId;
let threadStartedSessionId: string | null = null; let threadStartedSessionId: string | null = null;
let notice: AgentTimelineItem | null = null;
if (existingSessionId === null) { if (existingSessionId === null) {
this.claudeSessionId = newSessionId; this.claudeSessionId = newSessionId;
@@ -3153,6 +3182,7 @@ class ClaudeAgentSession implements AgentSession {
); );
this.claudeSessionId = newSessionId; this.claudeSessionId = newSessionId;
threadStartedSessionId = newSessionId; threadStartedSessionId = newSessionId;
notice = this.createClaudeSessionChangedNotice(existingSessionId, newSessionId);
} }
this.availableModes = DEFAULT_MODES; this.availableModes = DEFAULT_MODES;
this.currentMode = message.permissionMode; this.currentMode = message.permissionMode;
@@ -3174,7 +3204,7 @@ class ClaudeAgentSession implements AgentSession {
this.lastRuntimeModel = message.model; this.lastRuntimeModel = message.model;
this.cachedRuntimeInfo = null; this.cachedRuntimeInfo = null;
} }
return threadStartedSessionId; return { threadStartedSessionId, notice };
} }
private readMissingResumedConversationError(message: SDKMessage): string | null { private readMissingResumedConversationError(message: SDKMessage): string | null {

View File

@@ -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<AgentTimelineItem, { type: "assistant_message" }> =>
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<string> {
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);
});