Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 57 additions & 1 deletion packages/core/src/cloud-task/cloud-task-engine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ const SSE_RECONNECT_BASE_DELAY_MS = 500;
const SSE_RECONNECT_FLAT_ATTEMPTS = 3;
const SSE_RECONNECT_MAX_DELAY_MS = 30_000;
const SSE_HEALTHY_CONNECTION_MS = 60_000;
// The backend emits a keepalive at least every ~25-30s (see SSE_KEEPALIVE_INTERVAL_MS in
// packages/agent). A half-open socket (laptop sleep, unplugged NIC, NAT rebind) neither errors
// nor EOFs, so `reader.read()` awaits forever with nothing to trigger reconnect. This timeout
// treats "no bytes at all for a few keepalive intervals" as a disconnect so it flows into the
// existing reconnect/backoff machinery instead of hanging the watcher indefinitely.
const SSE_IDLE_TIMEOUT_MS = 90_000;
const EVENT_BATCH_FLUSH_MS = 16;
const EVENT_BATCH_MAX_SIZE = 50;
const SESSION_LOG_PAGE_LIMIT = 5_000;
Expand Down Expand Up @@ -1344,6 +1350,24 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
let streamWasEstablished = false;
let bytesReceived = 0;
let eventsReceived = 0;
let idleTimedOut = false;
let idleTimeoutHandle: ReturnType<typeof setTimeout> | null = null;

const clearIdleTimeout = () => {
if (idleTimeoutHandle) {
clearTimeout(idleTimeoutHandle);
idleTimeoutHandle = null;
}
};
// Re-armed on every read that returns a value (data or keepalive bytes), so it only fires
// when the transport has gone completely silent, not merely between infrequent events.
const armIdleTimeout = () => {
clearIdleTimeout();
idleTimeoutHandle = setTimeout(() => {
idleTimedOut = true;
controller.abort();
}, SSE_IDLE_TIMEOUT_MS);
};

try {
// The proxy authenticates with the run-scoped Bearer token; the Django leg uses the session.
Expand Down Expand Up @@ -1401,13 +1425,16 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
});

const reader = response.body.getReader();
armIdleTimeout();

