zeitlich 0.2.37 → 0.2.39

package/src/adapters/thread/langchain/fork-transform.test.ts

@@ -0,0 +1,142 @@
+import { describe, expect, it, vi } from "vitest";
+import {
+  HumanMessage,
+  AIMessage,
+  type StoredMessage,
+} from "@langchain/core/messages";
+import { createLangChainThreadManager } from "./thread-manager";
+
+function createStatefulRedis() {
+  const lists = new Map<string, string[]>();
+  const strings = new Map<string, string>();
+
+  return {
+    exists: vi.fn(async (...keys: string[]) =>
+      keys.reduce(
+        (acc, k) => acc + (lists.has(k) || strings.has(k) ? 1 : 0),
+        0
+      )
+    ),
+    lrange: vi.fn(async (key: string, start: number, stop: number) => {
+      const list = lists.get(key) ?? [];
+      const end = stop === -1 ? list.length : stop + 1;
+      return list.slice(start, end);
+    }),
+    rpush: vi.fn(async (key: string, ...values: string[]) => {
+      const list = lists.get(key) ?? [];
+      list.push(...values);
+      lists.set(key, list);
+      return list.length;
+    }),
+    ltrim: vi.fn(async () => "OK"),
+    del: vi.fn(async (...keys: string[]) => {
+      let removed = 0;
+      for (const k of keys) {
+        if (lists.delete(k)) removed++;
+        if (strings.delete(k)) removed++;
+      }
+      return removed;
+    }),
+    set: vi.fn(async (key: string, value: string) => {
+      strings.set(key, value);
+      return "OK";
+    }),
+    get: vi.fn(async (key: string) => strings.get(key) ?? null),
+    expire: vi.fn(async () => 1),
+    llen: vi.fn(async (key: string) => (lists.get(key) ?? []).length),
+    eval: vi.fn(
+      async (_script: string, _numKeys: number, ...args: string[]) => {
+        const [dedupKey, listKey, , ...serialised] = args;
+        if (!dedupKey || !listKey) return 0;
+        if (strings.has(dedupKey)) return 0;
+        const list = lists.get(listKey) ?? [];
+        list.push(...serialised);
+        lists.set(listKey, list);
+        strings.set(dedupKey, "1");
+        return 1;
+      }
+    ),
+  };
+}
+
+const humanMsg = new HumanMessage({ id: "msg-1", content: "Hello" }).toDict();
+const aiMsg = new AIMessage({ id: "msg-2", content: "Hi there!" }).toDict();
+const humanMsg2 = new HumanMessage({
+  id: "msg-3",
+  content: "Again please",
+}).toDict();
+
+async function seed(
+  redis: ReturnType<typeof createStatefulRedis>,
+  threadId: string,
+  messages: StoredMessage[]
+): Promise<void> {
+  const tm = createLangChainThreadManager({
+    redis: redis as never,
+    threadId,
+  });
+  await tm.initialize();
+  await tm.append(messages);
+}
+
+describe("LangChain fork + transform hooks", () => {
+  it("falls back to plain fork when no hooks are set", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [humanMsg, aiMsg]);
+
+    const tm = createLangChainThreadManager({
+      redis: redis as never,
+      threadId: "src",
+    });
+    const forked = await tm.fork("dst");
+    expect(await forked.load()).toEqual([humanMsg, aiMsg]);
+  });
+
+  it("applies onForkPrepareThread then onForkTransform in order", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [humanMsg, aiMsg, humanMsg2]);
+
+    const order: string[] = [];
+    const tm = createLangChainThreadManager({
+      redis: redis as never,
+      threadId: "src",
+      hooks: {
+        onForkPrepareThread: async (messages) => {
+          order.push("prepare");
+          return messages.slice(0, -1);
+        },
+        onForkTransform: (msg, i) => {
+          order.push("transform");
+          return { ...msg, data: { ...msg.data, content: `[x${i}]` } };
+        },
+      },
+    });
+
+    const forked = await tm.fork("dst");
+    const loaded = await forked.load();
+
+    expect(order).toEqual(["prepare", "transform", "transform"]);
+    expect(loaded).toHaveLength(2);
+    expect(loaded[0]?.data.content).toBe("[x0]");
+    expect(loaded[1]?.data.content).toBe("[x1]");
+  });
+
+  it("keeps the source thread untouched", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [humanMsg, aiMsg]);
+
+    const tm = createLangChainThreadManager({
+      redis: redis as never,
+      threadId: "src",
+      hooks: {
+        onForkTransform: (msg) => ({
+          ...msg,
+          data: { ...msg.data, content: "mutated" },
+        }),
+      },
+    });
+
+    await tm.fork("dst");
+    expect(await tm.load()).toEqual([humanMsg, aiMsg]);
+  });
+});

