zeitlich 0.2.15 → 0.2.17

package/src/lib/subagent/workflow.ts

@@ -0,0 +1,114 @@
+import type { z } from "zod";
+import type {
+  SubagentDefinition,
+  SubagentHandlerResponse,
+  SubagentWorkflowInput,
+  SubagentSessionInput,
+} from "./types";
+
+/**
+ * Defines a subagent workflow with embedded metadata (name, description, resultSchema).
+ * The returned value can be passed directly to `defineSubagent` — no need to repeat
+ * the name, description, or resultSchema in the parent workflow.
+ *
+ * Internally maps `SubagentWorkflowInput` fields to session-compatible `SubagentSessionInput`.
+ *
+ * @example
+ * ```ts
+ * import {
+ *   defineSubagentWorkflow,
+ *   defineSubagent,
+ *   createSession,
+ *   createAgentStateManager,
+ * } from 'zeitlich/workflow';
+ *
+ * // Define once — carries name, description, resultSchema
+ * export const researcherWorkflow = defineSubagentWorkflow(
+ *   {
+ *     name: "researcher",
+ *     description: "Researches topics on the web",
+ *     resultSchema: z.object({ findings: z.string() }),
+ *   },
+ *   async (prompt, sessionInput) => {
+ *     const stateManager = createAgentStateManager({
+ *       initialState: { systemPrompt: "You are a researcher." },
+ *     });
+ *
+ *     const session = await createSession({
+ *       ...sessionInput,
+ *       agentName: "researcher",
+ *       runAgent: runAgentActivity,
+ *       buildContextMessage: () => [{ type: "text", text: prompt }],
+ *     });
+ *
+ *     const { finalMessage, threadId } = await session.runSession({ stateManager });
+ *     return { toolResponse: finalMessage ?? "No response", data: null, threadId };
+ *   },
+ * );
+ *
+ * // Use in parent — only configure what's parent-specific
+ * export const researcher = defineSubagent(researcherWorkflow, {
+ *   hooks: { onPostExecution: ({ result }) => console.log(result) },
+ * });
+ * ```
+ */
+// Without resultSchema — data is null
+export function defineSubagentWorkflow<
+  TContext extends Record<string, unknown> = Record<string, unknown>,
+>(
+  config: { name: string; description: string },
+  fn: (
+    prompt: string,
+    sessionInput: SubagentSessionInput,
+    context: TContext
+  ) => Promise<SubagentHandlerResponse<null>>
+): SubagentDefinition<z.ZodNull, TContext>;
+// With resultSchema — data is inferred from the schema
+export function defineSubagentWorkflow<
+  TResult extends z.ZodType,
+  TContext extends Record<string, unknown> = Record<string, unknown>,
+>(
+  config: { name: string; description: string; resultSchema: TResult },
+  fn: (
+    prompt: string,
+    sessionInput: SubagentSessionInput,
+    context: TContext
+  ) => Promise<SubagentHandlerResponse<z.infer<TResult> | null>>
+): SubagentDefinition<TResult, TContext>;
+export function defineSubagentWorkflow(
+  config: { name: string; description: string; resultSchema?: z.ZodType },
+  fn: (
+    prompt: string,
+    sessionInput: SubagentSessionInput,
+    context: Record<string, unknown>
+  ) => Promise<SubagentHandlerResponse<unknown>>
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+): SubagentDefinition<any, any> {
+  const workflow = async (
+    prompt: string,
+    workflowInput: SubagentWorkflowInput,
+    context?: Record<string, unknown>
+  ): Promise<SubagentHandlerResponse<unknown>> => {
+    const sessionInput: SubagentSessionInput = {
+      agentName: config.name,
+      ...(workflowInput.previousThreadId && {
+        threadId: workflowInput.previousThreadId,
+        continueThread: true,
+      }),
+      ...(workflowInput.sandboxId && { sandboxId: workflowInput.sandboxId }),
+    };
+    return fn(prompt, sessionInput, context ?? {});
+  };
+
+  // for temporal workflow name
+  Object.defineProperty(workflow, "name", { value: config.name });
+
+  return Object.assign(workflow, {
+    agentName: config.name,
+    description: config.description,
+    ...(config.resultSchema !== undefined && {
+      resultSchema: config.resultSchema,
+    }),
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  }) as SubagentDefinition<any, any>;
+}

