feat(core): add runInMockTaskContext test infrastructure

In-memory managers for locals, lifecycle hooks, runtime, input streams,
and realtime streams, plus a mock TaskContext. Lets task code be driven
end-to-end without hitting the Trigger.dev runtime — send data into
input streams, inspect chunks written to output streams, and pre-seed
locals for dependency injection.

Author: ericallam

.changeset/mock-task-context-test-infra.md

@@ -0,0 +1,5 @@
+---
+"@trigger.dev/core": patch
+---
+
+Add `runInMockTaskContext` test harness at `@trigger.dev/core/v3/test` for unit-testing task code offline. Installs in-memory managers for `locals`, `lifecycleHooks`, `runtime`, `inputStreams`, and `realtimeStreams`, plus a mock `TaskContext`, so tasks can be driven end-to-end without hitting the Trigger.dev runtime. Provides drivers to send data into input streams and inspect chunks written to output streams.

packages/core/package.json

@@ -42,6 +42,7 @@
       "./v3/utils/retries": "./src/v3/utils/retries.ts",
       "./v3/utils/structuredLogger": "./src/v3/utils/structuredLogger.ts",
       "./v3/chat-client": "./src/v3/chat-client.ts",
+      "./v3/test": "./src/v3/test/index.ts",
       "./v3/zodfetch": "./src/v3/zodfetch.ts",
       "./v3/zodMessageHandler": "./src/v3/zodMessageHandler.ts",
       "./v3/zodNamespace": "./src/v3/zodNamespace.ts",
@@ -153,6 +154,9 @@
       ],
       "v3/isomorphic": [
         "dist/commonjs/v3/isomorphic/index.d.ts"
+      ],
+      "v3/test": [
+        "dist/commonjs/v3/test/index.d.ts"
       ]
     }
   },
@@ -458,6 +462,17 @@
         "default": "./dist/commonjs/v3/chat-client.js"
       }
     },
