test(core,webapp): in-process waitpoint backend for run-lifecycle e2e

Extends the chat.agent e2e harness to cover the run-lifecycle scenarios that previously needed the run-engine and a supervisor: idle-timeout suspend and resume, chat.endRun() continuation, and chat.requestUpgrade().

session.in.wait() suspends a run on a waitpoint whose only task-visible effect is that runtime.waitUntil() eventually resolves with the next .in record. A TestRuntimeManager plus a SessionWaitpointBackend reproduce that in one process: the two waitpoint apiClient calls are stubbed, and the backend opens its own tail on the session channel and resolves the wait with the next record, so the same run() invocation continues in place (matching local-dev warm resume and the task-observable semantics of a deployed CRIU restore). endRun/requestUpgrade exit the run; a small session orchestrator then spawns the next run as a continuation, mirroring the server re-triggering on the next append.

It deliberately does not model server-side waitpoint bookkeeping or process checkpoint/restore, none of which are visible to task code on the resume path.

Author: ericallam

apps/webapp/test/helpers/agentHarness.ts

@@ -1,7 +1,11 @@
 import { apiClientManager, resourceCatalog } from "@trigger.dev/core/v3";
 import type { LocalsKey } from "@trigger.dev/core/v3";
 import type { LanguageModel } from "ai";
-import { runInMockTaskContext, StandardSessionStreamManager } from "@trigger.dev/core/v3/test";
+import {
+  installSessionWaitpointBackend,
+  runInMockTaskContext,
+  StandardSessionStreamManager,
+} from "@trigger.dev/core/v3/test";
 
 export type RunRealChatAgentOptions = {
   agentId: string;
@@ -22,6 +26,12 @@ export type RunRealChatAgentOptions = {
    */
   continuation?: boolean;
   previousRunId?: string;
+  /**
+   * Idle window (seconds) before the turn loop falls through from the SSE
+   * once() to the suspending `session.in.wait()`. Set this low to force the
+   * suspend/resume path in a test.
+   */
+  idleTimeoutInSeconds?: number;
 };
 
 export type RunningAgent = {
@@ -45,9 +55,11 @@ export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
   });
   const apiClient = apiClientManager.clientOrThrow();
   const manager = new StandardSessionStreamManager(apiClient, opts.baseUrl);