package/src/adapters/thread/langchain/index.ts

@@ -15,6 +15,9 @@
  * ```
  */
 
+// Adapter identity (wire format — matches Temporal activity-name prefix)
+export { ADAPTER_ID, type AdapterId } from "./adapter-id";
+
 // Adapter (primary API)
 export {
   createLangChainAdapter,

package/src/adapters/thread/langchain/model-invoker.ts

@@ -47,7 +47,8 @@ export function createLangChainModelInvoker<
   return async function invokeLangChainModel(
     config: ModelInvokerConfig
   ): Promise<AgentResponse<StoredMessage>> {
-    const { threadId, threadKey, agentName, state, metadata } = config;
+    const { threadId, threadKey, agentName, state, metadata, assistantMessageId } =
+      config;
     const { heartbeat, signal } = getActivityContext();
 
     const thread = createLangChainThreadManager({
@@ -58,6 +59,11 @@ export function createLangChainModelInvoker<
     });
     const runId = uuidv4();
 
+    // Truncate the thread starting at the id the assistant message
+    // will be stored under. No-op on the first attempt; on rewind
+    // retry / Temporal reset it wipes the prior attempt's assistant
+    // + tool results so the LLM sees the original pre-call state.
+    await thread.truncateFromId(assistantMessageId);
     const { messages } = await thread.prepareForInvocation();
 
     const heartbeatInterval = heartbeat

package/src/adapters/thread/langchain/proxy.ts

@@ -22,15 +22,16 @@ import { type ActivityInterfaceFor } from "@temporalio/workflow";
 import type { ThreadOps } from "../../../lib/session/types";
 import type { LangChainContent } from "./thread-manager";
 import { createThreadOpsProxy } from "../../../lib/thread/proxy";
+import { ADAPTER_ID } from "./adapter-id";
 
-const ADAPTER_PREFIX = "langChain";
+export { ADAPTER_ID, type AdapterId } from "./adapter-id";
 
 export function proxyLangChainThreadOps(
   scope?: string,
   options?: Parameters<typeof createThreadOpsProxy>[2]
 ): ActivityInterfaceFor<ThreadOps<LangChainContent>> {
   return createThreadOpsProxy(
-    ADAPTER_PREFIX,
+    ADAPTER_ID,
     scope,
     options
   ) as ActivityInterfaceFor<ThreadOps<LangChainContent>>;

package/src/adapters/thread/langchain/thread-manager.ts

@@ -143,5 +143,31 @@ export function createLangChainThreadManager(
     },
   };
 
-  return Object.assign(base, helpers);
+  const manager = Object.assign(base, helpers);
+
+  const originalFork = manager.fork.bind(manager);
+  manager.fork = async (
+    newThreadId: string
+  ): Promise<LangChainThreadManager> => {
+    await originalFork(newThreadId);
+    const forked = createLangChainThreadManager({
+      ...config,
+      threadId: newThreadId,
+    });
+    const { onForkPrepareThread, onForkTransform } = config.hooks ?? {};
+    if (!onForkPrepareThread && !onForkTransform) {
+      return forked;
+    }
+    let next = await forked.load();
+    if (onForkPrepareThread) {
+      next = await onForkPrepareThread(next);
+    }
+    if (onForkTransform) {
+      next = next.map((msg, i) => onForkTransform(msg, i, next));
+    }
+    await forked.replaceAll(next);
+    return forked;
+  };
+
+  return manager;
 }

package/src/lib/lifecycle.ts

@@ -8,7 +8,9 @@
  * - `"new"` — start a fresh thread (optionally specify its ID).
  * - `"continue"` — append directly to an existing thread in-place.
  * - `"fork"` — copy all messages from an existing thread into a new one and
- *   continue there.
+ *   continue there. When the adapter has `onForkPrepareThread` and/or
+ *   `onForkTransform` hooks configured, they are applied once to the forked
+ *   thread before the session starts.
  */
 export type ThreadInit =
   | { mode: "new"; threadId?: string }
@@ -19,7 +21,7 @@ export type ThreadInit =
 // Sandbox lifecycle
 // ============================================================================
 
-import type { SandboxSnapshot } from "./sandbox/types";
+import type { SandboxCreateOptions, SandboxSnapshot } from "./sandbox/types";
 
 /**
  * Sandbox initialization strategy.
@@ -31,17 +33,24 @@ import type { SandboxSnapshot } from "./sandbox/types";
  *   Paused sandboxes are automatically resumed. The shutdown policy applies
  *   on exit.
  * - `"fork"` — fork from an existing (or paused) sandbox; a new sandbox is
- *   created and owned by this session.
+ *   created and owned by this session. `options` is an optional per-call
+ *   override merged on top of the provider's static defaults.
  * - `"from-snapshot"` — restore a fresh sandbox from a previously captured
  *   {@link SandboxSnapshot}. The new sandbox is owned by this session.
+ *   `options` is an optional per-call override merged on top of the
+ *   provider's static defaults.
  * - `"inherit"` — use a sandbox owned by someone else (e.g. a parent agent).
  *   The session will **not** manage its lifecycle on exit.
  */
 export type SandboxInit =
   | { mode: "new"; ctx?: unknown }
   | { mode: "continue"; sandboxId: string }
-  | { mode: "fork"; sandboxId: string }
-  | { mode: "from-snapshot"; snapshot: SandboxSnapshot }
+  | { mode: "fork"; sandboxId: string; options?: SandboxCreateOptions }
+  | {
+      mode: "from-snapshot";
+      snapshot: SandboxSnapshot;
+      options?: SandboxCreateOptions;
+    }
   | {
       mode: "inherit";
       sandboxId: string;

package/src/lib/model/types.ts

@@ -29,6 +29,13 @@ export interface ModelInvokerConfig {
   agentName: string;
   state: BaseAgentState;
   metadata?: Record<string, unknown>;
+  /**
+   * The id the assistant message produced by this call will be stored
+   * under. Invokers truncate the thread from this id on entry so that
+   * rewind retries and Temporal workflow resets restore the pre-call
+   * state before re-invoking the LLM. See {@link RunAgentConfig}.
+   */
+  assistantMessageId: string;
 }
 
 /**

package/src/lib/sandbox/manager.ts

@@ -120,9 +120,7 @@ export class SandboxManager<
   async create(
     options?: TOptions,
     ctx?: TCtx
-  ): Promise<{
-    sandboxId: string;
-  } | null> {
+  ): Promise<{ sandboxId: string } | null> {
     let providerOptions = options;
 
     if (this.hooks.onPreCreate) {
@@ -175,12 +173,15 @@ export class SandboxManager<
     await this.provider.resume(id);
   }
 
-  async snapshot(id: string): Promise<SandboxSnapshot> {
-    return this.provider.snapshot(id);
+  async snapshot(id: string, options?: TOptions): Promise<SandboxSnapshot> {
+    return this.provider.snapshot(id, options);
   }
 
-  async restore(snapshot: SandboxSnapshot): Promise<string> {
-    const sandbox = await this.provider.restore(snapshot);
+  async restore(
+    snapshot: SandboxSnapshot,
+    options?: TOptions
+  ): Promise<string> {
+    const sandbox = await this.provider.restore(snapshot, options);
     return sandbox.id;
   }
 
@@ -188,8 +189,8 @@ export class SandboxManager<
     await this.provider.deleteSnapshot(snapshot);
   }
 
-  async fork(sandboxId: string): Promise<string> {
-    const sandbox = await this.provider.fork(sandboxId);
+  async fork(sandboxId: string, options?: TOptions): Promise<string> {
+    const sandbox = await this.provider.fork(sandboxId, options);
     return sandbox.id;
   }
 
@@ -221,9 +222,7 @@ export class SandboxManager<
       createSandbox: async (
         options?: TOptions,
         ctx?: TCtx
-      ): Promise<{
-        sandboxId: string;
-      } | null> => {
+      ): Promise<{ sandboxId: string } | null> => {
         return this.create(options, ctx);
       },
       destroySandbox: async (sandboxId: string): Promise<void> => {
@@ -238,19 +237,28 @@ export class SandboxManager<
       resumeSandbox: async (sandboxId: string): Promise<void> => {
         await this.resume(sandboxId);
       },
-      snapshotSandbox: async (sandboxId: string): Promise<SandboxSnapshot> => {
-        return this.snapshot(sandboxId);
+      snapshotSandbox: async (
+        sandboxId: string,
+        options?: TOptions
+      ): Promise<SandboxSnapshot> => {
+        return this.snapshot(sandboxId, options);
       },
-      restoreSandbox: async (snapshot: SandboxSnapshot): Promise<string> => {
-        return this.restore(snapshot);
+      restoreSandbox: async (
+        snapshot: SandboxSnapshot,
+        options?: TOptions
+      ): Promise<string> => {
+        return this.restore(snapshot, options);
       },
       deleteSandboxSnapshot: async (
         snapshot: SandboxSnapshot
       ): Promise<void> => {
         await this.deleteSnapshot(snapshot);
       },
-      forkSandbox: async (sandboxId: string): Promise<string> => {
-        return this.fork(sandboxId);
+      forkSandbox: async (
+        sandboxId: string,
+        options?: TOptions
+      ): Promise<string> => {
+        return this.fork(sandboxId, options);
       },
     };
     const cap = (s: string): string => s.charAt(0).toUpperCase() + s.slice(1);

package/src/lib/sandbox/types.ts

@@ -158,11 +158,23 @@ export interface SandboxProvider<
   pause(sandboxId: string, ttlSeconds?: number): Promise<void>;
   /** Resume a paused sandbox. No-op if already running. */
   resume(sandboxId: string): Promise<void>;
-  snapshot(sandboxId: string): Promise<SandboxSnapshot>;
-  restore(snapshot: SandboxSnapshot): Promise<Sandbox>;
+  /**
+   * Capture a snapshot of a running sandbox. `options` is a per-call override
+   * merged on top of the provider's static defaults.
+   */
+  snapshot(sandboxId: string, options?: TOptions): Promise<SandboxSnapshot>;
+  /**
+   * Restore a sandbox from a snapshot. `options` is a per-call override
+   * merged on top of the provider's static defaults.
+   */
+  restore(snapshot: SandboxSnapshot, options?: TOptions): Promise<Sandbox>;
   /** Delete a previously captured snapshot. No-op if already deleted. */
   deleteSnapshot(snapshot: SandboxSnapshot): Promise<void>;
-  fork(sandboxId: string): Promise<Sandbox>;
+  /**
+   * Fork a running sandbox into a new one. `options` is a per-call override
+   * merged on top of the provider's static defaults.
+   */
+  fork(sandboxId: string, options?: TOptions): Promise<Sandbox>;
 }
 
 // ============================================================================
@@ -181,12 +193,20 @@ export interface SandboxOps<
   pauseSandbox(sandboxId: string): Promise<void>;
   /** Resume a paused sandbox. No-op if already running. */
   resumeSandbox(sandboxId: string): Promise<void>;
-  snapshotSandbox(sandboxId: string): Promise<SandboxSnapshot>;
-  /** Create a fresh sandbox from a previously captured snapshot. */
-  restoreSandbox(snapshot: SandboxSnapshot): Promise<string>;
+  /** Capture a snapshot. `options` is a per-call override merged on top of provider defaults. */
+  snapshotSandbox(
+    sandboxId: string,
+    options?: TOptions
+  ): Promise<SandboxSnapshot>;
+  /** Create a fresh sandbox from a snapshot. `options` is a per-call override merged on top of provider defaults. */
+  restoreSandbox(
+    snapshot: SandboxSnapshot,
+    options?: TOptions
+  ): Promise<string>;
   /** Delete a previously captured snapshot. No-op if already deleted. */
   deleteSandboxSnapshot(snapshot: SandboxSnapshot): Promise<void>;
-  forkSandbox(sandboxId: string): Promise<string>;
+  /** Fork a running sandbox. `options` is a per-call override merged on top of provider defaults. */
+  forkSandbox(sandboxId: string, options?: TOptions): Promise<string>;
 }
 
 /**