zeitlich 0.2.38 → 0.2.40

package/src/adapters/thread/google-genai/fork-transform.test.ts

@@ -0,0 +1,149 @@
+import { describe, expect, it, vi } from "vitest";
+import type { StoredContent } from "./thread-manager";
+import { createGoogleGenAIThreadManager } 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 userContent: StoredContent = {
+  id: "msg-1",
+  content: { role: "user", parts: [{ text: "Hello" }] },
+};
+
+const modelContent: StoredContent = {
+  id: "msg-2",
+  content: { role: "model", parts: [{ text: "Hi there!" }] },
+};
+
+const userContent2: StoredContent = {
+  id: "msg-3",
+  content: { role: "user", parts: [{ text: "Again please" }] },
+};
+
+async function seed(
+  redis: ReturnType<typeof createStatefulRedis>,
+  threadId: string,
+  messages: StoredContent[]
+): Promise<void> {
+  const tm = createGoogleGenAIThreadManager({
+    redis: redis as never,
+    threadId,
+  });
+  await tm.initialize();
+  await tm.append(messages);
+}
+
+describe("Google GenAI fork + transform hooks", () => {
+  it("falls back to plain fork when no hooks are set", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [userContent, modelContent]);
+
+    const tm = createGoogleGenAIThreadManager({
+      redis: redis as never,
+      threadId: "src",
+    });
+    const forked = await tm.fork("dst");
+    expect(await forked.load()).toEqual([userContent, modelContent]);
+  });
+
+  it("applies onForkPrepareThread then onForkTransform in order", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [userContent, modelContent, userContent2]);
+
+    const order: string[] = [];
+    const tm = createGoogleGenAIThreadManager({
+      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,
+            content: { ...msg.content, parts: [{ text: `[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]?.content.parts).toEqual([{ text: "[x0]" }]);
+    expect(loaded[1]?.content.parts).toEqual([{ text: "[x1]" }]);
+  });
+
+  it("keeps the source thread untouched", async () => {
+    const redis = createStatefulRedis();
+    await seed(redis, "src", [userContent, modelContent]);
+
+    const tm = createGoogleGenAIThreadManager({
+      redis: redis as never,
+      threadId: "src",
+      hooks: {
+        onForkTransform: (msg) => ({
+          ...msg,
+          content: { ...msg.content, parts: [{ text: "mutated" }] },
+        }),
+      },
+    });
+
+    await tm.fork("dst");
+    expect(await tm.load()).toEqual([userContent, modelContent]);
+  });
+});

package/src/adapters/thread/google-genai/index.ts

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

package/src/adapters/thread/google-genai/model-invoker.ts

@@ -64,7 +64,7 @@ export function createGoogleGenAIModelInvoker({
   return async function invokeGoogleGenAIModel(
     config: ModelInvokerConfig
   ): Promise<AgentResponse<Content>> {
-    const { threadId, threadKey, state } = config;
+    const { threadId, threadKey, state, assistantMessageId } = config;
     const { heartbeat, signal } = getActivityContext();
 
     const thread = createGoogleGenAIThreadManager({
@@ -73,7 +73,12 @@ export function createGoogleGenAIModelInvoker({
       key: threadKey,
       hooks,
     });
-    const { contents, systemInstruction, storedLength } =
+    // 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 { contents, systemInstruction } =
       await thread.prepareForInvocation();
 
     const functionDeclarations = toFunctionDeclarations(state.tools);
@@ -117,7 +122,6 @@ export function createGoogleGenAIModelInvoker({
         outputTokens: lastChunk.usageMetadata?.candidatesTokenCount,
         cachedReadTokens: lastChunk.usageMetadata?.cachedContentTokenCount,
       },
-      threadLengthAtCall: storedLength,
     };
   };
 }

package/src/adapters/thread/google-genai/proxy.ts

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

package/src/adapters/thread/google-genai/thread-manager.ts

@@ -37,8 +37,6 @@ export interface GoogleGenAIThreadManagerConfig {
 export interface GoogleGenAIInvocationPayload {
   contents: Content[];
   systemInstruction?: Part[];
-  /** Number of stored messages loaded from Redis before preparation. */
-  storedLength: number;
 }
 
 /** Thread manager with Google GenAI Content convenience helpers */
@@ -198,10 +196,35 @@ export function createGoogleGenAIThreadManager(
         ...(systemInstruction && systemInstruction.length > 0
           ? { systemInstruction }
           : {}),
-        storedLength: stored.length,
       };
     },
   };
 
-  return Object.assign(base, helpers);
+  const manager = Object.assign(base, helpers);
+
+  const originalFork = manager.fork.bind(manager);
+  manager.fork = async (
+    newThreadId: string
+  ): Promise<GoogleGenAIThreadManager> => {
+    await originalFork(newThreadId);
+    const forked = createGoogleGenAIThreadManager({
+      ...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/adapters/thread/index.ts

@@ -0,0 +1,39 @@
+/**
+ * Barrel re-exports for every built-in thread adapter's public identity.
+ *
+ * Downstream consumers reading persisted threads can import a narrow
+ * discriminated union of adapter identifiers without pulling the full
+ * adapter implementation (Redis, provider SDKs, etc.) as a dependency —
+ * each individual re-export resolves to an `adapter-id.ts` module with
+ * no runtime dependencies.
+ *
+ * @example
+ * ```typescript
+ * import {
+ *   LANGCHAIN_ADAPTER_ID,
+ *   GOOGLE_GENAI_ADAPTER_ID,
+ *   ANTHROPIC_ADAPTER_ID,
+ *   type ThreadAdapterId,
+ * } from 'zeitlich/adapters/thread';
+ *
+ * interface ThreadIdentity {
+ *   adapter: ThreadAdapterId;
+ *   threadKey: string;
+ *   threadId: string;
+ * }
+ * ```
+ */
+
+export { ADAPTER_ID as LANGCHAIN_ADAPTER_ID } from "./langchain/adapter-id";
+export { ADAPTER_ID as GOOGLE_GENAI_ADAPTER_ID } from "./google-genai/adapter-id";
+export { ADAPTER_ID as ANTHROPIC_ADAPTER_ID } from "./anthropic/adapter-id";
+
+import type { ADAPTER_ID as LANGCHAIN } from "./langchain/adapter-id";
+import type { ADAPTER_ID as GOOGLE_GENAI } from "./google-genai/adapter-id";
+import type { ADAPTER_ID as ANTHROPIC } from "./anthropic/adapter-id";
+
+/** Narrow discriminated union of every built-in thread adapter id. */
+export type ThreadAdapterId =
+  | typeof LANGCHAIN
+  | typeof GOOGLE_GENAI
+  | typeof ANTHROPIC;

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

@@ -1,5 +1,6 @@
 import type Redis from "ioredis";
 import type { ToolResultConfig } from "../../../lib/types";
+import type { PersistedThreadState } from "../../../lib/state/types";
 import type { MessageContent } from "@langchain/core/messages";
 import type {
   ActivityToolHandler,
@@ -21,11 +22,10 @@ import {
   type LangChainThreadManagerHooks,
 } from "./thread-manager";
 import { createLangChainModelInvoker } from "./model-invoker";
-
-const ADAPTER_PREFIX = "langChain" as const;
+import { ADAPTER_ID } from "./adapter-id";
 
 export type LangChainThreadOps<TScope extends string = ""> = PrefixedThreadOps<
-  ScopedPrefix<TScope, typeof ADAPTER_PREFIX>,
+  ScopedPrefix<TScope, typeof ADAPTER_ID>,
   LangChainContent
 >;
 
@@ -192,17 +192,43 @@ export function createLangChainAdapter(
         redis,
         threadId: sourceThreadId,
         key: threadKey,
+        hooks: config.hooks,
       });
       await thread.fork(targetThreadId);
     },
 
     async truncateThread(
       threadId: string,
-      length: number,
+      messageId: string,
       threadKey?: string,
     ): Promise<void> {
       const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
-      await thread.truncate(length);
+      await thread.truncateFromId(messageId);
+    },
+
+    async loadThreadState(
+      threadId: string,
+      threadKey?: string
+    ): Promise<PersistedThreadState | null> {
+      const thread = createLangChainThreadManager({
+        redis,
+        threadId,
+        key: threadKey,
+      });
+      return thread.loadState();
+    },
+
+    async saveThreadState(
+      threadId: string,
+      state: PersistedThreadState,
+      threadKey?: string
+    ): Promise<void> {
+      const thread = createLangChainThreadManager({
+        redis,
+        threadId,
+        key: threadKey,
+      });
+      await thread.saveState(state);
     },
   };
 
@@ -210,8 +236,8 @@ export function createLangChainAdapter(
     scope?: S
   ): LangChainThreadOps<S> {
     const prefix = scope
-      ? `${ADAPTER_PREFIX}${scope.charAt(0).toUpperCase()}${scope.slice(1)}`
-      : ADAPTER_PREFIX;
+      ? `${ADAPTER_ID}${scope.charAt(0).toUpperCase()}${scope.slice(1)}`
+      : ADAPTER_ID;
     const cap = (s: string): string => s.charAt(0).toUpperCase() + s.slice(1);
     return Object.fromEntries(
       Object.entries(threadOps).map(([k, v]) => [`${prefix}${cap(k)}`, v])

package/src/adapters/thread/langchain/adapter-id.ts

@@ -0,0 +1,16 @@
+/**
+ * Public adapter identity for the LangChain thread adapter.
+ *
+ * This value is wire format — it appears as the prefix for Temporal
+ * activity names (e.g. `langChainCodingAgentInitializeThread`) and must
+ * never change, since renaming it would orphan existing persisted
+ * threads and break in-flight workflows.
+ *
+ * Re-exported from `zeitlich/adapters/thread/langchain` so downstream
+ * consumers can use the exact same literal the adapter uses internally,
+ * typed as the narrow string literal `"langChain"`.
+ */
+export const ADAPTER_ID = "langChain" as const;
+
+/** Narrow string-literal type for {@link ADAPTER_ID}. */
+export type AdapterId = typeof ADAPTER_ID;

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,7 +59,12 @@ export function createLangChainModelInvoker<
     });
     const runId = uuidv4();
 
-    const { messages, storedLength } = await thread.prepareForInvocation();
+    // 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
       ? setInterval(() => heartbeat(), 30_000)
@@ -97,7 +103,6 @@ export function createLangChainModelInvoker<
             response.usage_metadata?.input_token_details?.cache_read ||
             (providerUsage.cacheReadInputTokens as number | undefined),
         },
-        threadLengthAtCall: storedLength,
       };
     } finally {
       if (heartbeatInterval) clearInterval(heartbeatInterval);

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

@@ -39,8 +39,6 @@ export interface LangChainThreadManagerConfig {
 /** Prepared payload ready to send to a LangChain chat model */
 export interface LangChainInvocationPayload {
   messages: BaseMessage[];
-  /** Number of stored messages loaded from Redis before preparation. */
-  storedLength: number;
 }
 
 /** Thread manager with LangChain StoredMessage convenience helpers */
@@ -141,10 +139,35 @@ export function createLangChainThreadManager(
         messages: onPreparedMessage
           ? messages.map((msg, i) => onPreparedMessage(msg, i, messages))
           : messages,
-        storedLength: stored.length,
       };
     },
   };
 
-  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 }

package/src/lib/model/types.ts

@@ -8,16 +8,6 @@ export interface AgentResponse<M = unknown> {
   message: M;
   rawToolCalls: RawToolCall[];
   usage?: TokenUsage;
-  /**
-   * Number of stored messages in the thread at the moment the LLM was
-   * invoked — i.e. *before* the assistant message is appended. The
-   * session uses this as a rewind snapshot so it can roll the thread
-   * back to this exact state if a tool requests a rewind.
-   *
-   * Adapters compute this for free from the array of stored messages
-   * they load when preparing the payload.
-   */
-  threadLengthAtCall?: number;
 }
 
 /**
@@ -39,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;
 }
 
 /**