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

Commit de581bc

Browse files
authored
fix(cloud-agent): improve cloud run reconnection (#1921)
1 parent 892cf92 commit de581bc

9 files changed

Lines changed: 926 additions & 51 deletions

File tree

‎apps/code/src/main/services/cloud-task/service.test.ts‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -328,6 +328,137 @@ describe("CloudTaskService", () => {
328328
);
329329
});
330330

331+
it("replays a current snapshot when a subscriber attaches to an existing watcher", async () => {
332+
const updates: unknown[] = [];
333+
service.on(CloudTaskEvent.Update, (payload) => updates.push(payload));
334+
335+
const historicalEntry = {
336+
type: "notification",
337+
timestamp: "2026-01-01T00:00:00Z",
338+
notification: {
339+
jsonrpc: "2.0",
340+
method: "_posthog/console",
341+
params: {
342+
sessionId: "run-1",
343+
level: "info",
344+
message: "older history",
345+
},
346+
},
347+
};
348+
const liveEntry = {
349+
type: "notification",
350+
timestamp: "2026-01-01T00:00:01Z",
351+
notification: {
352+
jsonrpc: "2.0",
353+
method: "_posthog/console",
354+
params: {
355+
sessionId: "run-1",
356+
level: "info",
357+
message: "live tail",
358+
},
359+
},
360+
};
361+
362+
const runResponse = {
363+
id: "run-1",
364+
status: "in_progress",
365+
stage: "build",
366+
output: null,
367+
error_message: null,
368+
branch: "main",
369+
updated_at: "2026-01-01T00:00:00Z",
370+
};
371+
372+
mockNetFetch
373+
.mockResolvedValueOnce(createJsonResponse(runResponse))
374+
.mockResolvedValueOnce(
375+
createJsonResponse([historicalEntry], 200, { "X-Has-More": "false" }),
376+
)
377+
.mockResolvedValueOnce(createJsonResponse(runResponse))
378+
.mockResolvedValueOnce(
379+
createJsonResponse([historicalEntry], 200, { "X-Has-More": "false" }),
380+
);
381+
382+
mockStreamFetch.mockResolvedValueOnce(
383+
createOpenSseResponse(`id: 1\ndata: ${JSON.stringify(liveEntry)}\n\n`),
384+
);
385+
386+
service.watch({
387+
taskId: "task-1",
388+
runId: "run-1",
389+
apiHost: "https://app.example.com",
390+
teamId: 2,
391+
});
392+
393+
await waitFor(() => updates.length >= 2);
394+
395+
service.watch({
396+
taskId: "task-1",
397+
runId: "run-1",
398+
apiHost: "https://app.example.com",
399+
teamId: 2,
400+
});
401+
402+
await waitFor(() =>
403+
updates.some(
404+
(update) =>
405+
typeof update === "object" &&
406+
update !== null &&
407+
(update as { kind?: string; totalEntryCount?: number }).kind ===
408+
"snapshot" &&
409+
(update as { totalEntryCount?: number }).totalEntryCount === 2,
410+
),
411+
);
412+
413+
const replayedSnapshot = updates.find(
414+
(update) =>
415+
typeof update === "object" &&
416+
update !== null &&
417+
(update as { kind?: string; totalEntryCount?: number }).kind ===
418+
"snapshot" &&
419+
(update as { totalEntryCount?: number }).totalEntryCount === 2,
420+
);
421+
422+
expect(replayedSnapshot).toEqual({
423+
taskId: "task-1",
424+
runId: "run-1",
425+
kind: "snapshot",
426+
newEntries: [historicalEntry, liveEntry],
427+
totalEntryCount: 2,
428+
status: "in_progress",
429+
stage: "build",
430+
output: null,
431+
errorMessage: null,
432+
branch: "main",
433+
});
434+
435+
const getWatcherEmittedEntryCount = (): number => {
436+
const watcher = (
437+
service as unknown as {
438+
watchers: Map<string, { emittedLogEntries: unknown[] }>;
439+
}
440+
).watchers.get("task-1:run-1");
441+
return watcher?.emittedLogEntries.length ?? 0;
442+
};
443+
444+
expect(getWatcherEmittedEntryCount()).toBe(1);
445+
446+
mockNetFetch.mockResolvedValueOnce(
447+
createJsonResponse([historicalEntry, liveEntry], 200, {
448+
"X-Has-More": "false",
449+
}),
450+
);
451+
452+
service.watch({
453+
taskId: "task-1",
454+
runId: "run-1",
455+
apiHost: "https://app.example.com",
456+
teamId: 2,
457+
});
458+
459+
await waitFor(() => getWatcherEmittedEntryCount() === 0);
460+
});
461+
331462
it("ignores keepalive SSE events while keeping the stream open", async () => {
332463
const updates: unknown[] = [];
333464
service.on(CloudTaskEvent.Update, (payload) => updates.push(payload));

‎apps/code/src/main/services/cloud-task/service.ts‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ interface WatcherState {
9292
isBootstrapping: boolean;
9393
hasEmittedSnapshot: boolean;
9494
bufferedLogBatches: StoredLogEntry[][];
95+
emittedLogEntries: StoredLogEntry[];
9596
failed: boolean;
9697
needsPostBootstrapReconnect: boolean;
9798
needsStopAfterBootstrap: boolean;
@@ -224,6 +225,7 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
224225
key,
225226
subscribers: existing.subscriberCount,
226227
});
228+
void this.emitCurrentSnapshot(key);
227229
return;
228230
}
229231

