Skip to content
This repository was archived by the owner on Aug 6, 2026. It is now read-only.

Commit e4d21d3

Browse files
committed
fix(code): never lose a local task's initial prompt
The initial prompt for a local task was only held in memory, so if the agent hadn't produced its first response yet and the user looked away (app backgrounded, reloaded, crashed, or a transient connect failure exhausted its silent retries), the prompt was lost and the task never started. Persist the prompt durably in the workspace-server task_metadata table keyed by task id the moment the task run is created, re-send it on resume when the agent hasn't consumed it (detected via the session/prompt echo in the replayed log), and clear it once consumed. clearSessionError now falls back to the durable copy so Retry and auto-retry recover after a reload wiped the in-memory session. Cloud tasks are unaffected — they self-fetch their prompt server-side. Generated-By: PostHog Code Task-Id: a4d03ce4-bd03-46db-8893-8a110e54e865
1 parent b529205 commit e4d21d3

13 files changed

Lines changed: 1560 additions & 14 deletions

File tree

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,208 @@
1+
import type { ContentBlock } from "@agentclientprotocol/sdk";
2+
import type { AcpMessage, AgentSession } from "@posthog/shared";
3+
import { describe, expect, it, vi } from "vitest";
4+
import { SessionService, type SessionServiceDeps } from "./sessionService";
5+
6+
const TASK_ID = "task-1";
7+
const RUN_ID = "run-1";
8+
9+
const PROMPT: ContentBlock[] = [{ type: "text", text: "do the thing" }];
10+
11+
function promptEcho(): AcpMessage {
12+
return {
13+
type: "acp_message",
14+
ts: 0,
15+
message: {
16+
jsonrpc: "2.0",
17+
id: 1,
18+
method: "session/prompt",
19+
params: {},
20+
},
21+
} as unknown as AcpMessage;
22+
}
23+
24+
function createHarness(
25+
overrides: {
26+
session?: AgentSession | null;
27+
getPendingInitialPrompt?: string | null;
28+
} = {},
29+
) {
30+
const sessions: Record<string, AgentSession> = {};
31+
if (overrides.session)
32+
sessions[overrides.session.taskRunId] = overrides.session;
33+
34+
const setPendingInitialPrompt = vi.fn().mockResolvedValue(undefined);
35+
const getPendingInitialPrompt = vi
36+
.fn()
37+
.mockResolvedValue(overrides.getPendingInitialPrompt ?? null);
38+
const clearPendingInitialPrompt = vi.fn().mockResolvedValue(undefined);
39+
40+
const store = {
41+
getSessions: () => sessions,
42+
getSessionByTaskId: (taskId: string) =>
43+
Object.values(sessions).find((s) => s.taskId === taskId),
44+
setSession: (session: AgentSession) => {
45+
sessions[session.taskRunId] = session;
46+
},
47+
updateSession: (taskRunId: string, updates: Partial<AgentSession>) => {
48+
const session = sessions[taskRunId];
49+
if (session) Object.assign(session, updates);
50+
},
51+
replaceOptimisticWithEvent: vi.fn(),
52+
appendEvents: vi.fn(),
53+
};
54+
55+
const deps = {
56+
store,
57+
log: { info: vi.fn(), warn: vi.fn(), error: vi.fn(), debug: vi.fn() },
58+
trpc: {
59+
agent: {
60+
onSessionIdleKilled: { subscribe: () => ({ unsubscribe: vi.fn() }) },
61+
},
62+
workspace: {
63+
setPendingInitialPrompt: { mutate: setPendingInitialPrompt },
64+
getPendingInitialPrompt: { query: getPendingInitialPrompt },
65+
clearPendingInitialPrompt: { mutate: clearPendingInitialPrompt },
66+
},
67+
},
68+
} as unknown as SessionServiceDeps;
69+
70+
const service = new SessionService(deps);
71+
72+
return {
73+
service,
74+
setPendingInitialPrompt,
75+
getPendingInitialPrompt,
76+
clearPendingInitialPrompt,
77+
};
78+
}
79+
80+
describe("initial prompt persistence", () => {
81+
describe("resendPendingPromptIfNeeded", () => {
82+
it("clears the durable prompt without resending when the log already has the echo", async () => {
83+
const h = createHarness({
84+
getPendingInitialPrompt: JSON.stringify(PROMPT),
85+
});
86+
const sendPrompt = vi
87+
.spyOn(h.service, "sendPrompt")
88+
.mockResolvedValue({ stopReason: "end_turn" });
89+
90+
await (
91+
h.service as unknown as {
92+
resendPendingPromptIfNeeded: (
93+
taskId: string,
94+
events: AcpMessage[],
95+
) => Promise<void>;
96+
}
97+
).resendPendingPromptIfNeeded(TASK_ID, [promptEcho()]);
98+
99+
expect(h.clearPendingInitialPrompt).toHaveBeenCalledWith({
100+
taskId: TASK_ID,
101+
});
102+
expect(h.getPendingInitialPrompt).not.toHaveBeenCalled();
103+
expect(sendPrompt).not.toHaveBeenCalled();
104+
});
105+
106+
it("resends the stored prompt exactly once when the log lacks the echo", async () => {
107+
const h = createHarness({
108+
getPendingInitialPrompt: JSON.stringify(PROMPT),
109+
});
110+
const sendPrompt = vi
111+
.spyOn(h.service, "sendPrompt")
112+
.mockResolvedValue({ stopReason: "end_turn" });
113+
114+
await (
115+
h.service as unknown as {
116+
resendPendingPromptIfNeeded: (
117+
taskId: string,
118+
events: AcpMessage[],
119+
) => Promise<void>;
120+
}
121+
).resendPendingPromptIfNeeded(TASK_ID, []);
122+
123+
expect(sendPrompt).toHaveBeenCalledTimes(1);
124+
expect(sendPrompt).toHaveBeenCalledWith(TASK_ID, PROMPT);
125+
});
126+
127+
it("does nothing when there is no stored prompt", async () => {
128+
const h = createHarness({ getPendingInitialPrompt: null });
129+
const sendPrompt = vi
130+
.spyOn(h.service, "sendPrompt")
131+
.mockResolvedValue({ stopReason: "end_turn" });
132+
133+
await (
134+
h.service as unknown as {
135+
resendPendingPromptIfNeeded: (
136+
taskId: string,
137+
events: AcpMessage[],
138+
) => Promise<void>;
139+
}
140+
).resendPendingPromptIfNeeded(TASK_ID, []);
141+
142+
expect(sendPrompt).not.toHaveBeenCalled();
143+
});
144+
});
145+
146+
describe("handleSessionEvent", () => {
147+
it("clears the durable prompt on the prompt echo", () => {
148+
const session = {
149+
taskRunId: RUN_ID,
150+
taskId: TASK_ID,
151+
events: [],
152+
messageQueue: [],
153+
optimisticItems: [],
154+
} as unknown as AgentSession;
155+
const h = createHarness({ session });
156+
157+
(
158+
h.service as unknown as {
159+
handleSessionEvent: (runId: string, msg: AcpMessage) => void;
160+
}
161+
).handleSessionEvent(RUN_ID, promptEcho());
162+
163+
expect(h.clearPendingInitialPrompt).toHaveBeenCalledWith({
164+
taskId: TASK_ID,
165+
});
166+
});
167+
});
168+
169+
describe("clearSessionError", () => {
170+
it("recovers the durable prompt when the in-memory session is gone", async () => {
171+
const h = createHarness({
172+
session: null,
173+
getPendingInitialPrompt: JSON.stringify(PROMPT),
174+
});
175+
const createNewLocalSession = vi
176+
.spyOn(
177+
h.service as unknown as {
178+
createNewLocalSession: (...args: unknown[]) => Promise<void>;
179+
},
180+
"createNewLocalSession",
181+
)
182+
.mockResolvedValue(undefined);
183+
vi.spyOn(
184+
h.service as unknown as {
185+
getAuthCredentialsStatus: () => Promise<unknown>;
186+
},
187+
"getAuthCredentialsStatus",
188+
).mockResolvedValue({ kind: "ready", auth: { client: {} } });
189+
190+
await h.service.clearSessionError(TASK_ID, "/repo");
191+
192+
expect(h.getPendingInitialPrompt).toHaveBeenCalledWith({
193+
taskId: TASK_ID,
194+
});
195+
expect(createNewLocalSession).toHaveBeenCalledWith(
196+
TASK_ID,
197+
"Task",
198+
"/repo",
199+
{ client: {} },
200+
PROMPT,
201+
undefined,
202+
undefined,
203+
undefined,
204+
undefined,
205+
);
206+
});
207+
});
208+
});

packages/core/src/sessions/sessionEventBatching.test.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,11 @@ function createHarness() {
109109
subscribe: () => ({ unsubscribe: vi.fn() }),
110110
},
111111
},
112+
workspace: {
113+
clearPendingInitialPrompt: {
114+
mutate: vi.fn().mockResolvedValue(undefined),
115+
},
116+
},
112117
},
113118
} as unknown as SessionServiceDeps;
114119

0 commit comments

Comments
 (0)