test(webapp): agent e2e legs for tools, HITL, and snapshot restore

Author: ericallam

apps/webapp/test/helpers/agentHarness.ts

@@ -15,6 +15,13 @@ export type RunRealChatAgentOptions = {
   model: LanguageModel;
   modelLocal: LocalsKey<LanguageModel>;
   runId?: string;
+  /**
+   * Boot as a continuation of a previous run for the same session. Gates the
+   * snapshot + `.out`/`.in` replay boot path, so the agent restores prior
+   * history instead of treating the chat as brand new.
+   */
+  continuation?: boolean;
+  previousRunId?: string;
 };
 
 export type RunningAgent = {
@@ -54,10 +61,15 @@ export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
   const done = runInMockTaskContext(
     async (drivers) => {
       drivers.locals.set(opts.modelLocal, opts.model);
-      await runFn(
-        { chatId: opts.addressingKey, trigger: "preload", metadata: {} },
-        { ctx: drivers.ctx, signal: runSignal.signal }
-      );
+      const payload = opts.continuation
+        ? {
+            chatId: opts.addressingKey,
+            continuation: true,
+            metadata: {},
+            ...(opts.previousRunId ? { previousRunId: opts.previousRunId } : {}),
+          }
+        : { chatId: opts.addressingKey, trigger: "preload", metadata: {} };
+      await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal });
     },
     { ctx: { run: { id: runId } }, sessionStreamManager: manager }
   ) as Promise<void>;

apps/webapp/test/helpers/testChatAgent.ts

@@ -1,6 +1,6 @@
 import { chat } from "@trigger.dev/sdk/ai";
 import { locals } from "@trigger.dev/core/v3";
-import { streamText, type LanguageModel, type UIMessage } from "ai";
+import { stepCountIs, streamText, tool, type LanguageModel, type UIMessage } from "ai";
 import { z } from "zod";
 
 /**
@@ -61,3 +61,73 @@ export const testChatAgent = chat
       return streamText({ model, messages, abortSignal: signal });
     },
   });
+
+/**
+ * A minimal agent with no lifecycle hooks. With `hydrateMessages` absent, the
+ * default snapshot + `.out`/`.in` replay boot path is what restores prior
+ * history on a continuation run.
+ */
+export const testPlainChatAgent = chat.agent({
+  id: "e2e-test-chat-plain",
+  run: async ({ messages, signal }) => {
+    const model = locals.get(testChatModelLocal);
+    if (!model) {
+      throw new Error("test model not injected via locals");
+    }
+    return streamText({ model, messages, abortSignal: signal });
+  },
+});
+
+/**
+ * A tool with a server-side `execute`: the agent runs it automatically and
+ * feeds the result back to the model, so a single turn covers the whole
+ * tool loop.
+ */
+const weatherTool = tool({
+  description: "Get the current weather for a city.",
+  inputSchema: z.object({ city: z.string() }),
+  execute: async ({ city }) => ({ city, tempC: 21, summary: "clear" }),
+});
+
+export const testToolChatAgent = chat.agent({
+  id: "e2e-test-chat-tool",
+  run: async ({ messages, signal }) => {
+    const model = locals.get(testChatModelLocal);
+    if (!model) {
+      throw new Error("test model not injected via locals");
+    }
+    return streamText({
+      model,
+      messages,
+      tools: { getWeather: weatherTool },
+      stopWhen: stepCountIs(5),
+      abortSignal: signal,
+    });
+  },
+});
+
+/**
+ * A tool with no `execute`: the model's call parks the turn on a
+ * human-in-the-loop round-trip. The client supplies the tool output on the
+ * next message and the agent continues from there.
+ */
+const askUserTool = tool({
+  description: "Ask the user a question and wait for their answer.",
+  inputSchema: z.object({ question: z.string() }),
+});
+
+export const testHitlChatAgent = chat.agent({
+  id: "e2e-test-chat-hitl",
+  run: async ({ messages, signal }) => {
+    const model = locals.get(testChatModelLocal);
+    if (!model) {
+      throw new Error("test model not injected via locals");
+    }
+    return streamText({
+      model,
+      messages,
+      tools: { askUser: askUserTool },
+      abortSignal: signal,
+    });
+  },
+});