while (true) {
const { done, value } = await reader.read();
if (done) {
break;
}

armIdleTimeout();

if (!value) {
continue;
}
Expand Down Expand Up @@ -1463,10 +1490,38 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
} catch (error) {
this.flushLogBatch(key);

if (controller.signal.aborted) {
// An idle-timeout abort must fall through to the reconnect machinery below rather than
// return here like a deliberate cancel (disconnectSse/stopWatching), since nothing else
// will ever notice this connection went silent.
if (controller.signal.aborted && !idleTimedOut) {
return;
}

if (idleTimedOut) {
const idleWatcher = this.watchers.get(key);
this.log.warn("Cloud task stream idle timeout, no bytes received", {
key,
leg,
streamUrl: url.toString(),
idleTimeoutMs: SSE_IDLE_TIMEOUT_MS,
bytesReceived,
eventsReceived,
connectionDurationMs: streamWasEstablished
? Date.now() - connectedAt
: 0,
});
if (idleWatcher) {
this.analytics.track(ANALYTICS_EVENTS.CLOUD_STREAM_IDLE_TIMEOUT, {
task_id: idleWatcher.taskId,
run_id: idleWatcher.runId,
team_id: idleWatcher.teamId,
idle_timeout_ms: SSE_IDLE_TIMEOUT_MS,
bytes_received: bytesReceived,
events_received: eventsReceived,
});
}
}

// Proxy-leg 401: the read token expired or its signing key rotated. Re-resolve to mint a
// fresh token (or route back to Django) instead of failing. Django-leg 401 stays fatal below.
const unauthorizedWatcher = this.watchers.get(key);
Expand Down Expand Up @@ -1548,6 +1603,7 @@ export class CloudTaskEngine extends TypedEventEmitter<CloudTaskEvents> {
countReconnectAttempt: !isBackendError && !wasHealthyStream,
});
} finally {
clearIdleTimeout();
const currentWatcher = this.watchers.get(key);
if (currentWatcher?.sseAbortController === controller) {
currentWatcher.sseAbortController = null;
Expand Down
112 changes: 111 additions & 1 deletion packages/core/src/cloud-task/cloud-task.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ async function waitFor(

describe("CloudTaskEngine", () => {
let service: CloudTaskEngine;
let analyticsMock: { track: ReturnType<typeof vi.fn> };

beforeEach(() => {
const scopedLog = {
Expand All @@ -103,7 +104,7 @@ describe("CloudTaskEngine", () => {
error: vi.fn(),
};
const loggerMock = { ...scopedLog, scope: vi.fn(() => scopedLog) };
const analyticsMock = { track: vi.fn() };
analyticsMock = { track: vi.fn() };
service = createCloudTaskEngine({
auth: mockAuthService as never,
analytics: analyticsMock as never,
Expand Down Expand Up @@ -2404,6 +2405,115 @@ describe("CloudTaskEngine", () => {
).toBe(false);
});

it("aborts and reconnects a stream that goes silent with no bytes or keepalives", async () => {
vi.useFakeTimers();

const updates: unknown[] = [];
service.on(CloudTaskEvent.Update, (payload) => updates.push(payload));

const makeInProgressRun = () =>
createJsonResponse({
id: "run-1",
status: "in_progress",
stage: null,
output: null,
error_message: null,
branch: "main",
updated_at: "2026-01-01T00:00:00Z",
});

mockNetFetch
.mockResolvedValueOnce(makeInProgressRun())
.mockResolvedValueOnce(
createJsonResponse([], 200, { "X-Has-More": "false" }),
)
.mockImplementation(() => Promise.resolve(makeInProgressRun()));

// First connection hangs forever: no bytes, no error, no EOF, simulating a half-open
// socket (laptop sleep, NAT rebind). The second connection stays open and delivers a
// keepalive so recovery is observable once the idle watchdog aborts the first.
let streamCall = 0;
const encoder = new TextEncoder();
const abortedFirstConnection = { value: false };
mockStreamFetch.mockImplementation(
(_input: unknown, init?: RequestInit) => {
streamCall += 1;
if (streamCall === 1) {
const stream = new ReadableStream<Uint8Array>({
start(controller) {
// Never enqueue or close on our own; the read() promise awaits forever until the
// idle watchdog aborts it below, mirroring how a real fetch's reader rejects once
// its AbortSignal fires.
init?.signal?.addEventListener("abort", () => {
abortedFirstConnection.value = true;
controller.error(new DOMException("Aborted", "AbortError"));
});
},
});
return Promise.resolve(
new Response(stream, {
status: 200,
headers: { "Content-Type": "text/event-stream" },
}),
);
}
const stream = new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(
encoder.encode(
'event: keepalive\ndata: {"type":"keepalive"}\n\n',
),
);
},
});
return Promise.resolve(
new Response(stream, {
status: 200,
headers: { "Content-Type": "text/event-stream" },
}),
);
},
);

service.watch({
taskId: "task-1",
runId: "run-1",
apiHost: "https://app.example.com",
teamId: 2,
});

await waitFor(() => mockStreamFetch.mock.calls.length === 1);

// Nothing throws or EOFs; without the idle watchdog this would hang forever.
await vi.advanceTimersByTimeAsync(60_000);
expect(abortedFirstConnection.value).toBe(false);

await vi.advanceTimersByTimeAsync(40_000);
await waitFor(() => abortedFirstConnection.value, 20_000);
await waitFor(() => mockStreamFetch.mock.calls.length >= 2, 20_000);

expect(
analyticsMock.track.mock.calls.some(
([eventName]) => eventName === "Cloud stream idle timeout",
),
).toBe(true);

const watcher = (
service as unknown as {
watchers: Map<string, { failed: boolean }>;
}
).watchers.get("task-1:run-1");
expect(watcher?.failed).toBe(false);
expect(
updates.some(
(u) =>
typeof u === "object" &&
u !== null &&
(u as { kind?: string }).kind === "error",
),
).toBe(false);
});

it("stops a cloud run through the run cancel endpoint", async () => {
mockNetFetch.mockResolvedValueOnce(
createJsonResponse({ id: "run-1", status: "in_progress" }, 202),
Expand Down
11 changes: 11 additions & 0 deletions packages/shared/src/analytics-events.ts
Original file line number Diff line number Diff line change
Expand Up @@ -327,6 +327,15 @@ export interface CloudStreamDisconnectedProperties {
was_bootstrapping: boolean;
}

export interface CloudStreamIdleTimeoutProperties {
task_id: string;
run_id: string;
team_id: number;
idle_timeout_ms: number;
bytes_received: number;
events_received: number;
}

// Permission events
export interface PermissionRespondedProperties {
task_id: string;
Expand Down Expand Up @@ -1378,6 +1387,7 @@ export const ANALYTICS_EVENTS = {
TASK_CREATION_FAILED: "Task creation failed",
AGENT_SESSION_ERROR: "Agent session error",
CLOUD_STREAM_DISCONNECTED: "Cloud stream disconnected",
CLOUD_STREAM_IDLE_TIMEOUT: "Cloud stream idle timeout",

// Inbox events
INBOX_VIEWED: "Inbox viewed",
Expand Down Expand Up @@ -1556,6 +1566,7 @@ export type EventPropertyMap = {
[ANALYTICS_EVENTS.TASK_CREATION_FAILED]: TaskCreationFailedProperties;
[ANALYTICS_EVENTS.AGENT_SESSION_ERROR]: AgentSessionErrorProperties;
[ANALYTICS_EVENTS.CLOUD_STREAM_DISCONNECTED]: CloudStreamDisconnectedProperties;
[ANALYTICS_EVENTS.CLOUD_STREAM_IDLE_TIMEOUT]: CloudStreamIdleTimeoutProperties;

// Inbox events
[ANALYTICS_EVENTS.INBOX_VIEWED]: InboxViewedProperties;
Expand Down
Loading