package/src/lib/thread/id.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, it, vi, beforeEach } from "vitest";
+
+let uuidCounter = 0;
+
+vi.mock("@temporalio/workflow", () => ({
+  uuid4: () => {
+    uuidCounter++;
+    const bytes = Array.from({ length: 16 }, (_, i) =>
+      ((uuidCounter * 31 + i * 7 + uuidCounter * i) & 0xff)
+        .toString(16)
+        .padStart(2, "0"),
+    ).join("");
+    return `${bytes.slice(0, 8)}-${bytes.slice(8, 12)}-${bytes.slice(12, 16)}-${bytes.slice(16, 20)}-${bytes.slice(20, 32)}`;
+  },
+}));
+
+import { getShortId } from "./id";
+
+describe("getShortId", () => {
+  beforeEach(() => {
+    uuidCounter = 0;
+  });
+
+  it("returns a string of default length 12", () => {
+    const id = getShortId();
+    expect(id).toHaveLength(12);
+  });
+
+  it("returns a string of custom length", () => {
+    const id = getShortId(6);
+    expect(id).toHaveLength(6);
+  });
+
+  it("contains only base-62 characters", () => {
+    const base62Regex = /^[A-Za-z0-9]+$/;
+    for (let i = 0; i < 10; i++) {
+      expect(getShortId()).toMatch(base62Regex);
+    }
+  });
+
+  it("generates unique IDs on successive calls", () => {
+    const ids = new Set(Array.from({ length: 20 }, () => getShortId()));
+    expect(ids.size).toBe(20);
+  });
+
+  it("returns empty string for length 0", () => {
+    const id = getShortId(0);
+    expect(id).toBe("");
+  });
+});

package/src/lib/tool-router/auto-append-sandbox.integration.test.ts

