Skip to content

Commit e1a416e

Browse files
gtrempergtremper
committed
test(chat): cover endAndContinue with real sessions
1 parent 03b84d1 commit e1a416e

3 files changed

Lines changed: 161 additions & 184 deletions

File tree

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -234,6 +234,41 @@ export const testUpgradeOnceChatAgent = chat.agent({
234234
},
235235
});
236236

237+
/**
238+
* Hands an unconsumed Session input record to a continuation run using the
239+
* public custom-agent lifecycle primitive. The continuation echoes the input
240+
* to `.out`, which lets the full-stack Session E2E assert durable delivery.
241+
*/
242+
export const testEndAndContinueCustomAgent = chat.customAgent({
243+
id: "e2e-test-chat-custom-end-and-continue",
244+
run: async (payload) => {
245+
if (!payload.continuation) {
246+
await chat.endAndContinue();
247+
return;
248+
}
249+
250+
const next = await chat.messages.waitWithIdleTimeout({
251+
idleTimeoutInSeconds: 2,
252+
timeout: "1m",
253+
});
254+
if (!next.ok) {
255+
throw next.error;
256+
}
257+
258+
const message = next.output.message as UIMessage | undefined;
259+
const text = message ? firstText(message) : "";
260+
const { waitUntilComplete } = chat.stream.writer({
261+
execute: ({ write }) => {
262+
write({ type: "text-start", id: "handoff-result" });
263+
write({ type: "text-delta", id: "handoff-result", delta: `received:${text}` });
264+
write({ type: "text-end", id: "handoff-result" });
265+
},
266+
});
267+
await waitUntilComplete();
268+
await chat.writeTurnComplete();
269+
},
270+
});
271+
237272
/**
238273
* A tool with a server-side `execute`: the agent runs it automatically and
239274
* feeds the result back to the model, so a single turn covers the whole

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

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import {
3333
testApprovalChatAgent,
3434
testChatAgent,
3535
testChatModelLocal,
36+
testEndAndContinueCustomAgent,
3637
testEndRunChatAgent,
3738
testHitlChatAgent,
3839
testHitlIdleChatAgent,
@@ -123,6 +124,34 @@ async function setupSession(agentId: string = testChatAgent.id) {
123124
return { addressingKey, token, apiKey, baseUrl: server.webapp.baseUrl };
124125
}
125126

127+
async function setupStartedSession(agentId: string) {
128+
const { environment, apiKey } = await seedTestEnvironment(server.prisma);
129+
const addressingKey = `chat-${randomBytes(6).toString("hex")}`;
130+
const createRes = await fetch(`${server.webapp.baseUrl}/api/v1/sessions`, {
131+
method: "POST",
132+
headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
133+
body: JSON.stringify({
134+
type: "chat.agent",
135+
externalId: addressingKey,
136+
taskIdentifier: agentId,
137+
triggerConfig: { basePayload: {} },
138+
}),
139+
});
140+
141+
expect(createRes.ok).toBe(true);
142+
const created = (await createRes.json()) as {
143+
runId: string;
144+
publicAccessToken: string;
145+
};
146+
return {
147+
...created,
148+
addressingKey,
149+
apiKey,
150+
environment,
151+
baseUrl: server.webapp.baseUrl,
152+
};
153+
}
154+
126155
function promptText(prompt: unknown): string {
127156
if (!Array.isArray(prompt)) return "";
128157
let out = "";
@@ -1533,4 +1562,97 @@ describe("session agent e2e (real chat.agent loop)", () => {
15331562
await agent.close();
15341563
}
15351564
});
1565+
1566+
it("EA23: custom endAndContinue hands pending input to a fresh run", async () => {
1567+
const { addressingKey, publicAccessToken, runId, apiKey, environment, baseUrl } =
1568+
await setupStartedSession(testEndAndContinueCustomAgent.id);
1569+
const initialRun = await server.prisma.taskRun.findFirstOrThrow({
1570+
where: { friendlyId: runId },
1571+
select: { id: true },
1572+
});
1573+
1574+
const append = await appendInput({
1575+
baseUrl,
1576+
addressingKey,
1577+
token: publicAccessToken,
1578+
partId: "pending-handoff-input",
1579+
body: submitBody(
1580+
addressingKey,
1581+
userMessage("deliver after endAndContinue", "pending-handoff-input")
1582+
),
1583+
});
1584+
expect(append.status).toBe(200);
1585+
1586+
const oldRun = runRealChatAgent({
1587+
agentId: testEndAndContinueCustomAgent.id,
1588+
baseUrl,
1589+
addressingKey,
1590+
secretKey: apiKey,
1591+
model: textModel("unused"),
1592+
modelLocal: testChatModelLocal,
1593+
runId,
1594+
});
1595+
let continuation: ReturnType<typeof runRealChatAgent> | undefined;
1596+
1597+
try {
1598+
await expect(oldRun.done).resolves.toBeUndefined();
1599+
1600+
const session = await server.prisma.session.findFirstOrThrow({
1601+
where: { runtimeEnvironmentId: environment.id, externalId: addressingKey },
1602+
select: { currentRunId: true, currentRunVersion: true },
1603+
});
1604+
expect(session.currentRunId).not.toBe(initialRun.id);
1605+
expect(session.currentRunVersion).toBeGreaterThan(1);
1606+
1607+
const successor = await server.prisma.taskRun.findUniqueOrThrow({
1608+
where: { id: session.currentRunId! },
1609+
select: { friendlyId: true },
1610+
});
1611+
continuation = runRealChatAgent({
1612+
agentId: testEndAndContinueCustomAgent.id,
1613+
baseUrl,
1614+
addressingKey,
1615+
secretKey: apiKey,
1616+
model: textModel("unused"),
1617+
modelLocal: testChatModelLocal,
1618+
runId: successor.friendlyId,
1619+
continuation: true,
1620+
previousRunId: runId,
1621+
});
1622+
1623+
const { parts } = await collectSessionOut({
1624+
baseUrl,
1625+
addressingKey,
1626+
token: publicAccessToken,
1627+
until: (p) => p.some(isTurnComplete),
1628+
maxMs: 30_000,
1629+
});
1630+
expect(joinChunks(parts)).toContain("received:deliver after endAndContinue");
1631+
await expect(continuation.done).resolves.toBeUndefined();
1632+
} finally {
1633+
await continuation?.close();
1634+
await oldRun.close();
1635+
}
1636+
});
1637+
1638+
it("EA24: custom endAndContinue rejects when the server rejects the handoff", async () => {
1639+
const { addressingKey, apiKey, baseUrl } = await setupStartedSession(
1640+
testEndAndContinueCustomAgent.id
1641+
);
1642+
const agent = runRealChatAgent({
1643+
agentId: testEndAndContinueCustomAgent.id,
1644+
baseUrl,
1645+
addressingKey,
1646+
secretKey: apiKey,
1647+
model: textModel("unused"),
1648+
modelLocal: testChatModelLocal,
1649+
runId: "run_missing_end_and_continue",
1650+
});
1651+
1652+
try {
1653+
await expect(agent.done).rejects.toThrow("callingRunId not found in this environment");
1654+
} finally {
1655+
await agent.close();
1656+
}
1657+
});
15361658
});

‎packages/trigger-sdk/test/chat-end-and-continue.test.ts‎

Lines changed: 4 additions & 184 deletions
Original file line numberDiff line numberDiff line change
@@ -1,190 +1,10 @@
1-
// Import the test entry point first so chat.customAgent() registers its task.
2-
import "../src/v3/test/index.js";
3-
4-
import { afterEach, describe, expect, it, vi } from "vitest";
5-
import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3";
6-
import { runInMockTaskContext, TestSessionStreamManager } from "@trigger.dev/core/v3/test";
1+
import { describe, expect, it } from "vitest";
72
import { chat } from "../src/v3/ai.js";
83

9-
const CHAT_ID = "chat-end-and-continue";
10-
const CALLING_RUN_ID = "run_before_handoff";
11-
const CONTINUATION_RUN_ID = "run_after_handoff";
12-
13-
type CustomAgentRun = (
14-
payload: Record<string, unknown>,
15-
options: { ctx: unknown; signal: AbortSignal }
16-
) => Promise<unknown>;
17-
18-
function getCustomAgentRun(id: string): CustomAgentRun {
19-
const taskEntry = resourceCatalog.getTask(id);
20-
if (!taskEntry) {
21-
throw new Error(`Task ${id} was not registered`);
22-
}
23-
24-
return taskEntry.fns.run as CustomAgentRun;
25-
}
26-
27-
class DurableTestSessionStreamManager extends TestSessionStreamManager {
28-
override reset(): void {
29-
// The Session stream outlives either task run. Drop run-local listeners,
30-
// but preserve buffered input for the continuation run.
31-
this.clearHandlers();
32-
}
33-
34-
dispose(): void {
35-
super.reset();
36-
}
37-
}
38-
394
describe("chat.endAndContinue", () => {
40-
afterEach(() => {
41-
vi.restoreAllMocks();
42-
});
43-
44-
it("ends cleanly and leaves unconsumed input for the continuation run", async () => {
45-
let continuationMessage: unknown;
46-
47-
const agent = chat.customAgent({
48-
id: "end-and-continue-custom-agent",
49-
run: async (payload) => {
50-
if (!payload.continuation) {
51-
return chat.endAndContinue();
52-
}
53-
54-
const next = await chat.messages.waitWithIdleTimeout({
55-
idleTimeoutInSeconds: 1,
56-
timeout: "1m",
57-
});
58-
if (!next.ok) {
59-
throw next.error;
60-
}
61-
62-
continuationMessage = next.output.message;
63-
},
64-
});
65-
66-
const runFn = getCustomAgentRun(agent.id);
67-
68-
const readSessionStreamRecords = vi.fn(async () => ({ records: [] }));
69-
const endAndContinueSession = vi.fn(async () => ({
70-
runId: CONTINUATION_RUN_ID,
71-
swapped: true,
72-
}));
73-
vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({
74-
readSessionStreamRecords,
75-
endAndContinueSession,
76-
} as never);
77-
78-
const sessionStreams = new DurableTestSessionStreamManager();
79-
const pendingPayload = {
80-
chatId: CHAT_ID,
81-
trigger: "submit-message",
82-
message: {
83-
id: "pending-user-message",
84-
role: "user",
85-
parts: [{ type: "text", text: "deliver after handoff" }],
86-
},
87-
metadata: {},
88-
};
89-
90-
try {
91-
await runInMockTaskContext(
92-
async (drivers) => {
93-
// This record is durable Session input, not run-local input. It is
94-
// written before the old run requests its handoff.
95-
await drivers.sessions.in.send(CHAT_ID, {
96-
kind: "message",
97-
payload: pendingPayload,
98-
});
99-
100-
await expect(
101-
runFn(
102-
{ chatId: CHAT_ID, trigger: "preload", metadata: {} },
103-
{ ctx: drivers.ctx, signal: new AbortController().signal }
104-
)
105-
).resolves.toBeUndefined();
106-
},
107-
{
108-
ctx: { run: { id: CALLING_RUN_ID } },
109-
sessionStreamManager: sessionStreams,
110-
}
111-
);
112-
113-
expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, {
114-
callingRunId: CALLING_RUN_ID,
115-
reason: "upgrade",
116-
});
117-
118-
await runInMockTaskContext(
119-
async (drivers) => {
120-
await expect(
121-
runFn(
122-
{ chatId: CHAT_ID, continuation: true, metadata: {} },
123-
{ ctx: drivers.ctx, signal: new AbortController().signal }
124-
)
125-
).resolves.toBeUndefined();
126-
},
127-
{
128-
ctx: { run: { id: CONTINUATION_RUN_ID } },
129-
sessionStreamManager: sessionStreams,
130-
}
131-
);
132-
133-
expect(continuationMessage).toEqual(pendingPayload.message);
134-
} finally {
135-
sessionStreams.dispose();
136-
}
137-
});
138-
139-
it("rejects when the server handoff fails", async () => {
140-
const agent = chat.customAgent({
141-
id: "end-and-continue-failure-agent",
142-
run: async () => {
143-
return chat.endAndContinue();
144-
},
145-
});
146-
147-
const runFn = getCustomAgentRun(agent.id);
148-
149-
const readSessionStreamRecords = vi.fn(async () => ({ records: [] }));
150-
const endAndContinueSession = vi.fn(async () => {
151-
throw new Error("handoff failed");
152-
});
153-
vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({
154-
readSessionStreamRecords,
155-
endAndContinueSession,
156-
} as never);
157-
158-
await runInMockTaskContext(
159-
async (drivers) => {
160-
await expect(
161-
runFn(
162-
{ chatId: CHAT_ID, trigger: "preload", metadata: {} },
163-
{ ctx: drivers.ctx, signal: new AbortController().signal }
164-
)
165-
).rejects.toThrow("handoff failed");
166-
},
167-
{ ctx: { run: { id: CALLING_RUN_ID } } }
168-
);
169-
170-
expect(endAndContinueSession).toHaveBeenCalledWith(CHAT_ID, {
171-
callingRunId: CALLING_RUN_ID,
172-
reason: "upgrade",
173-
});
174-
});
175-
1765
it("rejects calls outside a custom agent run", async () => {
177-
const endAndContinueSession = vi.fn();
178-
vi.spyOn(apiClientManager, "clientOrThrow").mockReturnValue({
179-
endAndContinueSession,
180-
} as never);
181-
182-
await runInMockTaskContext(async () => {
183-
await expect(chat.endAndContinue()).rejects.toThrow(
184-
"chat.endAndContinue() can only be called from inside a chat.customAgent() run"
185-
);
186-
});
187-
188-
expect(endAndContinueSession).not.toHaveBeenCalled();
6+
await expect(chat.endAndContinue()).rejects.toThrow(
7+
"chat.endAndContinue() can only be called from inside a chat.customAgent() run"
8+
);
1899
});
19010
});

0 commit comments

Comments
 (0)