@@ -399,6 +401,7 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
399401
isBootstrapping: false,
400402
hasEmittedSnapshot: false,
401403
bufferedLogBatches: [],
404+
emittedLogEntries: [],
402405
failed: false,
403406
needsPostBootstrapReconnect: false,
404407
needsStopAfterBootstrap: false,
@@ -767,6 +770,7 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
767770
}
768771

769772
watcher.totalEntryCount += entries.length;
773+
this.rememberEmittedLogEntries(watcher, entries);
770774

771775
this.emit(CloudTaskEvent.Update, {
772776
taskId: watcher.taskId,
@@ -812,6 +816,7 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
812816
}
813817

814818
watcher.totalEntryCount += dedupedEntries.length;
819+
this.rememberEmittedLogEntries(watcher, dedupedEntries);
815820
this.emit(CloudTaskEvent.Update, {
816821
taskId: watcher.taskId,
817822
runId: watcher.runId,
@@ -824,6 +829,92 @@ export class CloudTaskService extends TypedEventEmitter<CloudTaskEvents> {
824829
watcher.bufferedLogBatches = [];
825830
}
826831

832+
private rememberEmittedLogEntries(
833+
watcher: WatcherState,
834+
entries: StoredLogEntry[],
835+
): void {
836+
watcher.emittedLogEntries.push(...entries);
837+
}
838+
839+
private mergeHistoricalAndEmittedEntries(
840+
historicalEntries: StoredLogEntry[],
841+
emittedEntries: StoredLogEntry[],
842+
): {
843+
snapshotEntries: StoredLogEntry[];
844+
missingEmittedEntries: StoredLogEntry[];
845+
} {
846+
if (emittedEntries.length === 0) {
847+
return { snapshotEntries: historicalEntries, missingEmittedEntries: [] };
848+
}
849+
850+
const historicalCounts = new Map<string, number>();
851+
for (const entry of historicalEntries) {
852+
const serialized = JSON.stringify(entry);
853+
historicalCounts.set(
854+
serialized,
855+
(historicalCounts.get(serialized) ?? 0) + 1,
856+
);
857+
}
858+
859+
const missingEmittedEntries = emittedEntries.filter((entry) => {
860+
const serialized = JSON.stringify(entry);
861+
const remaining = historicalCounts.get(serialized) ?? 0;
862+
if (remaining <= 0) {
863+
return true;
864+
}
865+
866+
historicalCounts.set(serialized, remaining - 1);
867+
return false;
868+
});
869+
870+
return {
871+
snapshotEntries: [...historicalEntries, ...missingEmittedEntries],
872+
missingEmittedEntries,
873+
};
874+
}
875+
876+
private async emitCurrentSnapshot(key: string): Promise<void> {
877+
const watcher = this.watchers.get(key);
878+
if (!watcher || watcher.failed) return;
879+
880+
const historicalEntries = await this.fetchAllSessionLogs(watcher);
881+
const currentWatcher = this.watchers.get(key);
882+
if (!currentWatcher || currentWatcher !== watcher || watcher.failed) {
883+
return;
884+
}
885+
886+
if (!historicalEntries) {
887+
log.warn("Cloud task snapshot replay failed", {
888+
taskId: watcher.taskId,
889+
runId: watcher.runId,
890+
});
891+
return;
892+
}
893+
894+
const { snapshotEntries, missingEmittedEntries } =
895+
this.mergeHistoricalAndEmittedEntries(
896+
historicalEntries,
897+
watcher.emittedLogEntries,
898+
);
899+
watcher.emittedLogEntries = missingEmittedEntries;
900+
if (snapshotEntries.length > watcher.totalEntryCount) {
901+
watcher.totalEntryCount = snapshotEntries.length;
902+
}
903+
904+
this.emit(CloudTaskEvent.Update, {
905+
taskId: watcher.taskId,
906+
runId: watcher.runId,
907+
kind: "snapshot",
908+
newEntries: snapshotEntries,
909+
totalEntryCount: snapshotEntries.length,
910+
status: watcher.lastStatus ?? undefined,
911+
stage: watcher.lastStage,
912+
output: watcher.lastOutput,
913+
errorMessage: watcher.lastErrorMessage,
914+
branch: watcher.lastBranch,
915+
});
916+
}
917+
827918
private failWatcher(key: string, error: CloudTaskConnectionError): void {
828919
const watcher = this.watchers.get(key);
829920
if (!watcher) return;

‎apps/code/src/main/trpc/routers/cloud-task.ts‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -37,15 +37,19 @@ export const cloudTaskRouter = router({
3737
.input(onUpdateInput)
3838
.subscription(async function* (opts) {
3939
const service = getService();
40-
for await (const data of service.toIterable(CloudTaskEvent.Update, {
41-
signal: opts.signal,
42-
})) {
43-
if (
44-
data.taskId === opts.input.taskId &&
45-
data.runId === opts.input.runId
46-
) {
47-
yield data;
40+
try {
41+
for await (const data of service.toIterable(CloudTaskEvent.Update, {
42+
signal: opts.signal,
43+
})) {
44+
if (
45+
data.taskId === opts.input.taskId &&
46+
data.runId === opts.input.runId
47+
) {
48+
yield data;
49+
}
4850
}
51+
} finally {
52+
service.unwatch(opts.input.taskId, opts.input.runId);
4953
}
5054
}),
5155
});

‎apps/code/src/renderer/features/sessions/components/buildConversationItems.test.ts‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,18 @@ function promptResponseMsg(ts: number, id: number): AcpMessage {
6262
};
6363
}
6464

65+
function turnCompleteMsg(ts: number, stopReason = "end_turn"): AcpMessage {
66+
return {
67+
type: "acp_message",
68+
ts,
69+
message: {
70+
jsonrpc: "2.0",
71+
method: "_posthog/turn_complete",
72+
params: { sessionId: "session-1", stopReason },
73+
},
74+
};
75+
}
76+
6577
describe("buildConversationItems", () => {
6678
it("extracts cloud prompt attachments into user messages", () => {
6779
const uri = makeAttachmentUri("/tmp/hello world.txt");
@@ -109,6 +121,19 @@ describe("buildConversationItems", () => {
109121
]);
110122
});
111123

124+
it("marks cloud turns complete from structured turn completion notifications", () => {
125+
const result = buildConversationItems(
126+
[userPromptMsg(10, 42, "hello"), turnCompleteMsg(25)],
127+
null,
128+
);
129+
130+
expect(result.lastTurnInfo).toEqual({
131+
isComplete: true,
132+
durationMs: 15,
133+
stopReason: "end_turn",
134+
});
135+
});
136+
112137
it("keeps attachment-only prompts visible", () => {
113138
const uri = makeAttachmentUri("/tmp/test.txt");
114139

0 commit comments

Comments
 (0)