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

Commit 54a524b

Browse files
tatoalotatoalo
authored
fix(cloud-task): cover silent stream setup
Generated-By: PostHog Code Task-Id: 92b56d91-0d15-416f-8ab9-2b32e3cb3e34
1 parent 7137b92 commit 54a524b

2 files changed

Lines changed: 193 additions & 53 deletions

File tree

‎packages/core/src/cloud-task/cloud-task-engine.ts‎

Lines changed: 77 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -1262,12 +1262,75 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
12621262
const controller = new AbortController();
12631263
watcher.sseAbortController = controller;
12641264

1265+
let connectedAt = 0;
1266+
let streamWasEstablished = false;
1267+
let bytesReceived = 0;
1268+
let eventsReceived = 0;
1269+
let idleTimedOut = false;
1270+
let idleTimeoutHandle: ReturnType<typeof setTimeout> | null = null;
1271+
let idlePhase: "target_resolution" | "connection" | "stream" =
1272+
"target_resolution";
1273+
1274+
const clearIdleTimeout = () => {
1275+
if (idleTimeoutHandle) {
1276+
clearTimeout(idleTimeoutHandle);
1277+
idleTimeoutHandle = null;
1278+
}
1279+
};
1280+
const armIdleTimeout = () => {
1281+
clearIdleTimeout();
1282+
idleTimeoutHandle = setTimeout(() => {
1283+
idleTimedOut = true;
1284+
controller.abort();
1285+
}, SSE_IDLE_TIMEOUT_MS);
1286+
};
1287+
const recordIdleTimeout = (details: Record<string, unknown> = {}) => {
1288+
const idleWatcher = this.watchers.get(key);
1289+
this.log.warn("Cloud task stream idle timeout, no bytes received", {
1290+
key,
1291+
phase: idlePhase,
1292+
idleTimeoutMs: SSE_IDLE_TIMEOUT_MS,
1293+
bytesReceived,
1294+
eventsReceived,
1295+
connectionDurationMs: streamWasEstablished
1296+
? Date.now() - connectedAt
1297+
: 0,
1298+
...details,
1299+
});
1300+
if (idleWatcher) {
1301+
this.analytics.track(ANALYTICS_EVENTS.CLOUD_STREAM_IDLE_TIMEOUT, {
1302+
task_id: idleWatcher.taskId,
1303+
run_id: idleWatcher.runId,
1304+
team_id: idleWatcher.teamId,
1305+
idle_timeout_ms: SSE_IDLE_TIMEOUT_MS,
1306+
bytes_received: bytesReceived,
1307+
events_received: eventsReceived,
1308+
});
1309+
}
1310+
};
1311+
12651312
watcher.connStartedAt = 0;
12661313
watcher.connDataEventsReceived = 0;
12671314

