zeitlich 0.2.2 → 0.2.3

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/session.ts

@@ -1,36 +1,25 @@
 import { proxyActivities } from "@temporalio/workflow";
 import type { ZeitlichSharedActivities } from "../activities";
 import type {
+  ThreadOps,
   ZeitlichAgentConfig,
   SessionStartHook,
   SessionEndHook,
   SessionExitReason,
-  SubagentConfig,
 } from "./types";
+import type { StoredMessage } from "@langchain/core/messages";
 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 +32,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 +71,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 +81,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";
 
@@ -156,19 +114,19 @@ export const createSession = async <T extends ToolMap>({
             return message;
           }
 
-          const rawToolCalls: RawToolCall[] = await parseToolCalls(message);
+          const rawToolCalls: RawToolCall[] =
+            await threadOps.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 +134,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 +161,39 @@ 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<StoredMessage> {
+  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,
+    parseToolCalls: activities.parseToolCalls,
+  };
+}

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);
 }