zeitlich 0.2.25 → 0.2.27

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

@@ -0,0 +1,159 @@
+import { describe, expect, it, vi } from "vitest";
+import type { Content } from "@google/genai";
+import type { StoredContent } from "./thread-manager";
+import { createGoogleGenAIThreadManager } from "./thread-manager";
+
+function createMockRedis(stored: StoredContent[]) {
+  return {
+    exists: vi.fn().mockResolvedValue(1),
+    lrange: vi.fn().mockResolvedValue(stored.map((m) => JSON.stringify(m))),
+    del: vi.fn().mockResolvedValue(1),
+    set: vi.fn().mockResolvedValue("OK"),
+    rpush: vi.fn().mockResolvedValue(1),
+    expire: vi.fn().mockResolvedValue(1),
+    eval: vi.fn().mockResolvedValue(1),
+  };
+}
+
+const systemContent: StoredContent = {
+  id: "sys-1",
+  content: { role: "system", parts: [{ text: "You are helpful." }] },
+};
+
+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!" }] },
+};
+
+describe("Google GenAI thread manager hooks", () => {
+  describe("onPrepareMessage", () => {
+    it("transforms stored messages before system extraction and merge", async () => {
+      const hook = vi.fn((msg: StoredContent) => {
+        if (msg.content.role === "system") return msg;
+        return {
+          ...msg,
+          content: {
+            ...msg.content,
+            parts: [{ text: `[modified] ${msg.content.parts?.[0]?.text ?? ""}` }],
+          },
+        };
+      });
+
+      const redis = createMockRedis([systemContent, userContent, modelContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+        hooks: { onPrepareMessage: hook },
+      });
+
+      const { contents, systemInstruction } = await tm.prepareForInvocation();
+
+      expect(hook).toHaveBeenCalledTimes(3);
+      expect(hook).toHaveBeenCalledWith(systemContent, 0, [systemContent, userContent, modelContent]);
+      expect(systemInstruction).toBe("You are helpful.");
+      expect(contents[0]?.parts?.[0]?.text).toBe("[modified] Hello");
+      expect(contents[1]?.parts?.[0]?.text).toBe("[modified] Hi there!");
+    });
+
+    it("is not called when not configured", async () => {
+      const redis = createMockRedis([userContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+      });
+
+      const { contents } = await tm.prepareForInvocation();
+      expect(contents).toHaveLength(1);
+      expect(contents[0]?.parts?.[0]?.text).toBe("Hello");
+    });
+  });
+
+  describe("onPreparedMessage", () => {
+    it("transforms SDK-native Content after merge", async () => {
+      const hook = vi.fn((msg: Content) => ({
+        ...msg,
+        parts: [{ text: `[post] ${msg.parts?.[0]?.text ?? ""}` }],
+      }));
+
+      const redis = createMockRedis([userContent, modelContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+        hooks: { onPreparedMessage: hook },
+      });
+
+      const { contents } = await tm.prepareForInvocation();
+
+      expect(hook).toHaveBeenCalledTimes(2);
+      expect(contents[0]?.parts?.[0]?.text).toBe("[post] Hello");
+      expect(contents[1]?.parts?.[0]?.text).toBe("[post] Hi there!");
+    });
+
+    it("receives the full prepared contents array", async () => {
+      const hook = vi.fn((msg: Content) => msg);
+
+      const redis = createMockRedis([userContent, modelContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+        hooks: { onPreparedMessage: hook },
+      });
+
+      await tm.prepareForInvocation();
+
+      const args = hook.mock.calls[0] as unknown as [Content, number, Content[]];
+      expect(args[2]).toHaveLength(2);
+    });
+  });
+
+  describe("both hooks combined", () => {
+    it("runs onPrepareMessage before onPreparedMessage", async () => {
+      const order: string[] = [];
+
+      const redis = createMockRedis([userContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+        hooks: {
+          onPrepareMessage: (msg) => {
+            order.push("pre");
+            return msg;
+          },
+          onPreparedMessage: (msg) => {
+            order.push("post");
+            return msg;
+          },
+        },
+      });
+
+      await tm.prepareForInvocation();
+      expect(order).toEqual(["pre", "post"]);
+    });
+
+    it("onPreparedMessage sees results of onPrepareMessage", async () => {
+      const redis = createMockRedis([userContent]);
+      const tm = createGoogleGenAIThreadManager({
+        redis: redis as never,
+        threadId: "t1",
+        hooks: {
+          onPrepareMessage: (msg) => ({
+            ...msg,
+            content: { ...msg.content, parts: [{ text: "replaced" }] },
+          }),
+          onPreparedMessage: (msg) => {
+            expect(msg.parts?.[0]?.text).toBe("replaced");
+            return msg;
+          },
+        },
+      });
+
+      const { contents } = await tm.prepareForInvocation();
+      expect(contents[0]?.parts?.[0]?.text).toBe("replaced");
+    });
+  });
+});

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