+  const { runtimeManager, restore } = installSessionWaitpointBackend(apiClient);
 
   const taskEntry = resourceCatalog.getTask(opts.agentId);
   if (!taskEntry) {
+    restore();
     throw new Error(`runRealChatAgent: agent "${opts.agentId}" is not registered`);
   }
   const runFn = taskEntry.fns.run as (
@@ -58,21 +70,29 @@ export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
   const runSignal = new AbortController();
   const runId = opts.runId ?? `run_${opts.addressingKey}`;
 
-  const done = runInMockTaskContext(
-    async (drivers) => {
-      drivers.locals.set(opts.modelLocal, opts.model);
-      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>;
+  const idle =
+    opts.idleTimeoutInSeconds !== undefined
+      ? { idleTimeoutInSeconds: opts.idleTimeoutInSeconds }
+      : {};
+
+  const done = (
+    runInMockTaskContext(
+      async (drivers) => {
+        drivers.locals.set(opts.modelLocal, opts.model);
+        const payload = opts.continuation
+          ? {
+              chatId: opts.addressingKey,
+              continuation: true,
+              metadata: {},
+              ...idle,
+              ...(opts.previousRunId ? { previousRunId: opts.previousRunId } : {}),
+            }
+          : { chatId: opts.addressingKey, trigger: "preload", metadata: {}, ...idle };
+        await runFn(payload, { ctx: drivers.ctx, signal: runSignal.signal });
+      },
+      { ctx: { run: { id: runId } }, sessionStreamManager: manager, runtimeManager }
+    ) as Promise<void>
+  ).finally(restore);
 
   return {
     done,
@@ -102,3 +122,61 @@ export function runRealChatAgent(opts: RunRealChatAgentOptions): RunningAgent {
     },
   };
 }
+
+export type ChatAgentSessionOptions = Omit<
+  RunRealChatAgentOptions,
+  "runId" | "continuation" | "previousRunId"
+>;
+
+export type ChatAgentSession = {
+  /** How many runs the session has spawned so far (1 fresh + N continuations). */
+  runCount: () => number;
+  /** Close the currently-active run and stop spawning continuations. */
+  close: () => Promise<void>;
+};
+
+/**
+ * A session-scoped orchestrator that stands in for the run-engine's run
+ * lifecycle: it starts a run, and whenever that run exits on its own
+ * (`chat.endRun()` / `chat.requestUpgrade()`), spawns the next run as a
+ * continuation (new run id, `continuation: true`, `previousRunId` threaded)
+ * for the same session. That mirrors the server triggering a fresh run on the
+ * next append after the previous run went terminal, and lets each continuation
+ * restore prior history from the persisted snapshot. Runs never overlap: the
+ * next spawn is chained on the previous run's `done` (after its manager
+ * teardown), so the process-global managers are never installed twice at once.
+ */
+export function runChatAgentSession(opts: ChatAgentSessionOptions): ChatAgentSession {
+  let closed = false;
+  let index = 0;
+  let current: RunningAgent | undefined;
+  let previousRunId: string | undefined;
+
+  const spawn = () => {
+    index += 1;
+    const runId = `run_${opts.addressingKey}_${index}`;
+    current = runRealChatAgent({
+      ...opts,
+      runId,
+      continuation: index > 1,
+      previousRunId,
+    });
+    previousRunId = runId;
+    const settle = () => {
+      if (!closed) {
+        spawn();
+      }
+    };
+    current.done.then(settle, settle);
+  };
+
+  spawn();
+
+  return {
+    runCount: () => index,
+    close: async () => {
+      closed = true;
+      await current?.close();
+    },
+  };
+}

apps/webapp/test/helpers/sessionStream.ts

@@ -92,6 +92,10 @@ export function isTurnComplete(part: CollectedPart): boolean {
   return (part.headers ?? []).some(([k, v]) => k === "trigger-control" && v === "turn-complete");
 }
 
+export function isUpgradeRequired(part: CollectedPart): boolean {
+  return (part.headers ?? []).some(([k, v]) => k === "trigger-control" && v === "upgrade-required");
+}
+
 export type SubscribeOptions = {
   baseUrl: string;
   addressingKey: string;

apps/webapp/test/helpers/testChatAgent.ts

@@ -78,6 +78,65 @@ export const testPlainChatAgent = chat.agent({
   },
 });
 
+/**
+ * A plain agent with a 1-second idle window, so the turn loop falls through to
+ * the suspending `session.in.wait()` almost immediately instead of catching the
+ * next message in the warm once() window. Used to exercise the suspend/resume
+ * waitpoint path.
+ */
+export const testIdleChatAgent = chat.agent({
+  id: "e2e-test-chat-idle",
+  idleTimeoutInSeconds: 1,
+  preloadIdleTimeoutInSeconds: 1,
+  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 });
+  },
+});
+
+/**
+ * Ends the run after every turn via `chat.endRun()`. The next message on the
+ * same chat starts a fresh continuation run (the orchestrator drives that).
+ */
+export const testEndRunChatAgent = chat.agent({
+  id: "e2e-test-chat-endrun",
+  idleTimeoutInSeconds: 2,
+  preloadIdleTimeoutInSeconds: 2,
+  run: async ({ messages, signal }) => {
+    const model = locals.get(testChatModelLocal);
+    if (!model) {
+      throw new Error("test model not injected via locals");
+    }
+    const result = streamText({ model, messages, abortSignal: signal });
+    chat.endRun();
+    return result;
+  },
+});
+
+/**
+ * Requests an upgrade from `onTurnStart` (the pre-turn path): `run()` is
+ * skipped, an `upgrade-required` control record lands on `.out`, and the run
+ * exits so a fresh run on the new version handles the message.
+ */
+export const testUpgradeChatAgent = chat.agent({
+  id: "e2e-test-chat-upgrade",
+  idleTimeoutInSeconds: 2,
+  preloadIdleTimeoutInSeconds: 2,
+  onTurnStart: async () => {
+    chat.requestUpgrade();
+  },
+  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

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

@@ -22,20 +22,32 @@ import {
   appendInput,
   collectSessionOut,
   isTurnComplete,
+  isUpgradeRequired,
   mintSessionToken,
   subscribeSessionOut,
   type CollectedPart,
 } from "./helpers/sessionStream";
-import { runRealChatAgent } from "./helpers/agentHarness";
+import { runChatAgentSession, runRealChatAgent } from "./helpers/agentHarness";
 import {
   testApprovalChatAgent,
   testChatAgent,
   testChatModelLocal,
+  testEndRunChatAgent,
   testHitlChatAgent,
+  testIdleChatAgent,
   testPlainChatAgent,
   testToolChatAgent,
+  testUpgradeChatAgent,
 } from "./helpers/testChatAgent";
 
+async function waitFor(predicate: () => boolean, maxMs: number): Promise<void> {
+  const deadline = performance.now() + maxMs;
+  while (performance.now() < deadline) {
+    if (predicate()) return;
+    await new Promise((r) => setTimeout(r, 100));
+  }
+}
+
 vi.setConfig({ testTimeout: 120_000, hookTimeout: 240_000 });
 
 let server: SessionStreamTestServer;
@@ -965,4 +977,143 @@ describe("session agent e2e (real chat.agent loop)", () => {
     expect(outB).toContain("answer-for-B");
     expect(outB, "B's stream never sees A's output").not.toContain("answer-for-A");
   });
+
+  it("EA13: the run suspends on the idle waitpoint, then the next message resumes it in place", async () => {
+    const { addressingKey, token, apiKey, baseUrl } = await setupSession(testIdleChatAgent.id);
+    const agent = runRealChatAgent({
+      agentId: testIdleChatAgent.id,
+      baseUrl,
+      addressingKey,
+      secretKey: apiKey,
+      model: textModel("resumed after suspend"),
+      modelLocal: testChatModelLocal,
+    });
+
+    try {
+      await new Promise((r) => setTimeout(r, 2500));
+      await appendInput({
+        baseUrl,
+        addressingKey,
+        token,
+        partId: "u1",
+        body: submitBody(addressingKey, userMessage("hi after idle", "u1")),
+      });
+
+      const { parts } = await collectSessionOut({
+        baseUrl,
+        addressingKey,
+        token,
+        until: (p) => p.some(isTurnComplete),
+        maxMs: 30_000,
+      });
+
+      expect(
+        joinChunks(parts),
+        "a message sent well after the idle window still produces a turn, so the waitpoint resume delivered it"
+      ).toContain("resumed after suspend");
+      expect(parts.some(isTurnComplete)).toBe(true);
+    } finally {
+      await agent.close();
+    }
+  });
+
+  it("EA14: chat.endRun() ends the run; the next message continues on a fresh run", async () => {
+    const { addressingKey, token, apiKey, baseUrl } = await setupSession(testEndRunChatAgent.id);
+    const session = runChatAgentSession({
+      agentId: testEndRunChatAgent.id,
+      baseUrl,
+      addressingKey,
+      secretKey: apiKey,
+      model: echoModel(),
+      modelLocal: testChatModelLocal,
+    });
+
+    try {
+      await appendInput({
+        baseUrl,
+        addressingKey,
+        token,
+        partId: "u1",
+        body: submitBody(addressingKey, userMessage("MARKER-ONE", "u1")),
+      });
+      await collectSessionOut({
+        baseUrl,
+        addressingKey,
+        token,
+        until: (p) => p.some(isTurnComplete),
+        maxMs: 30_000,
+      });
+
+      await waitFor(() => session.runCount() >= 2, 10_000);
+      expect(
+        session.runCount(),
+        "endRun exited run 1, so the orchestrator spawned a continuation"
+      ).toBeGreaterThanOrEqual(2);
+
+      await appendInput({
+        baseUrl,
+        addressingKey,
+        token,
+        partId: "u2",
+        body: submitBody(addressingKey, userMessage("second msg", "u2")),
+      });
+      const { parts } = await collectSessionOut({
+        baseUrl,
+        addressingKey,
+        token,
+        until: (p) => p.filter(isTurnComplete).length >= 2,
+        maxMs: 30_000,
+      });
+
+      const blob = joinChunks(parts);
+      expect(blob, "the continuation run restored run 1's history").toContain("MARKER-ONE");
+      expect(blob, "the continuation run ran the new turn").toContain("second msg");
+    } finally {
+      await session.close();
+    }
+  });
+
+  it("EA15: chat.requestUpgrade() emits upgrade-required on .out and exits the run", async () => {
+    const { addressingKey, token, apiKey, baseUrl } = await setupSession(testUpgradeChatAgent.id);
+    const agent = runRealChatAgent({
+      agentId: testUpgradeChatAgent.id,
+      baseUrl,
+      addressingKey,
+      secretKey: apiKey,
+      model: echoModel(),
+      modelLocal: testChatModelLocal,
+    });
+
+    let exitedOnItsOwn = false;
+    agent.done.then(() => {
+      exitedOnItsOwn = true;
+    });
+
+    try {
+      await appendInput({
+        baseUrl,
+        addressingKey,
+        token,
+        partId: "u1",
+        body: submitBody(addressingKey, userMessage("upgrade me", "u1")),
+      });
+      const { parts } = await collectSessionOut({
+        baseUrl,
+        addressingKey,
+        token,
+        until: (p) => p.some(isUpgradeRequired),
+        maxMs: 30_000,
+      });
+
+      expect(
+        parts.some(isUpgradeRequired),
+        "an upgrade-required control record is written to .out"
+      ).toBe(true);
+
+      await waitFor(() => exitedOnItsOwn, 10_000);
+      expect(exitedOnItsOwn, "the run exits itself after requesting the upgrade").toBe(true);
+    } finally {
+      await agent.close();
+    }
+  });
 });