@@ -0,0 +1,344 @@
+import { describe, expect, it, vi } from "vitest";
+import { withAutoAppend } from "./auto-append";
+import { withSandbox } from "./with-sandbox";
+import type { RouterContext, ToolHandlerResponse } from "./types";
+import type { ToolResultConfig } from "../types";
+import type { Sandbox } from "../sandbox/types";
+
+// ---------------------------------------------------------------------------
+// withAutoAppend
+// ---------------------------------------------------------------------------
+
+describe("withAutoAppend", () => {
+  it("appends tool result via threadHandler and sets resultAppended", async () => {
+    const appended: ToolResultConfig[] = [];
+    const threadHandler = async (config: ToolResultConfig) => {
+      appended.push(config);
+    };
+
+    const innerHandler = async (
+      args: { text: string },
+      _ctx: RouterContext,
+    ): Promise<ToolHandlerResponse<{ echoed: string }>> => ({
+      toolResponse: `Echo: ${args.text}`,
+      data: { echoed: args.text },
+    });
+
+    const wrapped = withAutoAppend(threadHandler, innerHandler);
+
+    const result = await wrapped(
+      { text: "hello" },
+      {
+        threadId: "thread-1",
+        toolCallId: "tc-1",
+        toolName: "Echo",
+      },
+    );
+
+    expect(result.resultAppended).toBe(true);
+    expect(result.toolResponse).toBe("Response appended via withAutoAppend");
+    expect(result.data).toEqual({ echoed: "hello" });
+
+    expect(appended).toHaveLength(1);
+    const firstAppended = appended.at(0);
+    if (!firstAppended) throw new Error("expected appended item");
+    expect(firstAppended.threadId).toBe("thread-1");
+    expect(firstAppended.toolCallId).toBe("tc-1");
+    expect(firstAppended.toolName).toBe("Echo");
+    expect(firstAppended.content).toBe("Echo: hello");
+  });
+
+  it("preserves original data but replaces toolResponse", async () => {
+    const threadHandler = async () => {};
+
+    const innerHandler = async (): Promise<
+      ToolHandlerResponse<{ large: string }>
+    > => ({
+      toolResponse: "A".repeat(10000),
+      data: { large: "summary" },
+    });
+
+    const wrapped = withAutoAppend(threadHandler, innerHandler);
+
+    const result = await wrapped(
+      {},
+      { threadId: "t", toolCallId: "tc", toolName: "BigTool" },
+    );
+
+    expect(result.toolResponse).toBe("Response appended via withAutoAppend");
+    expect(result.data).toEqual({ large: "summary" });
+    expect(result.resultAppended).toBe(true);
+  });
+
+  it("propagates handler errors without appending", async () => {
+    const appendSpy = vi.fn();
+    const threadHandler = appendSpy;
+
+    const innerHandler = async (): Promise<ToolHandlerResponse<null>> => {
+      throw new Error("handler failed");
+    };
+
+    const wrapped = withAutoAppend(threadHandler, innerHandler);
+
+    await expect(
+      wrapped(
+        {},
+        { threadId: "t", toolCallId: "tc", toolName: "Fail" },
+      ),
+    ).rejects.toThrow("handler failed");
+
+    expect(appendSpy).not.toHaveBeenCalled();
+  });
+
+  it("uses correct context fields for thread handler config", async () => {
+    let capturedConfig: ToolResultConfig | null = null;
+    const threadHandler = async (config: ToolResultConfig) => {
+      capturedConfig = config;
+    };
+
+    const innerHandler = async (): Promise<ToolHandlerResponse<null>> => ({
+      toolResponse: "result content",
+      data: null,
+    });
+
+    const wrapped = withAutoAppend(threadHandler, innerHandler);
+
+    await wrapped(
+      {},
+      {
+        threadId: "my-thread",
+        toolCallId: "my-tc",
+        toolName: "MyTool",
+        sandboxId: "sb-1",
+      },
+    );
+
+    expect(capturedConfig).toEqual({
+      threadId: "my-thread",
+      toolCallId: "my-tc",
+      toolName: "MyTool",
+      content: "result content",
+    });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// withSandbox
+// ---------------------------------------------------------------------------
+
+describe("withSandbox", () => {
+  function createMockSandbox(): Sandbox {
+    return {
+      id: "mock-sandbox",
+      capabilities: {
+        filesystem: true,
+        execution: true,
+        persistence: false,
+      },
+      fs: {
+        workspaceBase: "/",
+        exists: async () => false,
+        stat: async () => ({
+          isFile: false,
+          isDirectory: false,
+          isSymbolicLink: false,
+          size: 0,
+          mtime: new Date(),
+        }),
+        readdir: async () => [],
+        readdirWithFileTypes: async () => [],
+        readFile: async () => "",
+        readFileBuffer: async () => new Uint8Array(),
+        writeFile: async () => {},
+        appendFile: async () => {},
+        mkdir: async () => {},
+        rm: async () => {},
+        cp: async () => {},
+        mv: async () => {},
+        readlink: async () => "",
+        resolvePath: (base: string, path: string) => base + "/" + path,
+      },
+      exec: async () => ({ stdout: "", stderr: "", exitCode: 0 }),
+      destroy: async () => {},
+    };
+  }
+
+  it("resolves sandbox and passes it to handler", async () => {
+    const mockSandbox = createMockSandbox();
+    const manager = {
+      getSandbox: async (id: string) => {
+        expect(id).toBe("sb-42");
+        return mockSandbox;
+      },
+    };
+
+    let capturedSandbox: Sandbox | null = null;
+    let capturedSandboxId: string | null = null;
+
+    const handler = async (
+      _args: { text: string },
+      ctx: RouterContext & { sandbox: Sandbox; sandboxId: string },
+    ): Promise<ToolHandlerResponse<null>> => {
+      capturedSandbox = ctx.sandbox;
+      capturedSandboxId = ctx.sandboxId;
+      return { toolResponse: "ok", data: null };
+    };
+
+    const wrapped = withSandbox(manager, handler);
+
+    const result = await wrapped(
+      { text: "hello" },
+      {
+        threadId: "thread-1",
+        toolCallId: "tc-1",
+        toolName: "Test",
+        sandboxId: "sb-42",
+      },
+    );
+
+    expect(result.toolResponse).toBe("ok");
+    expect(capturedSandbox).toBe(mockSandbox);
+    expect(capturedSandboxId).toBe("sb-42");
+  });
+
+  it("returns error when no sandboxId is present", async () => {
+    const manager = {
+      getSandbox: vi.fn(),
+    };
+
+    const handler = async (): Promise<ToolHandlerResponse<string>> => ({
+      toolResponse: "should not run",
+      data: "nope",
+    });
+
+    const wrapped = withSandbox(manager, handler);
+
+    const result = await wrapped(
+      {},
+      {
+        threadId: "thread-1",
+        toolCallId: "tc-1",
+        toolName: "Bash",
+      },
+    );
+
+    expect(result.toolResponse).toContain("No sandbox configured");
+    expect(result.toolResponse).toContain("Bash");
+    expect(result.data).toBeNull();
+    expect(manager.getSandbox).not.toHaveBeenCalled();
+  });
+
+  it("returns error when sandboxId is undefined explicitly", async () => {
+    const manager = {
+      getSandbox: vi.fn(),
+    };
+
+    const handler = async (): Promise<ToolHandlerResponse<null>> => ({
+      toolResponse: "nope",
+      data: null,
+    });
+
+    const wrapped = withSandbox(manager, handler);
+
+    const result = await wrapped(
+      {},
+      {
+        threadId: "t",
+        toolCallId: "tc",
+        toolName: "Grep",
+        sandboxId: undefined,
+      },
+    );
+
+    expect(result.toolResponse).toContain("No sandbox configured");
+    expect(result.data).toBeNull();
+  });
+
+  it("propagates getSandbox errors", async () => {
+    const manager = {
+      getSandbox: async () => {
+        throw new Error("sandbox not found");
+      },
+    };
+
+    const handler = async (): Promise<ToolHandlerResponse<null>> => ({
+      toolResponse: "ok",
+      data: null,
+    });
+
+    const wrapped = withSandbox(manager, handler);
+
+    await expect(
+      wrapped(
+        {},
+        {
+          threadId: "t",
+          toolCallId: "tc",
+          toolName: "Test",
+          sandboxId: "sb-missing",
+        },
+      ),
+    ).rejects.toThrow("sandbox not found");
+  });
+
+  it("passes all RouterContext fields through to inner handler", async () => {
+    const mockSandbox = createMockSandbox();
+    const manager = { getSandbox: async () => mockSandbox };
+
+    let capturedCtx: (RouterContext & { sandbox: Sandbox; sandboxId: string }) | null = null;
+
+    const handler = async (
+      _args: unknown,
+      ctx: RouterContext & { sandbox: Sandbox; sandboxId: string },
+    ): Promise<ToolHandlerResponse<null>> => {
+      capturedCtx = ctx;
+      return { toolResponse: "ok", data: null };
+    };
+
+    const wrapped = withSandbox(manager, handler);
+
+    await wrapped(
+      {},
+      {
+        threadId: "my-thread",
+        toolCallId: "my-tc",
+        toolName: "MyTool",
+        sandboxId: "my-sandbox",
+      },
+    );
+
+    expect(capturedCtx).toEqual(
+      expect.objectContaining({
+        threadId: "my-thread",
+        toolCallId: "my-tc",
+        toolName: "MyTool",
+        sandboxId: "my-sandbox",
+        sandbox: mockSandbox,
+      }),
+    );
+  });
+
+  it("handles sandboxId with empty string as falsy", async () => {
+    const manager = { getSandbox: vi.fn() };
+
+    const handler = async (): Promise<ToolHandlerResponse<null>> => ({
+      toolResponse: "nope",
+      data: null,
+    });
+
+    const wrapped = withSandbox(manager, handler);
+
+    const result = await wrapped(
+      {},
+      {
+        threadId: "t",
+        toolCallId: "tc",
+        toolName: "Test",
+        sandboxId: "",
+      },
+    );
+
+    expect(result.toolResponse).toContain("No sandbox configured");
+    expect(manager.getSandbox).not.toHaveBeenCalled();
+  });
+});