@@ -2,10 +2,14 @@ import type Redis from "ioredis";
 import type { Content, Part } from "@google/genai";
 import {
   createThreadManager,
-  type BaseThreadManager,
+  type ProviderThreadManager,
   type ThreadManagerConfig,
+  type ThreadManagerHooks,
 } from "../../../lib/thread";
-import type { MessageContent, ToolMessageContent } from "../../../lib/types";
+import type { GoogleGenAIToolResponse } from "./activities";
+
+/** SDK-native content type for Google GenAI human messages */
+export type GoogleGenAIContent = string | Part[];
 
 /** A Content with a unique ID for idempotent Redis storage */
 export interface StoredContent {
@@ -13,78 +17,65 @@ export interface StoredContent {
   content: Content;
 }
 
+export type GoogleGenAIThreadManagerHooks = ThreadManagerHooks<StoredContent, Content>;
+
 export interface GoogleGenAIThreadManagerConfig {
   redis: Redis;
   threadId: string;
   /** Thread key, defaults to 'messages' */
   key?: string;
+  hooks?: GoogleGenAIThreadManagerHooks;
+}
+
+/** Prepared payload ready to send to the Google GenAI API */
+export interface GoogleGenAIInvocationPayload {
+  contents: Content[];
+  systemInstruction?: string;
 }
 
 /** Thread manager with Google GenAI Content convenience helpers */
-export interface GoogleGenAIThreadManager extends BaseThreadManager<StoredContent> {
-  createUserContent(
-    id: string,
-    content: string | MessageContent
-  ): StoredContent;
-  createSystemContent(id: string, content: string): StoredContent;
-  createModelContent(id: string, parts: Part[]): StoredContent;
-  createToolResponseContent(
-    id: string,
-    toolCallId: string,
-    toolName: string,
-    content: ToolMessageContent
-  ): StoredContent;
-  appendUserMessage(
-    id: string,
-    content: string | MessageContent
-  ): Promise<void>;
-  appendSystemMessage(id: string, content: string): Promise<void>;
+export interface GoogleGenAIThreadManager
+  extends ProviderThreadManager<StoredContent, GoogleGenAIContent, GoogleGenAIToolResponse> {
   appendModelContent(id: string, parts: Part[]): Promise<void>;
-  appendToolResult(
-    id: string,
-    toolCallId: string,
-    toolName: string,
-    content: ToolMessageContent
-  ): Promise<void>;
+  prepareForInvocation(): Promise<GoogleGenAIInvocationPayload>;
 }
 
 function storedContentId(msg: StoredContent): string {
   return msg.id;
 }
 
-/** Convert zeitlich MessageContent to Google GenAI Part[] */
-export function messageContentToParts(
-  content: string | MessageContent
-): Part[] {
+/** Normalise content into Part[] */
+function toParts(content: GoogleGenAIContent): Part[] {
   if (typeof content === "string") {
     return [{ text: content }];
   }
-  if (Array.isArray(content)) {
-    return content.map((part) => {
-      if (part.type === "text") {
-        return { text: part.text as string };
-      }
-      return part as unknown as Part;
-    });
+  return content;
+}
+
+/** Convert a string or object into a Record suitable for functionResponse.response */
+function toFunctionResponse(content: string | Record<string, unknown>): Record<string, unknown> {
+  if (typeof content === "object") {
+    return content;
   }
-  return [{ text: String(content) }];
+  return { result: content };
 }
 
-/** Parse ToolMessageContent into a Record suitable for functionResponse */
-function parseToolResponse(
-  content: ToolMessageContent
-): Record<string, unknown> {
-  if (typeof content === "string") {
-    try {
-      const parsed: unknown = JSON.parse(content);
-      return typeof parsed === "object" && parsed !== null
-        ? (parsed as Record<string, unknown>)
-        : { result: content };
-    } catch {
-      return { result: content };
+/**
+ * Merge consecutive Content objects sharing the same role.
+ * The Gemini API requires alternating user/model turns; without
+ * merging, multiple sequential tool-result messages would violate this.
+ */
+function mergeConsecutiveContents(contents: Content[]): Content[] {
+  const merged: Content[] = [];
+  for (const content of contents) {
+    const last = merged[merged.length - 1];
+    if (last && last.role === content.role) {
+      last.parts = [...(last.parts ?? []), ...(content.parts ?? [])];
+    } else {
+      merged.push({ ...content, parts: [...(content.parts ?? [])] });
     }
   }
-  return { result: content };
+  return merged;
 }
 
 /**
@@ -93,7 +84,7 @@ function parseToolResponse(
  * appending typed Content messages.
  */
 export function createGoogleGenAIThreadManager(
-  config: GoogleGenAIThreadManagerConfig
+  config: GoogleGenAIThreadManagerConfig,
 ): GoogleGenAIThreadManager {
   const baseConfig: ThreadManagerConfig<StoredContent> = {
     redis: config.redis,
@@ -104,79 +95,79 @@ export function createGoogleGenAIThreadManager(
 
   const base = createThreadManager(baseConfig);
 
-  const helpers = {
-    createUserContent(
+  const helpers: Omit<GoogleGenAIThreadManager, keyof typeof base> = {
+    async appendUserMessage(
       id: string,
-      content: string | MessageContent
-    ): StoredContent {
-      return {
+      content: GoogleGenAIContent,
+    ): Promise<void> {
+      await base.append([{
         id,
-        content: { role: "user", parts: messageContentToParts(content) },
-      };
+        content: { role: "user", parts: toParts(content) },
+      }]);
     },
 
-    createSystemContent(id: string, content: string): StoredContent {
-      return {
+    async appendSystemMessage(id: string, content: string): Promise<void> {
+      await base.initialize();
+      await base.append([{
         id,
         content: { role: "system", parts: [{ text: content }] },
-      };
+      }]);
     },
 
-    createModelContent(id: string, parts: Part[]): StoredContent {
-      return {
+    async appendModelContent(id: string, parts: Part[]): Promise<void> {
+      await base.append([{
         id,
         content: { role: "model", parts },
-      };
+      }]);
     },
 
-    createToolResponseContent(
+    async appendToolResult(
       id: string,
       toolCallId: string,
       toolName: string,
-      content: ToolMessageContent
-    ): StoredContent {
-      return {
-        id,
-        content: {
-          role: "user",
-          parts: [
-            {
-              functionResponse: {
-                id: toolCallId,
-                name: toolName,
-                response: parseToolResponse(content),
-              },
-            },
-          ],
-        },
-      };
-    },
-
-    async appendUserMessage(
-      id: string,
-      content: string | MessageContent
+      content: GoogleGenAIToolResponse,
     ): Promise<void> {
-      await base.append([helpers.createUserContent(id, content)]);
-    },
+      const parts: Part[] = Array.isArray(content)
+        ? content as Part[]
+        : [{
+            functionResponse: {
+              id: toolCallId,
+              name: toolName,
+              response: toFunctionResponse(content),
+            },
+          }];
 
-    async appendSystemMessage(id: string, content: string): Promise<void> {
-      await base.initialize();
-      await base.append([helpers.createSystemContent(id, content)]);
+      await base.append([{
+        id,
+        content: { role: "user", parts },
+      }]);
     },
 
-    async appendModelContent(id: string, parts: Part[]): Promise<void> {
-      await base.append([helpers.createModelContent(id, parts)]);
-    },
+    async prepareForInvocation(): Promise<GoogleGenAIInvocationPayload> {
+      const stored = await base.load();
+      const { onPrepareMessage, onPreparedMessage } = config.hooks ?? {};
+      const mapped = onPrepareMessage
+        ? stored.map((msg, i) => onPrepareMessage(msg, i, stored))
+        : stored;
+
+      let systemInstruction: string | undefined;
+      const conversationContents: Content[] = [];
+
+      for (const item of mapped) {
+        if (item.content.role === "system") {
+          systemInstruction = item.content.parts?.[0]?.text;
+        } else {
+          conversationContents.push(item.content);
+        }
+      }
 
-    async appendToolResult(
-      id: string,
-      toolCallId: string,
-      toolName: string,
-      content: ToolMessageContent
-    ): Promise<void> {
-      await base.append([
-        helpers.createToolResponseContent(id, toolCallId, toolName, content),
-      ]);
+      const contents = mergeConsecutiveContents(conversationContents);
+      return {
+        contents: onPreparedMessage
+          ? contents.map((msg, i) => onPreparedMessage(msg, i, contents))
+          : contents,
+        ...(systemInstruction ? { systemInstruction } : {}),
+      };
     },
   };
 

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

@@ -1,6 +1,11 @@
 import type Redis from "ioredis";
 import type { ToolResultConfig } from "../../../lib/types";
 import type { MessageContent } from "@langchain/core/messages";
+import type {
+  ActivityToolHandler,
+  RouterContext,
+  ToolHandlerResponse,
+} from "../../../lib/tool-router/types";
 import type {
   ThreadOps,
   PrefixedThreadOps,
@@ -9,21 +14,35 @@ import type {
 import type { ModelInvoker } from "../../../lib/model";
 import type { StoredMessage } from "@langchain/core/messages";
 import type { BaseChatModel } from "@langchain/core/language_models/chat_models";
-import { createLangChainThreadManager } from "./thread-manager";
+import {
+  createLangChainThreadManager,
+  type LangChainContent,
+  type LangChainThreadManagerHooks,
+} from "./thread-manager";
 import { createLangChainModelInvoker } from "./model-invoker";
 
 const ADAPTER_PREFIX = "langChain" as const;
 
 export type LangChainThreadOps<TScope extends string = ""> =
-  PrefixedThreadOps<ScopedPrefix<TScope, typeof ADAPTER_PREFIX>>;
+  PrefixedThreadOps<ScopedPrefix<TScope, typeof ADAPTER_PREFIX>, LangChainContent>;
 
 export interface LangChainAdapterConfig {
   redis: Redis;
   /** Optional default model — if omitted, use `createModelInvoker()` to create invokers later */
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   model?: BaseChatModel<any>;
+  hooks?: LangChainThreadManagerHooks;
 }
 
+/**
+ * Tool response type accepted by the LangChain adapter.
+ *
+ * Content is passed directly to `ToolMessage` as `MessageContent`.
+ * Handlers can return a string or an array of content blocks
+ * (e.g. `{ type: "text", text: "..." }`, `{ type: "image_url", image_url: { ... } }`).
+ */
+export type LangChainToolResponse = MessageContent;
+
 export interface LangChainAdapter {
   /** Model invoker using the default model (only available when `model` was provided) */
   invoker: ModelInvoker<StoredMessage>;
@@ -42,8 +61,19 @@ export interface LangChainAdapter {
    * ```
    */
   createActivities<S extends string = "">(
-    scope?: S
+    scope?: S,
   ): LangChainThreadOps<S>;
+
+  /**
+   * Identity wrapper that types a tool handler for this adapter.
+   * Constrains `toolResponse` to {@link LangChainToolResponse}.
+   */
+  wrapHandler<TArgs, TResult, TContext extends RouterContext = RouterContext>(
+    handler: (
+      args: TArgs,
+      context: TContext,
+    ) => Promise<ToolHandlerResponse<TResult, LangChainToolResponse>>,
+  ): ActivityToolHandler<TArgs, TResult, TContext, LangChainToolResponse>;
 }
 
 /**
@@ -82,54 +112,58 @@ export interface LangChainAdapter {
  * ```
  */
 export function createLangChainAdapter(
-  config: LangChainAdapterConfig
+  config: LangChainAdapterConfig,
 ): LangChainAdapter {
   const { redis } = config;
 
-  const threadOps: ThreadOps = {
-    async initializeThread(threadId: string): Promise<void> {
-      const thread = createLangChainThreadManager({ redis, threadId });
+  const threadOps: ThreadOps<LangChainContent> = {
+    async initializeThread(threadId: string, threadKey?: string): Promise<void> {
+      const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
       await thread.initialize();
     },
 
     async appendHumanMessage(
       threadId: string,
       id: string,
-      content: string | MessageContent
+      content: LangChainContent,
+      threadKey?: string,
     ): Promise<void> {
-      const thread = createLangChainThreadManager({ redis, threadId });
-      await thread.appendHumanMessage(id, content);
+      const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
+      await thread.appendUserMessage(id, content);
     },
 
     async appendSystemMessage(
       threadId: string,
       id: string,
-      content: string
+      content: string,
+      threadKey?: string,
     ): Promise<void> {
-      const thread = createLangChainThreadManager({ redis, threadId });
+      const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
       await thread.appendSystemMessage(id, content);
     },
 
     async appendToolResult(id: string, cfg: ToolResultConfig): Promise<void> {
-      const { threadId, toolCallId, content } = cfg;
-      const thread = createLangChainThreadManager({ redis, threadId });
-      await thread.appendToolMessage(id, content, toolCallId);
+      const { threadId, threadKey, toolCallId, content } = cfg;
+      const thread = createLangChainThreadManager({ redis, threadId, key: threadKey });
+      await thread.appendToolResult(id, toolCallId, "", content);
     },
 
     async forkThread(
       sourceThreadId: string,
-      targetThreadId: string
+      targetThreadId: string,
+      threadKey?: string,
     ): Promise<void> {
       const thread = createLangChainThreadManager({
         redis,
         threadId: sourceThreadId,
+        key: threadKey,
       });
       await thread.fork(targetThreadId);
     },
   };
 
   function createActivities<S extends string = "">(
-    scope?: S
+    scope?: S,
   ): LangChainThreadOps<S> {
     const prefix = scope
       ? `${ADAPTER_PREFIX}${scope.charAt(0).toUpperCase()}${scope.slice(1)}`
@@ -137,22 +171,22 @@ export function createLangChainAdapter(
     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])
+      Object.entries(threadOps).map(([k, v]) => [`${prefix}${cap(k)}`, v]),
     ) as LangChainThreadOps<S>;
   }
 
   const makeInvoker = (
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    model: BaseChatModel<any>
+    model: BaseChatModel<any>,
   ): ModelInvoker<StoredMessage> =>
-    createLangChainModelInvoker({ redis, model });
+    createLangChainModelInvoker({ redis, model, hooks: config.hooks });
 
   const invoker: ModelInvoker<StoredMessage> = config.model
     ? makeInvoker(config.model)
     : () => {
         throw new Error(
           "No default model provided to createLangChainAdapter. " +
-            "Either pass `model` in the config or use `createModelInvoker(model)` instead."
+            "Either pass `model` in the config or use `createModelInvoker(model)` instead.",
         );
       };
 
@@ -160,5 +194,6 @@ export function createLangChainAdapter(
     createActivities,
     invoker,
     createModelInvoker: makeInvoker,
+    wrapHandler: (handler) => handler,
   };
 }

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

@@ -0,0 +1,37 @@
+import type { BaseMessage, MessageContent } from "@langchain/core/messages";
+
+type ContentBlock = MessageContent extends (infer U)[] | string ? U : never;
+
+/**
+ * Creates an `onPreparedMessage` hook that appends a cache-point content
+ * block to the last message in the thread.
+ *
+ * Skips appending if the last message already contains a block with the
+ * same `type`.
+ */
+export function appendCachePoint(
+  block: ContentBlock,
+): (message: BaseMessage, index: number, messages: readonly BaseMessage[]) => BaseMessage {
+  return (message, index, messages) => {
+    if (index !== messages.length - 1) {
+      return message;
+    }
+
+    const { content } = message;
+
+    if (Array.isArray(content)) {
+      if (content.some((b) => b.type === block.type)) {
+        return message;
+      }
+      message.content = [...content, block];
+      return message;
+    }
+
+    if (typeof content === "string") {
+      message.content = [{ type: "text", text: content }, block] satisfies MessageContent;
+      return message;
+    }
+
+    return message;
+  };
+}

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

@@ -21,6 +21,7 @@ export {
   type LangChainAdapter,
   type LangChainAdapterConfig,
   type LangChainThreadOps,
+  type LangChainToolResponse,
 } from "./activities";
 
 // Thread manager
@@ -28,7 +29,8 @@ export {
   createLangChainThreadManager,
   type LangChainThreadManager,
   type LangChainThreadManagerConfig,
-  type LangChainToolMessageContent,
+  type LangChainContent,
+  type LangChainInvocationPayload,
 } from "./thread-manager";
 
 // Model invoker (for advanced use — prefer adapter.createModelInvoker)
@@ -37,3 +39,6 @@ export {
   invokeLangChainModel,
   type LangChainModelInvokerConfig,
 } from "./model-invoker";
+
+// Hooks / utilities
+export { appendCachePoint } from "./hooks";