+    "./v3/test": {
+      "import": {
+        "@triggerdotdev/source": "./src/v3/test/index.ts",
+        "types": "./dist/esm/v3/test/index.d.ts",
+        "default": "./dist/esm/v3/test/index.js"
+      },
+      "require": {
+        "types": "./dist/commonjs/v3/test/index.d.ts",
+        "default": "./dist/commonjs/v3/test/index.js"
+      }
+    },
     "./v3/zodfetch": {
       "import": {
         "@triggerdotdev/source": "./src/v3/zodfetch.ts",

packages/core/src/v3/test/index.ts

@@ -0,0 +1,8 @@
+export {
+  runInMockTaskContext,
+  type MockTaskContextDrivers,
+  type MockTaskContextOptions,
+} from "./mock-task-context.js";
+export { TestInputStreamManager } from "./test-input-stream-manager.js";
+export { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js";
+export { TestRunMetadataManager } from "./test-run-metadata-manager.js";

packages/core/src/v3/test/mock-task-context.ts

@@ -0,0 +1,260 @@
+import { inputStreams } from "../input-streams-api.js";
+import { realtimeStreams } from "../realtime-streams-api.js";
+import { localsAPI } from "../locals-api.js";
+import { runMetadata } from "../run-metadata-api.js";
+import { taskContext } from "../task-context-api.js";
+import { lifecycleHooks } from "../lifecycle-hooks-api.js";
+import { runtime } from "../runtime-api.js";
+import { StandardLocalsManager } from "../locals/manager.js";
+import { StandardLifecycleHooksManager } from "../lifecycleHooks/manager.js";
+import { NoopRuntimeManager } from "../runtime/noopRuntimeManager.js";
+import { unregisterGlobal } from "../utils/globals.js";
+import type { ServerBackgroundWorker, TaskRunContext } from "../schemas/index.js";
+import type { LocalsKey } from "../locals/types.js";
+import { TestInputStreamManager } from "./test-input-stream-manager.js";
+import { TestRealtimeStreamsManager } from "./test-realtime-streams-manager.js";
+import { TestRunMetadataManager } from "./test-run-metadata-manager.js";
+
+/**
+ * Shallow-partial overrides applied on top of the default mock
+ * `TaskRunContext`. Each sub-object is a partial of its real shape —
+ * unset fields get sensible defaults.
+ */
+export type MockTaskRunContextOverrides = {
+  task?: Partial<TaskRunContext["task"]>;
+  attempt?: Partial<TaskRunContext["attempt"]>;
+  run?: Partial<TaskRunContext["run"]>;
+  machine?: Partial<TaskRunContext["machine"]>;
+  queue?: Partial<TaskRunContext["queue"]>;
+  environment?: Partial<TaskRunContext["environment"]>;
+  organization?: Partial<TaskRunContext["organization"]>;
+  project?: Partial<TaskRunContext["project"]>;
+  batch?: TaskRunContext["batch"];
+};
+
+/**
+ * Options for overriding parts of the mock task context.
+ */
+export type MockTaskContextOptions = {
+  /** Overrides applied on top of the default mock `TaskRunContext`. */
+  ctx?: MockTaskRunContextOverrides;
+  /** Overrides applied on top of the default `ServerBackgroundWorker`. */
+  worker?: Partial<ServerBackgroundWorker>;
+  /** Whether this is a warm start. */
+  isWarmStart?: boolean;
+};
+
+/**
+ * Drivers passed to the function running inside `runInMockTaskContext`.
+ */
+export type MockTaskContextDrivers = {
+  /** Push data into input streams — simulates realtime input from outside the task. */
+  inputs: {
+    /**
+     * Send `data` to the named input stream. Resolves when all `.on()`
+     * handlers have run.
+     */
+    send(streamId: string, data: unknown): Promise<void>;
+    /** Resolve any pending `.once()` waiters with a timeout error. */
+    close(streamId: string): void;
+  };
+  /** Inspect chunks written to output (realtime) streams. */
+  outputs: {
+    /** All chunks for a given stream, in the order they were written. */
+    chunks<T = unknown>(streamId: string): T[];
+    /** All chunks across every stream, keyed by stream id. */
+    all(): Record<string, unknown[]>;
+    /** Clear chunks for one stream, or all streams if no id is provided. */
+    clear(streamId?: string): void;
+    /**
+     * Register a listener fired for every chunk written to any stream.
+     * Returns an unsubscribe function.
+     */
+    onWrite(listener: (streamId: string, chunk: unknown) => void): () => void;
+  };
+  /** Read or seed locals for the run. */
+  locals: {
+    /** Read a local set by either the task or `set()` below. */
+    get<T>(key: LocalsKey<T>): T | undefined;
+    /**
+     * Pre-seed a local before the task runs. Use this for dependency
+     * injection — e.g. supply a test database client that the agent's
+     * hooks read via `locals.get()` instead of constructing the prod one.
+     */
+    set<T>(key: LocalsKey<T>, value: T): void;
+  };
+  /** The mock `TaskRunContext` assembled from defaults + user overrides. */
+  ctx: TaskRunContext;
+};
+
+function defaultTaskRunContext(overrides?: MockTaskRunContextOverrides): TaskRunContext {
+  return {
+    task: {
+      id: "test-task",
+      filePath: "test-task.ts",
+      ...overrides?.task,
+    },
+    attempt: {
+      number: 1,
+      startedAt: new Date(),
+      ...overrides?.attempt,
+    },
+    run: {
+      id: "run_test",
+      tags: [],
+      isTest: false,
+      createdAt: new Date(),
+      startedAt: new Date(),
+      ...overrides?.run,
+    },
+    machine: {
+      name: "micro",
+      cpu: 1,
+      memory: 0.5,
+      centsPerMs: 0,
+      ...overrides?.machine,
+    },
+    queue: {
+      name: "test-queue",
+      id: "test-queue-id",
+      ...overrides?.queue,
+    },
+    environment: {
+      id: "test-env-id",
+      slug: "test-env",
+      type: "DEVELOPMENT",
+      ...overrides?.environment,
+    },
+    organization: {
+      id: "test-org-id",
+      slug: "test-org",
+      name: "Test Org",
+      ...overrides?.organization,
+    },
+    project: {
+      id: "test-project-id",
+      ref: "test-project-ref",
+      slug: "test-project",
+      name: "Test Project",
+      ...overrides?.project,
+    },
+    batch: overrides?.batch,
+  };
+}
+
+function defaultWorker(overrides?: Partial<ServerBackgroundWorker>): ServerBackgroundWorker {
+  return {
+    id: "test-worker-id",
+    version: "test-version",
+    contentHash: "test-content-hash",
+    engine: "V2",
+    ...overrides,
+  };
+}
+
+/**
+ * Run a function inside a fully mocked task runtime context.
+ *
+ * Installs in-memory test managers for `locals`, `inputStreams`,
+ * `realtimeStreams`, `lifecycleHooks`, and `runtime`, sets a mock
+ * `TaskContext`, and tears everything down when the function returns.
+ *
+ * Inside the function, any code that reads from `locals`, `inputStreams`,
+ * `realtimeStreams`, or `taskContext.ctx` will see the mock context —
+ * so you can directly invoke the internal `run` function of any task
+ * (including `chat.agent`) without hitting the Trigger.dev runtime.
+ *
+ * @example
+ * ```ts
+ * import { runInMockTaskContext } from "@trigger.dev/core/v3/test";
+ *
+ * await runInMockTaskContext(
+ *   async ({ inputs, outputs, ctx }) => {
+ *     // Fire an input stream from the "outside"
+ *     setTimeout(() => {
+ *       inputs.send("chat-messages", { messages: [], chatId: "c1" });
+ *     }, 0);
+ *
+ *     // Run task code that reads from inputStreams.once(...)
+ *     await myTask.fns.run(payload, { ctx, signal: new AbortController().signal });
+ *
+ *     // Inspect chunks written to the output stream
+ *     expect(outputs.chunks("chat")).toContainEqual({ type: "text-delta", delta: "hi" });
+ *   },
+ *   { ctx: { run: { id: "run_abc" } } }
+ * );
+ * ```
+ */
+export async function runInMockTaskContext<T>(
+  fn: (drivers: MockTaskContextDrivers) => T | Promise<T>,
+  options?: MockTaskContextOptions
+): Promise<T> {
+  const ctx = defaultTaskRunContext(options?.ctx);
+  const worker = defaultWorker(options?.worker);
+
+  const localsManager = new StandardLocalsManager();
+  const lifecycleManager = new StandardLifecycleHooksManager();
+  const runtimeManager = new NoopRuntimeManager();
+  const metadataManager = new TestRunMetadataManager();
+  const inputManager = new TestInputStreamManager();
+  const outputManager = new TestRealtimeStreamsManager();
+
+  // Unregister any previously-installed managers so `setGlobal*` wins —
+  // `registerGlobal` returns false silently if an entry already exists.
+  unregisterGlobal("locals");
+  unregisterGlobal("lifecycle-hooks");
+  unregisterGlobal("runtime");
+  unregisterGlobal("run-metadata");
+  unregisterGlobal("input-streams");
+  unregisterGlobal("realtime-streams");
+  unregisterGlobal("task-context");
+
+  localsAPI.setGlobalLocalsManager(localsManager);
+  lifecycleHooks.setGlobalLifecycleHooksManager(lifecycleManager);
+  runtime.setGlobalRuntimeManager(runtimeManager);
+  runMetadata.setGlobalManager(metadataManager);
+  inputStreams.setGlobalManager(inputManager);
+  realtimeStreams.setGlobalManager(outputManager);
+  taskContext.setGlobalTaskContext({
+    ctx,
+    worker,
+    isWarmStart: options?.isWarmStart ?? false,
+  });
+
+  const drivers: MockTaskContextDrivers = {
+    inputs: {
+      send: (streamId, data) => inputManager.__sendFromTest(streamId, data),
+      close: (streamId) => inputManager.__closeFromTest(streamId),
+    },
+    outputs: {
+      chunks: (streamId) => outputManager.__chunksFromTest(streamId),
+      all: () => outputManager.__allChunksFromTest(),
+      clear: (streamId) => outputManager.__clearFromTest(streamId),
+      onWrite: (listener) => outputManager.onWrite(listener),
+    },
+    locals: {
+      get: <TValue>(key: LocalsKey<TValue>) => localsManager.getLocal(key),
+      set: <TValue>(key: LocalsKey<TValue>, value: TValue) =>
+        localsManager.setLocal(key, value),
+    },
+    ctx,
+  };
+
+  try {
+    return await fn(drivers);
+  } finally {
+    localsAPI.disable();
+    lifecycleHooks.disable();
+    runtime.disable();
+    // taskContext.disable() only sets a flag — unregister the global so
+    // `taskContext.ctx` returns undefined after the harness returns.
+    unregisterGlobal("task-context");
+    unregisterGlobal("input-streams");
+    unregisterGlobal("realtime-streams");
+    unregisterGlobal("run-metadata");
+    localsManager.reset();
+    inputManager.reset();
+    outputManager.reset();
+    metadataManager.reset();
+  }
+}