Skip to content

Commit 8c75caf

Browse files
ericallamericallam
committed
test(core,webapp): suspend-hook, OOM-retry, and idle-timeout e2e legs
Rounds out the run-lifecycle coverage: onChatSuspend/onChatResume fire around a real suspend/resume, an OutOfMemoryError on the first attempt fails the run the way a machine swap needs and the attempt-2 retry recovers the in-flight message, and a run with no next message times out on the waitpoint and exits. Adds waitpoint-timeout support to the in-process backend (honors the wire timeout, resolves ok:false) so the idle-timeout path is exercised.
1 parent 267ef71 commit 8c75caf

4 files changed

Lines changed: 307 additions & 4 deletions

File tree

‎apps/webapp/test/helpers/agentHarness.ts‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,12 @@ export type RunRealChatAgentOptions = {
3232
* suspend/resume path in a test.
3333
*/
3434
idleTimeoutInSeconds?: number;
35+
/**
36+
* `ctx.attempt.number`. A value greater than 1 makes the boot treat the run
37+
* as a retry (`couldHavePriorState`), restoring from the snapshot + `.in`
38+
* replay. Used to model an OOM retry re-dispatch.
39+
*/
40+
attemptNumber?: number;
3541
};
3642

3743
export type RunningAgent = {
@@ -90,7 +96,14 @@ export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
9096
: { chatId: opts.addressingKey, trigger: "preload", metadata: {}, ...idle };
9197
await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal });
9298
},
93-
{ ctx: { run: { id: runId } }, sessionStreamManager: manager, runtimeManager }
99+
{
100+
ctx: {
101+
run: { id: runId },
102+
...(opts.attemptNumber !== undefined ? { attempt: { number: opts.attemptNumber } } : {}),
103+
},
104+
sessionStreamManager: manager,
105+
runtimeManager,
106+
}
94107
) as Promise<void>
95108
).finally(restore);
96109

‎apps/webapp/test/helpers/testChatAgent.ts‎

Lines changed: 74 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { chat } from "@trigger.dev/sdk/ai";
2-
import { locals } from "@trigger.dev/core/v3";
2+
import { locals, OutOfMemoryError } from "@trigger.dev/core/v3";
33
import { stepCountIs, streamText, tool, type LanguageModel, type UIMessage } from "ai";
44
import { z } from "zod";
55

@@ -62,6 +62,79 @@ export const testChatAgent = chat
6262
},
6363
});
6464

