Skip to content

Commit 1182fdd

Browse files
richardsolomourichardsolomou
authored
fix(sessions): recover and resend prompts after "session not found" (#3170)
1 parent 6867e5d commit 1182fdd

3 files changed

Lines changed: 228 additions & 9 deletions

File tree

‎packages/core/src/sessions/sessionService.ts‎

Lines changed: 82 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2473,7 +2473,7 @@ export class SessionService {
24732473
session: AgentSession,
24742474
blocks: ContentBlock[],
24752475
promptText: string,
2476-
options: { optimisticApplied?: boolean } = {},
2476+
options: { optimisticApplied?: boolean; isRecoveryResend?: boolean } = {},
24772477
): Promise<{ stopReason: string }> {
24782478
if (!options.optimisticApplied) {
24792479
this.applyOptimisticPrompt(session.taskRunId, blocks, promptText);
@@ -2515,14 +2515,33 @@ export class SessionService {
25152515
errorMessage,
25162516
errorDetails,
25172517
});
2518-
this.startAutoRecoverLocalSession(
2519-
session.taskId,
2520-
session.taskRunId,
2521-
session.taskTitle,
2522-
errorDetails || errorMessage,
2523-
errorDetails ||
2524-
"Session connection lost. Please retry or start a new session.",
2525-
);
2518+
if (!options.isRecoveryResend) {
2519+
const resent = await this.recoverAndResendPrompt(
2520+
session,
2521+
blocks,
2522+
promptText,
2523+
errorDetails || errorMessage,
2524+
);
2525+
if (resent) {
2526+
return resent;
2527+
}
2528+
}
2529+
// Recovery failed (or this already was the post-recovery resend):
2530+
// surface the error state so the user can retry manually.
2531+
const latest = this.d.store.getSessionByTaskId(session.taskId);
2532+
if (
2533+
latest?.taskRunId === session.taskRunId &&
2534+
latest.status !== "error"
2535+
) {
2536+
this.setErrorSession(
2537+
session.taskId,
2538+
session.taskRunId,
2539+
session.taskTitle,
2540+
errorDetails ||
2541+
"Session connection lost. Please retry or start a new session.",
2542+
"Connection lost",
2543+
);
2544+
}
25262545
} else {
25272546
this.d.store.updateSession(session.taskRunId, {
25282547
isPromptPending: false,
@@ -2535,6 +2554,60 @@ export class SessionService {
25352554
}
25362555
}
25372556

2557+
/**
2558+
* A fatal prompt failure (e.g. "Session not found") usually means the
2559+
* backend agent was idle-killed or the host process restarted while the
2560+
* renderer still shows the session as connected. Recover the session in
2561+
* place and resend the prompt once, so a reply to a stale session lands
2562+
* instead of erroring and losing the user's message.
2563+
*
2564+
* Returns the resend result, or null when recovery (or the refreshed
2565+
* session lookup) failed and the caller should surface the original error.
2566+
*/
2567+
private async recoverAndResendPrompt(
2568+
session: AgentSession,
2569+
blocks: ContentBlock[],
2570+
promptText: string,
2571+
reason: string,
2572+
): Promise<{ stopReason: string } | null> {
2573+
let recovered = false;
2574+
try {
2575+
recovered = await this.tryAutoRecoverLocalSession(
2576+
session.taskId,
2577+
session.taskRunId,
2578+
reason,
2579+
);
2580+
} catch (recoveryError) {
2581+
this.d.log.warn("Session recovery threw while resending prompt", {
2582+
taskId: session.taskId,
2583+
taskRunId: session.taskRunId,
2584+
error:
2585+
recoveryError instanceof Error
2586+
? recoveryError.message
2587+
: String(recoveryError),
2588+
});
2589+
return null;
2590+
}
2591+
if (!recovered) return null;
2592+
2593+
const refreshed = this.d.store.getSessionByTaskId(session.taskId);
2594+
if (
2595+
!refreshed ||
2596+
refreshed.taskRunId !== session.taskRunId ||
2597+
refreshed.status !== "connected"
2598+
) {
2599+
return null;
2600+
}
2601+
2602+
this.d.log.info("Resending prompt after session recovery", {
2603+
taskId: session.taskId,
2604+
taskRunId: session.taskRunId,
2605+
});
2606+
return this.sendLocalPrompt(refreshed, blocks, promptText, {
2607+
isRecoveryResend: true,
2608+
});
2609+
}
2610+
25382611
/**
25392612
* Steer a single queued message into the running turn now: drop it from the
25402613
* queue and resend it as a steer. Native (Claude, local) injects at the next

‎packages/core/src/sessions/sessionServicePromptRecovery.test.ts‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
import type { AgentSession } from "@posthog/shared";
2+
import { describe, expect, it, vi } from "vitest";
3+
import { SessionService, type SessionServiceDeps } from "./sessionService";
4+
5+
const TASK_ID = "task-1";
6+
const TASK_RUN_ID = `run-${TASK_ID}`;
7+
8+
function makeSession(): AgentSession {
9+
return {
10+
taskRunId: TASK_RUN_ID,
11+
taskId: TASK_ID,
12+
taskTitle: "Test task",
13+
channel: "",
14+
events: [],
15+
startedAt: 1,
16+
status: "connected",
17+
isPromptPending: false,
18+
isCompacting: false,
19+
promptStartedAt: null,
20+
pendingPermissions: new Map(),
21+
pausedDurationMs: 0,
22+
messageQueue: [],
23+
optimisticItems: [],
24+
} as AgentSession;
25+
}
26+
27+
function createHarness() {
28+
const sessions: Record<string, AgentSession> = {
29+
[TASK_RUN_ID]: makeSession(),
30+
};
31+
const store = {
32+
getSessions: () => sessions,
33+
getSessionByTaskId: (taskId: string) =>
34+
Object.values(sessions).find((s) => s.taskId === taskId),
35+
setSession: vi.fn((session: AgentSession) => {
36+
sessions[session.taskRunId] = session;
37+
}),
38+
updateSession: vi.fn((taskRunId: string, patch: Partial<AgentSession>) => {
39+
const existing = sessions[taskRunId];
40+
if (existing) sessions[taskRunId] = { ...existing, ...patch };
41+
}),
42+
appendOptimisticItem: vi.fn(),
43+
clearOptimisticItems: vi.fn(),
44+
};
45+
const promptMutate = vi.fn();
46+
const deps = {
47+
store,
48+
h: { extractSkillButtonId: () => undefined },
49+
log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
50+
toast: { error: vi.fn(), info: vi.fn() },
51+
track: vi.fn(),
52+
getIsOnline: () => true,
53+
addDirectoryDialog: { open: false },
54+
usageLimit: { show: vi.fn() },
55+
trpc: {
56+
agent: {
57+
prompt: { mutate: promptMutate },
58+
cancel: { mutate: vi.fn().mockResolvedValue(undefined) },
59+
onSessionIdleKilled: {
60+
subscribe: () => ({ unsubscribe: vi.fn() }),
61+
},
62+
},
63+
},
64+
} as unknown as SessionServiceDeps;
65+
66+
const service = new SessionService(deps);
67+
const recoverSpy = vi.spyOn(
68+
service as unknown as {
69+
tryAutoRecoverLocalSession: (
70+
taskId: string,
71+
taskRunId: string,
72+
reason: string,
73+
) => Promise<boolean>;
74+
},
75+
"tryAutoRecoverLocalSession",
76+
);
77+
return { service, sessions, store, promptMutate, recoverSpy };
78+
}
79+
80+
type RecoverSpy = ReturnType<typeof createHarness>["recoverSpy"];
81+
82+
describe("SessionService prompt recovery on fatal session errors", () => {
83+
it("recovers the session and resends the prompt once", async () => {
84+
const { service, promptMutate, recoverSpy } = createHarness();
85+
promptMutate
86+
.mockRejectedValueOnce(new Error(`Session not found: ${TASK_RUN_ID}`))
87+
.mockResolvedValueOnce({ stopReason: "end_turn" });
88+
recoverSpy.mockResolvedValue(true);
89+
90+
const result = await service.sendPrompt(TASK_ID, "hello again");
91+
92+
expect(result).toEqual({ stopReason: "end_turn" });
93+
expect(recoverSpy).toHaveBeenCalledTimes(1);
94+
expect(promptMutate).toHaveBeenCalledTimes(2);
95+
});
96+
97+
it.each([
98+
{
99+
case: "recovery fails",
100+
setupRecovery: (spy: RecoverSpy) => spy.mockResolvedValue(false),
101+
expectedPromptCalls: 1,
102+
},
103+
{
104+
case: "the resend hits another fatal error",
105+
setupRecovery: (spy: RecoverSpy) => spy.mockResolvedValue(true),
106+
expectedPromptCalls: 2,
107+
},
108+
{
109+
case: "recovery throws",
110+
setupRecovery: (spy: RecoverSpy) =>
111+
spy.mockRejectedValue(new Error("auth restoring")),
112+
expectedPromptCalls: 1,
113+
},
114+
])(
115+
"surfaces the error state and rethrows when $case",
116+
async ({ setupRecovery, expectedPromptCalls }) => {
117+
const { service, store, promptMutate, recoverSpy } = createHarness();
118+
promptMutate.mockRejectedValue(
119+
new Error(`Session not found: ${TASK_RUN_ID}`),
120+
);
121+
setupRecovery(recoverSpy);
122+
123+
await expect(service.sendPrompt(TASK_ID, "hello again")).rejects.toThrow(
124+
"Session not found",
125+
);
126+
expect(recoverSpy).toHaveBeenCalledTimes(1);
127+
expect(promptMutate).toHaveBeenCalledTimes(expectedPromptCalls);
128+
expect(store.setSession).toHaveBeenCalledWith(
129+
expect.objectContaining({ taskRunId: TASK_RUN_ID, status: "error" }),
130+
);
131+
},
132+
);
133+
});

‎packages/ui/src/features/sessions/hooks/useSessionCallbacks.ts‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
import {
2+
isContentEmpty,
3+
xmlToContent,
4+
} from "@posthog/core/message-editor/content";
15
import {
26
combineQueuedCloudPrompts,
37
promptToQueuedEditorContent,
@@ -106,6 +110,13 @@ export function useSessionCallbacks({
106110
} catch (error) {
107111
const message =
108112
error instanceof Error ? error.message : "Failed to send message";
113+
// The composer clears optimistically on submit, so a failed send
114+
// would otherwise lose the typed prompt. Restore it — unless the
115+
// user has already started typing a new message.
116+
if (isContentEmpty(useDraftStore.getState().drafts[taskId] ?? null)) {
117+
setPendingContent(taskId, xmlToContent(text));
118+
requestFocus(taskId);
119+
}
109120
toast.error(message);
110121
log.error("Failed to send prompt", error);
111122
}
@@ -119,6 +130,8 @@ export function useSessionCallbacks({
119130
sessionService,
120131
hostClient,
121132
messagingMode,
133+
setPendingContent,
134+
requestFocus,
122135
],
123136
);
124137

0 commit comments

Comments
 (0)