fix(claude): preserve prompt delivery and rewind identity across query restarts

This commit is contained in:
louis1618
2026-10-13 16:42:51 +02:00
committed by paseo-bot[bot]
parent 5805b776c2
commit d3c1cc54bb
2 changed files with 132 additions and 27 deletions
@@ -2,6 +2,7 @@ import { afterEach, expect, test, vi } from "vitest";
import { createTestLogger } from "../../../../test-utils/test-logger.js";
import { ClaudeAgentClient } from "./agent.js";
import { AgentManager } from "../../agent-manager.js";
import { streamSession } from "../test-utils/session-stream-adapter.js";
import type { AgentSession, AgentStreamEvent } from "../../agent-sdk-types.js";
@@ -12,6 +13,7 @@ interface QueryMock {
close: ReturnType<typeof vi.fn>;
setPermissionMode: ReturnType<typeof vi.fn>;
setModel: ReturnType<typeof vi.fn>;
applyFlagSettings: ReturnType<typeof vi.fn>;
supportedModels: ReturnType<typeof vi.fn>;
supportedCommands: ReturnType<typeof vi.fn>;
rewindFiles: ReturnType<typeof vi.fn>;
@@ -148,6 +150,7 @@ function createScriptedQuery(params: {
close: vi.fn(() => undefined),
setPermissionMode: vi.fn(async () => undefined),
setModel: vi.fn(async () => undefined),
applyFlagSettings: vi.fn(async () => undefined),
supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]),
supportedCommands: vi.fn(async () => []),
rewindFiles: vi.fn(async () => ({ canRewind: true })),
@@ -276,6 +279,126 @@ afterEach(() => {
queryFactory.mockReset();
});
test.each([false, true])(
"submits after a thinking change with no live query (previous query ended: %s)",
async (previousQueryEnded) => {
const queries: ScriptedQuery[] = [];
queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
const scriptedQuery = createScriptedQuery({
prompt,
sessionId: "thinking-restart-session",
handlePrompt: ({ query }) => query.emit(buildSuccessResult("thinking-restart-session")),
});
queries.push(scriptedQuery);
return scriptedQuery;
});
const session = await new ClaudeAgentClient({
logger: createTestLogger(),
queryFactory,
resolveBinary: async () => "/test/claude/bin",
}).createSession({ provider: "claude", cwd: process.cwd() });
try {
if (previousQueryEnded) {
await collectUntilTerminal(streamSession(session, "first prompt"));
queries[0].end();
// Let the pump retire the old query before changing thinking.
await new Promise((resolve) => setTimeout(resolve, 0));
}
await session.setThinkingOption!("medium");
const events = await collectUntilTerminal(
streamSession(session, "continue", { clientMessageId: "client-continue" }),
);
expect(events.filter((event) => event.type === "turn_failed")).toEqual([]);
expect(events.at(-1)?.type).toBe("turn_completed");
expect(queryFactory).toHaveBeenCalledTimes(previousQueryEnded ? 2 : 1);
expect(queries.at(-1)?.prompts.map((prompt) => prompt.text)).toEqual(["continue"]);
} finally {
await session.close();
}
},
);
test("publishes the submitted prompt identity before an immediate provider failure", async () => {
queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable<unknown> }) =>
createScriptedQuery({
prompt,
sessionId: "immediate-failure-session",
handlePrompt: ({ query }) =>
query.emit({
type: "result",
subtype: "error_during_execution",
session_id: "immediate-failure-session",
errors: ["session limit reached"],
}),
}),
);
const session = await new ClaudeAgentClient({
logger: createTestLogger(),
queryFactory,
resolveBinary: async () => "/test/claude/bin",
}).createSession({ provider: "claude", cwd: process.cwd() });
try {
const events = await collectUntilTerminal(
streamSession(session, "continue", { clientMessageId: "client-failed-prompt" }),
);
expect(events.at(-1)?.type).toBe("turn_failed");
expect(events).toContainEqual(
expect.objectContaining({
type: "timeline",
item: expect.objectContaining({
type: "user_message",
text: "continue",
clientMessageId: "client-failed-prompt",
messageId: expect.any(String),
}),
}),
);
} finally {
await session.close();
}
});
test("rewind resolves a failed prompt's client identity to its Claude checkpoint", async () => {
const queries: ScriptedQuery[] = [];
queryFactory.mockImplementation(({ prompt }: { prompt: AsyncIterable<unknown> }) => {
const query = createScriptedQuery({
prompt,
sessionId: "failed-prompt-rewind-session",
handlePrompt: ({ query: activeQuery }) =>
activeQuery.emit({
type: "result",
subtype: "error_during_execution",
session_id: "failed-prompt-rewind-session",
errors: ["session limit reached"],
}),
});
queries.push(query);
return query;
});
const client = new ClaudeAgentClient({
logger: createTestLogger(),
queryFactory,
resolveBinary: async () => "/test/claude/bin",
});
const manager = new AgentManager({ clients: { claude: client }, logger: createTestLogger() });
const agent = await manager.createAgent({ provider: "claude", cwd: process.cwd() }, undefined, {
workspaceId: undefined,
});
try {
await collectUntilTerminal(
manager.streamAgent(agent.id, "continue", { clientMessageId: "client-failed-prompt" }),
);
expect(manager.getAgent(agent.id)?.lastError).toBe("session limit reached");
await manager.rewind(agent.id, "client-failed-prompt", "files");
expect(queries[0].rewindFiles).toHaveBeenCalledWith(queries[0].prompts[0].uuid, {
dryRun: false,
});
} finally {
await manager.closeAgent(agent.id);
}
});
test("interrupt only calls query.interrupt and leaves the query open", async () => {
const logger = createTestLogger();
const queries: ScriptedQuery[] = [];
@@ -2403,13 +2403,9 @@ class ClaudeAgentSession implements AgentSession {
}
this.activeForegroundQuery = this.query;
this.activeForegroundInput = this.input;
this.startQueryPump();
this.input.push(sdkMessage);
setTimeout(() => {
if (this.activeForegroundTurnId === turnId) {
this.emitSubmittedUserMessage(sdkMessage, turnId, options?.clientMessageId);
}
}, 0);
this.emitSubmittedUserMessage(sdkMessage, turnId, options?.clientMessageId);
this.startQueryPump();
} catch (error) {
if (sdkUserMessageId) this.unstartedMessageUuids.delete(sdkUserMessageId);
this.finishForegroundTurn(
@@ -3275,6 +3271,9 @@ class ClaudeAgentSession implements AgentSession {
}
}
// A fresh query already uses the current settings, even when there was no old query to retire.
this.queryRestartNeeded = false;
// Preserve claudeSessionId across query recreation so buildOptions() passes
// resume: sessionId and the new query continues the existing conversation.
this.persistence = null;
@@ -3877,11 +3876,12 @@ class ClaudeAgentSession implements AgentSession {
}
private startQueryPump(): void {
if (this.closed || this.queryPumpPromise) {
if (this.closed || this.queryPumpPromise || !this.query) {
return;
}
const pump = this.runQueryPump().catch((error) => {
// Drain the initialized query; a pending settings change belongs to the next turn.
const pump = this.runQueryPump(this.query).catch((error) => {
this.logger.trace(
{
agentId: this.agentId,
@@ -3902,25 +3902,7 @@ class ClaudeAgentSession implements AgentSession {
});
}
private async runQueryPump(): Promise<void> {
let activeQuery: Query;
try {
activeQuery = await this.ensureQuery();
} catch (error) {
this.logger.trace(
{
agentId: this.agentId,
provider: "claude",
sessionId: this.claudeSessionId,
turnId: this.activeForegroundTurnId ?? this.autonomousTurn?.id ?? undefined,
err: error,
},
"provider.claude.query_pump.init_failed",
);
this.failActiveTurns(error instanceof Error ? error.message : "Claude stream failed");
return;
}
private async runQueryPump(activeQuery: Query): Promise<void> {
let consecutiveInterruptAbortRecoveries = 0;
const logRawMessage = (message: SDKMessage): void => {
this.logger.trace(