65+
/**
66+
* Records suspend/resume lifecycle-hook fires (keyed by chatId so tests can
67+
* filter to their own session). Appended to by {@link testSuspendHooksChatAgent}.
68+
*/
69+
export const suspendResumeEvents: Array<{
70+
chatId: string;
71+
kind: "suspend" | "resume";
72+
phase: string;
73+
}> = [];
74+
75+
/**
76+
* Short-idle agent that records `onChatSuspend` / `onChatResume` fires, so a
77+
* test can assert the lifecycle hooks run around a real suspend/resume.
78+
*/
79+
export const testSuspendHooksChatAgent = chat.agent({
80+
id: "e2e-test-chat-suspend-hooks",
81+
idleTimeoutInSeconds: 1,
82+
preloadIdleTimeoutInSeconds: 1,
83+
onChatSuspend: async ({ chatId, phase }) => {
84+
suspendResumeEvents.push({ chatId, kind: "suspend", phase });
85+
},
86+
onChatResume: async ({ chatId, phase }) => {
87+
suspendResumeEvents.push({ chatId, kind: "resume", phase });
88+
},
89+
run: async ({ messages, signal }) => {
90+
const model = locals.get(testChatModelLocal);
91+
if (!model) {
92+
throw new Error("test model not injected via locals");
93+
}
94+
return streamText({ model, messages, abortSignal: signal });
95+
},
96+
});
97+
98+
/**
99+
* Throws `OutOfMemoryError` on the first attempt so the run fails the way the
100+
* runtime detects for a machine swap, then succeeds on attempt 2 (the retry
101+
* boots through the restore path because `ctx.attempt.number > 1`).
102+
*/
103+
export const testOomChatAgent = chat.agent({
104+
id: "e2e-test-chat-oom",
105+
idleTimeoutInSeconds: 1,
106+
preloadIdleTimeoutInSeconds: 1,
107+
run: async ({ messages, signal, ctx }) => {
108+
if (ctx.attempt.number === 1) {
109+
throw new OutOfMemoryError();
110+
}
111+
const model = locals.get(testChatModelLocal);
112+
if (!model) {
113+
throw new Error("test model not injected via locals");
114+
}
115+
return streamText({ model, messages, abortSignal: signal });
116+
},
117+
});
118+
119+
/**
120+
* Short idle window and a short `turnTimeout`, so a run with no incoming
121+
* message suspends and then times out on the waitpoint, ending the run.
122+
*/
123+
export const testTimeoutChatAgent = chat.agent({
124+
id: "e2e-test-chat-timeout",
125+
idleTimeoutInSeconds: 1,
126+
preloadIdleTimeoutInSeconds: 1,
127+
turnTimeout: "3s",
128+
preloadTimeout: "3s",
129+
run: async ({ messages, signal }) => {
130+
const model = locals.get(testChatModelLocal);
131+
if (!model) {
132+
throw new Error("test model not injected via locals");
133+
}
134+
return streamText({ model, messages, abortSignal: signal });
135+
},
136+
});
137+
65138
/**
66139
* A minimal agent with no lifecycle hooks. With `hydrateMessages` absent, the
67140
* default snapshot + `.out`/`.in` replay boot path is what restores prior

‎apps/webapp/test/session-agent.e2e.test.ts‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,18 @@ import {
2929
} from "./helpers/sessionStream";
3030
import { runChatAgentSession, runRealChatAgent } from "./helpers/agentHarness";
3131
import {
32+
suspendResumeEvents,
3233
testApprovalChatAgent,
3334
testChatAgent,
3435
testChatModelLocal,
3536
testEndRunChatAgent,
3637
testHitlChatAgent,
3738
testHitlIdleChatAgent,
3839
testIdleChatAgent,
40+
testOomChatAgent,
3941
testPlainChatAgent,
42+
testSuspendHooksChatAgent,
43+
testTimeoutChatAgent,
4044
testToolChatAgent,
4145
testUpgradeChatAgent,
4246
testUpgradeOnceChatAgent,
@@ -1350,4 +1354,166 @@ describe("session agent e2e (real chat.agent loop)", () => {
13501354
await session.close();
13511355
}
13521356
});
1357+
1358+
it("EA20: onChatSuspend and onChatResume fire around a suspend/resume", async () => {
1359+
const { addressingKey, token, apiKey, baseUrl } = await setupSession(
1360+
testSuspendHooksChatAgent.id
1361+
);
1362+
const agent = runRealChatAgent({
1363+
agentId: testSuspendHooksChatAgent.id,
1364+
baseUrl,
1365+
addressingKey,
1366+
secretKey: apiKey,
1367+
model: sequenceModel(["turn-a", "turn-b"]),
1368+
modelLocal: testChatModelLocal,
1369+
});
1370+
1371+
try {
1372+
await appendInput({
1373+
baseUrl,
1374+
addressingKey,
1375+
token,
1376+
partId: "u1",
1377+
body: submitBody(addressingKey, userMessage("first", "u1")),
1378+
});
1379+
await collectSessionOut({
1380+
baseUrl,
1381+
addressingKey,
1382+
token,
1383+
until: (p) => p.some(isTurnComplete),
1384+
maxMs: 30_000,
1385+
});
1386+
1387+
await new Promise((r) => setTimeout(r, 2500));
1388+
1389+
await appendInput({
1390+
baseUrl,
1391+
addressingKey,
1392+
token,
1393+
partId: "u2",
1394+
body: submitBody(addressingKey, userMessage("second", "u2")),
1395+
});
1396+
await collectSessionOut({
1397+
baseUrl,
1398+
addressingKey,
1399+
token,
1400+
until: (p) => joinChunks(p).includes("turn-b"),
1401+
maxMs: 30_000,
1402+
});
1403+
1404+
const mine = suspendResumeEvents.filter((e) => e.chatId === addressingKey);
1405+
expect(
1406+
mine.some((e) => e.kind === "suspend"),
1407+
"onChatSuspend fired when the run suspended"
1408+
).toBe(true);
1409+
expect(
1410+
mine.some((e) => e.kind === "resume"),
1411+
"onChatResume fired when the next message resumed it"
1412+
).toBe(true);
1413+
} finally {
1414+
await agent.close();
1415+
}
1416+
});
1417+
1418+
it("EA21: an OOM fails the run; the attempt-2 retry recovers the message", async () => {
1419+
const { addressingKey, token, apiKey, baseUrl } = await setupSession(testOomChatAgent.id);
1420+
const runId = `run_oom_${addressingKey}`;
1421+
1422+
const attempt1 = runRealChatAgent({
1423+
agentId: testOomChatAgent.id,
1424+
baseUrl,
1425+
addressingKey,
1426+
secretKey: apiKey,
1427+
model: echoModel(),
1428+
modelLocal: testChatModelLocal,
1429+
runId,
1430+
attemptNumber: 1,
1431+
});
1432+
let attempt1Error: unknown;
1433+
attempt1.done.catch((e) => {
1434+
attempt1Error = e;
1435+
});
1436+
1437+
await appendInput({
1438+
baseUrl,
1439+
addressingKey,
1440+
token,
1441+
partId: "u1",
1442+
body: submitBody(addressingKey, userMessage("OOM-THEN-OK", "u1")),
1443+
});
1444+
await waitFor(() => attempt1Error !== undefined, 15_000);
1445+
expect(
1446+
String(attempt1Error),
1447+
"attempt 1 fails with an OOM so the runtime can swap machines"
1448+
).toMatch(/OutOfMemory/i);
1449+
1450+
const attempt2 = runRealChatAgent({
1451+
agentId: testOomChatAgent.id,
1452+
baseUrl,
1453+
addressingKey,
1454+
secretKey: apiKey,
1455+
model: echoModel(),
1456+
modelLocal: testChatModelLocal,
1457+
runId,
1458+
attemptNumber: 2,
1459+
continuation: true,
1460+
});
1461+
try {
1462+
const { parts } = await collectSessionOut({
1463+
baseUrl,
1464+
addressingKey,
1465+
token,
1466+
until: (p) => joinChunks(p).includes("OOM-THEN-OK"),
1467+
maxMs: 30_000,
1468+
});
1469+
expect(
1470+
joinChunks(parts),
1471+
"the attempt-2 retry restored the unprocessed message and ran it"
1472+
).toContain("OOM-THEN-OK");
1473+
} finally {
1474+
await attempt2.close();
1475+
}
1476+
});
1477+
1478+
it("EA22: the idle wait times out and ends the run when no message arrives", async () => {
1479+
const { addressingKey, token, apiKey, baseUrl } = await setupSession(testTimeoutChatAgent.id);
1480+
const agent = runRealChatAgent({
1481+
agentId: testTimeoutChatAgent.id,
1482+
baseUrl,
1483+
addressingKey,
1484+
secretKey: apiKey,
1485+
model: textModel("turn one"),
1486+
modelLocal: testChatModelLocal,
1487+
});
1488+
1489+
let ended = false;
1490+
agent.done.then(() => {
1491+
ended = true;
1492+
});
1493+
1494+
try {
1495+
await appendInput({
1496+
baseUrl,
1497+
addressingKey,
1498+
token,
1499+
partId: "u1",
1500+
body: submitBody(addressingKey, userMessage("hello", "u1")),
1501+
});
1502+
await collectSessionOut({
1503+
baseUrl,
1504+
addressingKey,
1505+
token,
1506+
until: (p) => p.some(isTurnComplete),
1507+
maxMs: 30_000,
1508+
});
1509+
1510+
await waitFor(() => ended, 20_000);
1511+
expect(
1512+
ended,
1513+
"with no next message the between-turns wait times out and the run exits itself"
1514+
).toBe(true);
1515+
} finally {
1516+
await agent.close();
1517+
}
1518+
});
13531519
});

‎packages/core/src/v3/test/session-waitpoint-backend.ts‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,33 @@ type PendingWait = {
1212
session: string;
1313
io: "in" | "out";
1414
lastSeqNum?: number;
15+
timeout?: string;
1516
abort: AbortController;
1617
};
1718

19+
const TIMED_OUT = Symbol("session-waitpoint-timeout");
20+
21+
function parseTimeoutMs(timeout: string | undefined): number | undefined {
22+
if (!timeout) {
23+
return undefined;
24+
}
25+
const match = /^(\d+)(ms|s|m|h)$/.exec(timeout.trim());
26+
if (!match) {
27+
return undefined;
28+
}
29+
const value = Number(match[1]);
30+
switch (match[2]) {
31+
case "ms":
32+
return value;
33+
case "s":
34+
return value * 1_000;
35+
case "m":
36+
return value * 60_000;
37+
default:
38+
return value * 3_600_000;
39+
}
40+
}
41+
1842
/**
1943
* In-process stand-in for the run-engine's session-stream waitpoint machinery.
2044
*
@@ -45,6 +69,7 @@ export class SessionWaitpointBackend {
4569
session: body.session,
4670
io: body.io,
4771
lastSeqNum: body.lastSeqNum,
72+
timeout: body.timeout,
4873
abort: new AbortController(),
4974
});
5075
return { waitpointId, isCached: false };
@@ -57,16 +82,42 @@ export class SessionWaitpointBackend {
5782
}
5883
this.pending.delete(waitpointFriendlyId);
5984

85+
const timeoutMs = parseTimeoutMs(pending.timeout);
86+
let timer: ReturnType<typeof setTimeout> | undefined;
87+
6088
try {
61-
const record = await this.readNextRecord(pending);
62-
const output = typeof record === "string" ? record : JSON.stringify(record);
89+
const recordPromise = this.readNextRecord(pending);
90+
const result =
91+
timeoutMs === undefined
92+
? await recordPromise
93+
: await Promise.race([
94+
recordPromise,
95+
new Promise<typeof TIMED_OUT>((resolve) => {
96+
timer = setTimeout(() => resolve(TIMED_OUT), timeoutMs);
97+
}),
98+
]);
99+
100+
if (result === TIMED_OUT) {
101+
pending.abort.abort();
102+
return {
103+
ok: false,
104+
output: JSON.stringify({ message: "Timed out" }),
105+
outputType: "application/json",
106+
};
107+
}
108+
109+
const output = typeof result === "string" ? result : JSON.stringify(result);
63110
return { ok: true, output, outputType: "application/json" };
64111
} catch {
65112
return {
66113
ok: false,
67114
output: JSON.stringify({ message: "Session stream wait ended before a record arrived" }),
68115
outputType: "application/json",
69116
};
117+
} finally {
118+
if (timer) {
119+
clearTimeout(timer);
120+
}
70121
}
71122
}
72123

0 commit comments

Comments
 (0)