zeitlich 0.2.2 → 0.2.4

package/src/activities.ts

@@ -2,18 +2,14 @@ import type Redis from "ioredis";
 import { createThreadManager } from "./lib/thread-manager";
 import type { ToolResultConfig } from "./lib/types";
 import {
-  type AIMessage,
-  mapStoredMessageToChatMessage,
   type MessageContent,
   type StoredMessage,
 } from "@langchain/core/messages";
-import type { RawToolCall } from "./lib/tool-router";
 /**
  * Shared Zeitlich activities - thread management and message handling
  * Note: runAgent is workflow-specific and should be created per-workflow
  */
 export interface ZeitlichSharedActivities {
-  appendSystemMessage(threadId: string, content: string): Promise<void>;
   /**
    * Append a tool result to the thread.
    * Handles JSON serialization and optional cache points.
@@ -25,11 +21,6 @@ export interface ZeitlichSharedActivities {
    */
   initializeThread(threadId: string): Promise<void>;
 
-  /**
-   * Append a system message to a thread.
-   */
-  appendSystemMessage(threadId: string, content: string): Promise<void>;
-
   /**
    * Append messages to a thread.
    */
@@ -46,11 +37,6 @@ export interface ZeitlichSharedActivities {
     content: string | MessageContent
   ): Promise<void>;
 
-  /**
-   * Extract raw tool calls from a stored message.
-   * Returns unvalidated tool calls - use toolRegistry.parseToolCall() to validate.
-   */
-  parseToolCalls(storedMessage: StoredMessage): Promise<RawToolCall[]>;
 }
 
 /**
@@ -62,14 +48,6 @@ export interface ZeitlichSharedActivities {
  */
 export function createSharedActivities(redis: Redis): ZeitlichSharedActivities {
   return {
-    async appendSystemMessage(
-      threadId: string,
-      content: string
-    ): Promise<void> {
-      const thread = createThreadManager({ redis, threadId });
-      await thread.appendSystemMessage(content);
-    },
-
     async appendToolResult(config: ToolResultConfig): Promise<void> {
       const { threadId, toolCallId, content } = config;
       const thread = createThreadManager({ redis, threadId });
@@ -98,15 +76,5 @@ export function createSharedActivities(redis: Redis): ZeitlichSharedActivities {
       await thread.appendHumanMessage(content);
     },
 
-    async parseToolCalls(storedMessage: StoredMessage): Promise<RawToolCall[]> {
-      const message = mapStoredMessageToChatMessage(storedMessage) as AIMessage;
-      const toolCalls = message.tool_calls ?? [];
-
-      return toolCalls.map((toolCall) => ({
-        id: toolCall.id,
-        name: toolCall.name,
-        args: toolCall.args,
-      }));
-    },
   };
 }

package/src/index.ts

@@ -9,7 +9,7 @@
  * @example
  * ```typescript
  * // In your activities file
- * import { invokeModel, globHandler } from '@bead-ai/zeitlich';
+ * import { invokeModel, createGlobHandler } from '@bead-ai/zeitlich';
  *
  * // In your worker file
  * import { ZeitlichPlugin } from '@bead-ai/zeitlich';
@@ -28,25 +28,28 @@ export type { ZeitlichPluginOptions } from "./plugin";
 export { createSharedActivities } from "./activities";
 export type { ZeitlichSharedActivities } from "./activities";
 
+// Auto-append wrapper for large tool results (activity-side only)
+export { withAutoAppend } from "./lib/tool-router";
+
 // Model invocation (requires Redis, LangChain)
 export { invokeModel } from "./lib/model-invoker";
 export type { InvokeModelConfig } from "./lib/model-invoker";
 
 // Tool handlers (activity implementations)
-// These are direct functions that accept scopedNodes per-call for dynamic file trees
-export { handleAskUserQuestionToolResult } from "./tools/ask-user-question/handler";
-export { globHandler } from "./tools/glob/handler";
+// All handlers follow the factory pattern: createXHandler(deps) => handler(args)
+export { createAskUserQuestionHandler } from "./tools/ask-user-question/handler";
+export { createGlobHandler } from "./tools/glob/handler";
 
-export { editHandler } from "./tools/edit/handler";
-export type {
-  EditResult,
-  EditHandlerResponse,
-  EditHandlerOptions,
-} from "./tools/edit/handler";
+export { createEditHandler } from "./tools/edit/handler";
 
-export { handleBashTool } from "./tools/bash/handler";
+export { createBashHandler } from "./tools/bash/handler";
 
 export { toTree } from "./lib/fs";
 
 export { getStateQuery } from "./lib/state-manager";
 export { createThreadManager } from "./lib/thread-manager";
+export type {
+  BaseThreadManager,
+  ThreadManager,
+  ThreadManagerConfig,
+} from "./lib/thread-manager";

package/src/lib/model-invoker.ts

@@ -63,9 +63,15 @@ export async function invokeModel({
 
   await thread.append([response.toDict()]);
 
+  const toolCalls = response.tool_calls ?? [];
+
   return {
     message: response.toDict(),
-    stopReason: (response.response_metadata?.stop_reason as string) ?? null,
+    rawToolCalls: toolCalls.map((tc) => ({
+      id: tc.id,
+      name: tc.name,
+      args: tc.args,
+    })),
     usage: {
       input_tokens: response.usage_metadata?.input_tokens,
       output_tokens: response.usage_metadata?.output_tokens,

package/src/lib/session.ts

@@ -1,36 +1,23 @@
 import { proxyActivities } from "@temporalio/workflow";
 import type { ZeitlichSharedActivities } from "../activities";
 import type {
+  ThreadOps,
   ZeitlichAgentConfig,
   SessionStartHook,
   SessionEndHook,
   SessionExitReason,
-  SubagentConfig,
 } from "./types";
 import { type AgentStateManager, type JsonSerializable } from "./state-manager";
 import {
   createToolRouter,
-  type ParsedToolCall,
   type ParsedToolCallUnion,
-  type RawToolCall,
   type ToolMap,
 } from "./tool-router";
-import type { StoredMessage } from "@langchain/core/messages";
-import { createTaskTool, type TaskToolSchemaType } from "../tools/task/tool";
 
-export interface ZeitlichSession {
+export interface ZeitlichSession<M = unknown> {
   runSession<T extends JsonSerializable<T>>(args: {
     stateManager: AgentStateManager<T>;
-  }): Promise<StoredMessage | null>;
-}
-
-async function resolvePrompt(
-  prompt: string | (() => string | Promise<string>)
-): Promise<string> {
-  if (typeof prompt === "function") {
-    return prompt();
-  }
-  return prompt;
+  }): Promise<M | null>;
 }
 
 /**
@@ -43,48 +30,24 @@ export interface SessionLifecycleHooks {
   onSessionEnd?: SessionEndHook;
 }
 
-export const createSession = async <T extends ToolMap>({
+export const createSession = async <T extends ToolMap, M = unknown>({
   threadId,
   agentName,
   maxTurns = 50,
   metadata = {},
   runAgent,
-  baseSystemPrompt,
-  instructionsPrompt,
+  threadOps,
   buildContextMessage,
-  buildFileTree = async (): Promise<string> => "",
   subagents,
   tools = {} as T,
   processToolsInParallel = true,
-  buildInTools = {},
   hooks = {},
-}: ZeitlichAgentConfig<T>): Promise<ZeitlichSession> => {
-  const {
-    initializeThread,
-    appendHumanMessage,
-    parseToolCalls,
-    appendToolResult,
-    appendSystemMessage,
-  } = proxyActivities<ZeitlichSharedActivities>({
-    startToCloseTimeout: "30m",
-    retry: {
-      maximumAttempts: 6,
-      initialInterval: "5s",
-      maximumInterval: "15m",
-      backoffCoefficient: 4,
-    },
-    heartbeatTimeout: "5m",
-  });
-
-  const fileTree = await buildFileTree();
-
+}: ZeitlichAgentConfig<T, M>): Promise<ZeitlichSession<M>> => {
   const toolRouter = createToolRouter({
     tools,
-    appendToolResult,
+    appendToolResult: threadOps.appendToolResult,
     threadId,
     hooks,
-    buildInTools,
-    fileTree,
     subagents,
     parallel: processToolsInParallel,
   });
@@ -106,7 +69,7 @@ export const createSession = async <T extends ToolMap>({
   };
 
   return {
-    runSession: async ({ stateManager }): Promise<StoredMessage | null> => {
+    runSession: async ({ stateManager }): Promise<M | null> => {
       if (hooks.onSessionStart) {
         await hooks.onSessionStart({
           threadId,
@@ -116,15 +79,8 @@ export const createSession = async <T extends ToolMap>({
       }
       stateManager.setTools(toolRouter.getToolDefinitions());
 
-      await initializeThread(threadId);
-      await appendSystemMessage(
-        threadId,
-        [
-          await resolvePrompt(baseSystemPrompt),
-          await resolvePrompt(instructionsPrompt),
-        ].join("\n")
-      );
-      await appendHumanMessage(threadId, await buildContextMessage());
+      await threadOps.initializeThread(threadId);
+      await threadOps.appendHumanMessage(threadId, await buildContextMessage());
 
       let exitReason: SessionExitReason = "completed";
 
@@ -137,38 +93,29 @@ export const createSession = async <T extends ToolMap>({
           stateManager.incrementTurns();
           const currentTurn = stateManager.getTurns();
 
-          const { message, stopReason } = await runAgent({
+          const { message, rawToolCalls } = await runAgent({
             threadId,
             agentName,
             metadata,
           });
 
-          if (stopReason === "end_turn") {
-            stateManager.complete();
-            exitReason = "completed";
-            return message;
-          }
-
           // No tools configured - treat any non-end_turn as completed
-          if (!toolRouter.hasTools()) {
+          if (!toolRouter.hasTools() || rawToolCalls.length === 0) {
             stateManager.complete();
             exitReason = "completed";
             return message;
           }
 
-          const rawToolCalls: RawToolCall[] = await parseToolCalls(message);
-
-          // Parse tool calls, catching schema errors and returning them to the agent
+          // Parse all tool calls uniformly through the router
           const parsedToolCalls: ParsedToolCallUnion<T>[] = [];
-          for (const tc of rawToolCalls.filter(
-            (tc: RawToolCall) => tc.name !== "Task"
-          )) {
+          for (const tc of rawToolCalls) {
             try {
               parsedToolCalls.push(toolRouter.parseToolCall(tc));
             } catch (error) {
-              await appendToolResult({
+              await threadOps.appendToolResult({
                 threadId,
                 toolCallId: tc.id ?? "",
+                toolName: tc.name,
                 content: JSON.stringify({
                   error: `Invalid tool call for "${tc.name}": ${error instanceof Error ? error.message : String(error)}`,
                 }),
@@ -176,47 +123,10 @@ export const createSession = async <T extends ToolMap>({
             }
           }
 
-          const taskToolCalls: ParsedToolCall<
-            "Task",
-            TaskToolSchemaType<SubagentConfig[]>
-          >[] = [];
-          if (subagents && subagents.length > 0) {
-            for (const tc of rawToolCalls.filter(
-              (tc: RawToolCall) => tc.name === "Task"
-            )) {
-              try {
-                const parsedArgs = createTaskTool(subagents).schema.parse(
-                  tc.args
-                );
-                taskToolCalls.push({
-                  id: tc.id ?? "",
-                  name: tc.name,
-                  args: parsedArgs,
-                } as ParsedToolCall<
-                  "Task",
-                  TaskToolSchemaType<SubagentConfig[]>
-                >);
-              } catch (error) {
-                await appendToolResult({
-                  threadId,
-                  toolCallId: tc.id ?? "",
-                  content: JSON.stringify({
-                    error: `Invalid tool call for "Task": ${error instanceof Error ? error.message : String(error)}`,
-                  }),
-                });
-              }
-            }
-          }
-
           // Hooks can call stateManager.waitForInput() to pause the session
-          await toolRouter.processToolCalls(
-            [...parsedToolCalls, ...taskToolCalls] as ParsedToolCallUnion<
-              T & { Task: TaskToolSchemaType<SubagentConfig[]> }
-            >[],
-            {
-              turn: currentTurn,
-            }
-          );
+          await toolRouter.processToolCalls(parsedToolCalls, {
+            turn: currentTurn,
+          });
 
           if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
             exitReason = "waiting_for_input";
@@ -240,3 +150,38 @@ export const createSession = async <T extends ToolMap>({
     },
   };
 };
+
+/**
+ * Proxy the default ZeitlichSharedActivities as ThreadOps<StoredMessage>.
+ * Call this in workflow code for the standard LangChain/StoredMessage setup.
+ *
+ * @example
+ * ```typescript
+ * const session = await createSession({
+ *   threadOps: proxyDefaultThreadOps(),
+ *   // ...
+ * });
+ * ```
+ */
+export function proxyDefaultThreadOps(
+  options?: Parameters<typeof proxyActivities>[0]
+): ThreadOps {
+  const activities = proxyActivities<ZeitlichSharedActivities>(
+    options ?? {
+      startToCloseTimeout: "30m",
+      retry: {
+        maximumAttempts: 6,
+        initialInterval: "5s",
+        maximumInterval: "15m",
+        backoffCoefficient: 4,
+      },
+      heartbeatTimeout: "5m",
+    }
+  );
+
+  return {
+    initializeThread: activities.initializeThread,
+    appendHumanMessage: activities.appendHumanMessage,
+    appendToolResult: activities.appendToolResult,
+  };
+}

package/src/lib/thread-manager.ts

@@ -7,7 +7,6 @@ import {
   type MessageContent,
   type MessageStructure,
   type StoredMessage,
-  SystemMessage,
   ToolMessage,
 } from "@langchain/core/messages";
 import { v4 as uuidv4 } from "uuid";
@@ -24,28 +23,31 @@ function getThreadKey(threadId: string, key: string): string {
  */
 export type ToolMessageContent = $InferMessageContent<MessageStructure, "tool">;
 
-export interface ThreadManagerConfig {
+export interface ThreadManagerConfig<T = StoredMessage> {
   redis: Redis;
   threadId: string;
   /** Thread key, defaults to 'messages' */
   key?: string;
+  /** Custom serializer, defaults to JSON.stringify */
+  serialize?: (message: T) => string;
+  /** Custom deserializer, defaults to JSON.parse */
+  deserialize?: (raw: string) => T;
 }
 
-export interface ThreadManager {
-  /** Append a system message to the thread */
-  appendSystemMessage(content: string): Promise<void>;
+/** Generic thread manager for any message type */
+export interface BaseThreadManager<T> {
   /** Initialize an empty thread */
   initialize(): Promise<void>;
   /** Load all messages from the thread */
-  load(): Promise<StoredMessage[]>;
+  load(): Promise<T[]>;
   /** Append messages to the thread */
-  append(messages: StoredMessage[]): Promise<void>;
+  append(messages: T[]): Promise<void>;
   /** Delete the thread */
   delete(): Promise<void>;
+}
 
-  /** Create a SystemMessage (returns StoredMessage for storage) */
-  createSystemMessage(content: string): StoredMessage;
-
+/** Thread manager with StoredMessage convenience helpers */
+export interface ThreadManager extends BaseThreadManager<StoredMessage> {
   /** Create a HumanMessage (returns StoredMessage for storage) */
   createHumanMessage(content: string | MessageContent): StoredMessage;
 
@@ -74,26 +76,38 @@ export interface ThreadManager {
 
 /**
  * Creates a thread manager for handling conversation state in Redis.
+ * Without generic args, returns a full ThreadManager with StoredMessage helpers.
+ * With a custom type T, returns a BaseThreadManager<T>.
  */
-export function createThreadManager(
-  config: ThreadManagerConfig
-): ThreadManager {
-  const { redis, threadId, key = "messages" } = config;
+export function createThreadManager(config: ThreadManagerConfig): ThreadManager;
+export function createThreadManager<T>(
+  config: ThreadManagerConfig<T>
+): BaseThreadManager<T>;
+export function createThreadManager<T>(
+  config: ThreadManagerConfig<T>
+): BaseThreadManager<T> {
+  const {
+    redis,
+    threadId,
+    key = "messages",
+    serialize = (m: T): string => JSON.stringify(m),
+    deserialize = (raw: string): T => JSON.parse(raw) as T,
+  } = config;
   const redisKey = getThreadKey(threadId, key);
 
-  return {
+  const base: BaseThreadManager<T> = {
     async initialize(): Promise<void> {
       await redis.del(redisKey);
     },
 
-    async load(): Promise<StoredMessage[]> {
+    async load(): Promise<T[]> {
       const data = await redis.lrange(redisKey, 0, -1);
-      return data.map((item) => JSON.parse(item) as StoredMessage);
+      return data.map(deserialize);
     },
 
-    async append(messages: StoredMessage[]): Promise<void> {
+    async append(messages: T[]): Promise<void> {
       if (messages.length > 0) {
-        await redis.rpush(redisKey, ...messages.map((m) => JSON.stringify(m)));
+        await redis.rpush(redisKey, ...messages.map(serialize));
         await redis.expire(redisKey, THREAD_TTL_SECONDS);
       }
     },
@@ -101,7 +115,11 @@ export function createThreadManager(
     async delete(): Promise<void> {
       await redis.del(redisKey);
     },
+  };
 
+  // If no custom serialize/deserialize were provided and T defaults to StoredMessage,
+  // the overload guarantees the caller gets ThreadManager with convenience helpers.
+  const helpers = {
     createHumanMessage(content: string | MessageContent): StoredMessage {
       return new HumanMessage({
         id: uuidv4(),
@@ -131,39 +149,29 @@ export function createThreadManager(
       toolCallId: string
     ): StoredMessage {
       return new ToolMessage({
-        // Cast needed due to langchain type compatibility
         content: content as MessageContent,
         tool_call_id: toolCallId,
       }).toDict();
     },
 
-    createSystemMessage(content: string): StoredMessage {
-      return new SystemMessage({
-        content,
-      }).toDict();
-    },
-
-    async appendSystemMessage(content: string): Promise<void> {
-      const message = this.createSystemMessage(content);
-      await this.append([message]);
-    },
-
     async appendHumanMessage(content: string | MessageContent): Promise<void> {
-      const message = this.createHumanMessage(content);
-      await this.append([message]);
+      const message = helpers.createHumanMessage(content);
+      await (base as BaseThreadManager<StoredMessage>).append([message]);
     },
 
     async appendToolMessage(
       content: ToolMessageContent,
       toolCallId: string
     ): Promise<void> {
-      const message = this.createToolMessage(content, toolCallId);
-      await this.append([message]);
+      const message = helpers.createToolMessage(content, toolCallId);
+      await (base as BaseThreadManager<StoredMessage>).append([message]);
     },
 
     async appendAIMessage(content: string | MessageContent): Promise<void> {
-      const message = this.createAIMessage(content);
-      await this.append([message]);
+      const message = helpers.createAIMessage(content as string);
+      await (base as BaseThreadManager<StoredMessage>).append([message]);
     },
   };
+
+  return Object.assign(base, helpers);
 }