12681315
// Resolve the read target once (proxy URL + token, or Django), reused across reconnects.
12691316
if (!watcher.streamTargetResolved) {
1270-
await this.resolveStreamTarget(watcher);
1317+
armIdleTimeout();
1318+
try {
1319+
await this.resolveStreamTarget(watcher, controller.signal);
1320+
} catch (error) {
1321+
if (!idleTimedOut) {
1322+
return;
1323+
}
1324+
recordIdleTimeout();
1325+
await this.handleStreamCompletion(key, {
1326+
reconnectOnDisconnect: true,
1327+
reconnectError: error,
1328+
countReconnectAttempt: true,
1329+
});
1330+
return;
1331+
} finally {
1332+
clearIdleTimeout();
1333+
}
12711334
const resolvedWatcher = this.watchers.get(key);
12721335
if (
12731336
!resolvedWatcher ||
@@ -1346,30 +1409,11 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
13461409

13471410
// Track how long the body stayed open so healthy long-lived connections cut by churn
13481411
// aren't penalized as failed reconnects (see SSE_HEALTHY_CONNECTION_MS).
1349-
let connectedAt = 0;
1350-
let streamWasEstablished = false;
1351-
let bytesReceived = 0;
1352-
let eventsReceived = 0;
1353-
let idleTimedOut = false;
1354-
let idleTimeoutHandle: ReturnType<typeof setTimeout> | null = null;
1355-
1356-
const clearIdleTimeout = () => {
1357-
if (idleTimeoutHandle) {
1358-
clearTimeout(idleTimeoutHandle);
1359-
idleTimeoutHandle = null;
1360-
}
1361-
};
13621412
// Re-armed on every read that returns a value (data or keepalive bytes), so it only fires
13631413
// when the transport has gone completely silent, not merely between infrequent events.
1364-
const armIdleTimeout = () => {
1365-
clearIdleTimeout();
1366-
idleTimeoutHandle = setTimeout(() => {
1367-
idleTimedOut = true;
1368-
controller.abort();
1369-
}, SSE_IDLE_TIMEOUT_MS);
1370-
};
1371-
13721414
try {
1415+
idlePhase = "connection";
1416+
armIdleTimeout();
13731417
// The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session.
13741418
const response = usingProxy
13751419
? await this.streamFetch(url.toString(), {
@@ -1425,6 +1469,7 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
14251469
});
14261470

14271471
const reader = response.body.getReader();
1472+
idlePhase = "stream";
14281473
armIdleTimeout();
14291474

14301475
while (true) {
@@ -1498,28 +1543,10 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
14981543
}
14991544

15001545
if (idleTimedOut) {
1501-
const idleWatcher = this.watchers.get(key);
1502-
this.log.warn("Cloud task stream idle timeout, no bytes received", {
1503-
key,
1546+
recordIdleTimeout({
15041547
leg,
15051548
streamUrl: url.toString(),
1506-
idleTimeoutMs: SSE_IDLE_TIMEOUT_MS,
1507-
bytesReceived,
1508-
eventsReceived,
1509-
connectionDurationMs: streamWasEstablished
1510-
? Date.now() - connectedAt
1511-
: 0,
15121549
});
1513-
if (idleWatcher) {
1514-
this.analytics.track(ANALYTICS_EVENTS.CLOUD_STREAM_IDLE_TIMEOUT, {
1515-
task_id: idleWatcher.taskId,
1516-
run_id: idleWatcher.runId,
1517-
team_id: idleWatcher.teamId,
1518-
idle_timeout_ms: SSE_IDLE_TIMEOUT_MS,
1519-
bytes_received: bytesReceived,
1520-
events_received: eventsReceived,
1521-
});
1522-
}
15231550
}
15241551

15251552
// Proxy-leg 401: the read token expired or its signing key rotated. Re-resolve to mint a
@@ -1561,6 +1588,7 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
15611588
const isBackendError = error instanceof BackendStreamError;
15621589
const wasHealthyStream =
15631590
!isBackendError &&
1591+
!idleTimedOut &&
15641592
streamWasEstablished &&
15651593
Date.now() - connectedAt >= SSE_HEALTHY_CONNECTION_MS;
15661594

@@ -2254,11 +2282,15 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
22542282
}
22552283
}
22562284

2257-
private async resolveStreamTarget(watcher: WatcherState): Promise<void> {
2285+
private async resolveStreamTarget(
2286+
watcher: WatcherState,
2287+
signal: AbortSignal,
2288+
): Promise<void> {
22582289
const url = `${watcher.apiHost}/api/projects/${watcher.teamId}/tasks/${watcher.taskId}/runs/${watcher.runId}/stream_token/`;
22592290
try {
22602291
const response = await this.auth.authenticatedFetch(url, {
22612292
method: "GET",
2293+
signal,
22622294
});
22632295
if (!response.ok) {
22642296
watcher.streamBaseUrl = null;
@@ -2301,6 +2333,9 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
23012333
durableStream: watcher.durableStreamEnabled,
23022334
});
23032335
} catch (error) {
2336+
if (signal.aborted) {
2337+
throw error;
2338+
}
23042339
// Transient failure: leave unresolved so the next reconnect retries and falls back to Django.
23052340
watcher.streamBaseUrl = null;
23062341
watcher.streamReadToken = null;

‎packages/core/src/cloud-task/cloud-task.test.ts‎

Lines changed: 116 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2430,10 +2430,9 @@ describe("CloudTaskEngine", () => {
24302430
.mockImplementation(() => Promise.resolve(makeInProgressRun()));
24312431

24322432
// First connection hangs forever: no bytes, no error, no EOF, simulating a half-open
2433-
// socket (laptop sleep, NAT rebind). The second connection stays open and delivers a
2434-
// keepalive so recovery is observable once the idle watchdog aborts the first.
2433+
// socket (laptop sleep, NAT rebind). The second connection stays open so recovery is
2434+
// observable once the idle watchdog aborts the first.
24352435
let streamCall = 0;
2436-
const encoder = new TextEncoder();
24372436
const abortedFirstConnection = { value: false };
24382437
mockStreamFetch.mockImplementation(
24392438
(_input: unknown, init?: RequestInit) => {
@@ -2458,13 +2457,7 @@ describe("CloudTaskEngine", () => {
24582457
);
24592458
}
24602459
const stream = new ReadableStream<Uint8Array>({
2461-
start(controller) {
2462-
controller.enqueue(
2463-
encoder.encode(
2464-
'event: keepalive\ndata: {"type":"keepalive"}\n\n',
2465-
),
2466-
);
2467-
},
2460+
start() {},
24682461
});
24692462
return Promise.resolve(
24702463
new Response(stream, {
@@ -2500,10 +2493,13 @@ describe("CloudTaskEngine", () => {
25002493

25012494
const watcher = (
25022495
service as unknown as {
2503-
watchers: Map<string, { failed: boolean }>;
2496+
watchers: Map<string, { failed: boolean; reconnectAttempts: number }>;
25042497
}
25052498
).watchers.get("task-1:run-1");
25062499
expect(watcher?.failed).toBe(false);
2500+
// Silence is a broken transport, not a healthy long-lived connection. It must consume the
2501+
// reconnect budget so a persistently silent endpoint eventually reaches the circuit breaker.
2502+
expect(watcher?.reconnectAttempts).toBe(1);
25072503
expect(
25082504
updates.some(
25092505
(u) =>
@@ -2514,6 +2510,115 @@ describe("CloudTaskEngine", () => {
25142510
).toBe(false);
25152511
});
25162512

2513+
it("times out while resolving the stream target and reconnects", async () => {
2514+
vi.useFakeTimers();
2515+
2516+
mockNetFetch
2517+
.mockResolvedValueOnce(
2518+
createJsonResponse({ id: "run-1", status: "in_progress" }),
2519+
)
2520+
.mockResolvedValueOnce(
2521+
createJsonResponse([], 200, { "X-Has-More": "false" }),
2522+
)
2523+
.mockResolvedValue(
2524+
createJsonResponse({ id: "run-1", status: "in_progress" }),
2525+
);
2526+
2527+
let tokenCall = 0;
2528+
const abortedResolution = { value: false };
2529+
mockStreamTokenFetch.mockImplementation(
2530+
(_input: unknown, init?: RequestInit) => {
2531+
tokenCall += 1;
2532+
if (tokenCall === 1) {
2533+
return new Promise<Response>((_resolve, reject) => {
2534+
init?.signal?.addEventListener("abort", () => {
2535+
abortedResolution.value = true;
2536+
reject(new DOMException("Aborted", "AbortError"));
2537+
});
2538+
});
2539+
}
2540+
return Promise.resolve(
2541+
createJsonResponse({ token: "test-token", stream_base_url: null }),
2542+
);
2543+
},
2544+
);
2545+
mockStreamFetch.mockImplementation(
2546+
() =>
2547+
new Promise<Response>(() => {
2548+
// The connection-phase watchdog owns this second pending request.
2549+
}),
2550+
);
2551+
2552+
service.watch({
2553+
taskId: "task-1",
2554+
runId: "run-1",
2555+
apiHost: "https://app.example.com",
2556+
teamId: 2,
2557+
});
2558+
2559+
await waitFor(() => mockStreamTokenFetch.mock.calls.length === 1);
2560+
await vi.advanceTimersByTimeAsync(100_000);
2561+
await waitFor(() => abortedResolution.value, 20_000);
2562+
await waitFor(() => mockStreamTokenFetch.mock.calls.length >= 2, 20_000);
2563+
2564+
expect(mockStreamTokenFetch.mock.calls[0]?.[1]?.signal).toBeDefined();
2565+
expect(
2566+
analyticsMock.track.mock.calls.some(
2567+
([eventName]) => eventName === "Cloud stream idle timeout",
2568+
),
2569+
).toBe(true);
2570+
});
2571+
2572+
it("times out while waiting for stream response headers and reconnects", async () => {
2573+
vi.useFakeTimers();
2574+
2575+
mockNetFetch
2576+
.mockResolvedValueOnce(
2577+
createJsonResponse({ id: "run-1", status: "in_progress" }),
2578+
)
2579+
.mockResolvedValueOnce(
2580+
createJsonResponse([], 200, { "X-Has-More": "false" }),
2581+
)
2582+
.mockResolvedValue(
2583+
createJsonResponse({ id: "run-1", status: "in_progress" }),
2584+
);
2585+
2586+
let streamCall = 0;
2587+
const abortedConnection = { value: false };
2588+
mockStreamFetch.mockImplementation(
2589+
(_input: unknown, init?: RequestInit) => {
2590+
streamCall += 1;
2591+
if (streamCall === 1) {
2592+
return new Promise<Response>((_resolve, reject) => {
2593+
init?.signal?.addEventListener("abort", () => {
2594+
abortedConnection.value = true;
2595+
reject(new DOMException("Aborted", "AbortError"));
2596+
});
2597+
});
2598+
}
2599+
return Promise.resolve(createOpenSseResponse(""));
2600+
},
2601+
);
2602+
2603+
service.watch({
2604+
taskId: "task-1",
2605+
runId: "run-1",
2606+
apiHost: "https://app.example.com",
2607+
teamId: 2,
2608+
});
2609+
2610+
await waitFor(() => mockStreamFetch.mock.calls.length === 1);
2611+
await vi.advanceTimersByTimeAsync(100_000);
2612+
await waitFor(() => abortedConnection.value, 20_000);
2613+
await waitFor(() => mockStreamFetch.mock.calls.length >= 2, 20_000);
2614+
2615+
expect(
2616+
analyticsMock.track.mock.calls.some(
2617+
([eventName]) => eventName === "Cloud stream idle timeout",
2618+
),
2619+
).toBe(true);
2620+
});
2621+
25172622
it("stops a cloud run through the run cancel endpoint", async () => {
25182623
mockNetFetch.mockResolvedValueOnce(
25192624
createJsonResponse({ id: "run-1", status: "in_progress" }, 202),

0 commit comments

Comments
 (0)