zeitlich 0.2.0 → 0.2.2

package/src/lib/session.ts

@@ -114,7 +114,6 @@ export const createSession = async <T extends ToolMap>({
           metadata,
         });
       }
-
       stateManager.setTools(toolRouter.getToolDefinitions());
 
       await initializeThread(threadId);
@@ -158,29 +157,56 @@ export const createSession = async <T extends ToolMap>({
           }
 
           const rawToolCalls: RawToolCall[] = await parseToolCalls(message);
-          const parsedToolCalls = rawToolCalls
-            .filter((tc: RawToolCall) => tc.name !== "Task")
-            .map((tc: RawToolCall) => toolRouter.parseToolCall(tc));
-          const taskToolCalls =
-            subagents && subagents.length > 0
-              ? rawToolCalls
-                  .filter((tc: RawToolCall) => tc.name === "Task")
-                  .map((tc: RawToolCall) => {
-                    // Parse and validate args using the tool's schema
-                    const parsedArgs = createTaskTool(subagents).schema.parse(
-                      tc.args
-                    );
-
-                    return {
-                      id: tc.id ?? "",
-                      name: tc.name,
-                      args: parsedArgs,
-                    } as ParsedToolCall<
-                      "Task",
-                      TaskToolSchemaType<SubagentConfig[]>
-                    >;
-                  })
-              : [];
+
+          // Parse tool calls, catching schema errors and returning them to the agent
+          const parsedToolCalls: ParsedToolCallUnion<T>[] = [];
+          for (const tc of rawToolCalls.filter(
+            (tc: RawToolCall) => tc.name !== "Task"
+          )) {
+            try {
+              parsedToolCalls.push(toolRouter.parseToolCall(tc));
+            } catch (error) {
+              await appendToolResult({
+                threadId,
+                toolCallId: tc.id ?? "",
+                content: JSON.stringify({
+                  error: `Invalid tool call for "${tc.name}": ${error instanceof Error ? error.message : String(error)}`,
+                }),
+              });
+            }
+          }
+
+          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(

package/src/lib/state-manager.ts

@@ -6,6 +6,7 @@ import {
   isTerminalStatus,
 } from "./types";
 import type { ToolDefinition } from "./tool-router";
+import { z } from "zod";
 
 /**
  * JSON primitive types that Temporal can serialize
@@ -105,7 +106,7 @@ export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
   /** Delete a task by ID */
   deleteTask(id: string): boolean;
 
-  /** Set the tools */
+  /** Set the tools (converts Zod schemas to JSON Schema for serialization) */
   setTools(newTools: ToolDefinition[]): void;
 }
 
@@ -245,7 +246,13 @@ export function createAgentStateManager<
     },
 
     setTools(newTools: ToolDefinition[]): void {
-      tools = newTools;
+      tools = newTools.map((tool) => ({
+        name: tool.name,
+        description: tool.description,
+        schema: z.toJSONSchema(tool.schema) as Record<string, unknown>,
+        strict: tool.strict,
+        max_uses: tool.max_uses,
+      }));
     },
 
     deleteTask(id: string): boolean {

package/src/lib/tool-router.ts

@@ -1,6 +1,15 @@
 import type { MessageToolDefinition } from "@langchain/core/messages";
 import type { ToolMessageContent } from "./thread-manager";
-import type { Hooks, SubagentConfig, ToolResultConfig } from "./types";
+import type {
+  Hooks,
+  PostToolUseFailureHookResult,
+  PreToolUseHookResult,
+  SubagentConfig,
+  SubagentHooks,
+  ToolHooks,
+  ToolResultConfig,
+} from "./types";
+import type { GenericTaskToolSchemaType } from "../tools/task/tool";
 
 import type { z } from "zod";
 import { proxyActivities } from "@temporalio/workflow";
@@ -49,6 +58,8 @@ export interface ToolWithHandler<
   handler: ToolHandler<z.infer<TSchema>, TResult, TContext>;
   strict?: boolean;
   max_uses?: number;
+  /** Per-tool lifecycle hooks (run in addition to global hooks) */
+  hooks?: ToolHooks<z.infer<TSchema>, TResult>;
 }
 
 /**
@@ -68,6 +79,9 @@ export type ToolMap = Record<
     /* eslint-enable @typescript-eslint/no-explicit-any */
     strict?: boolean;
     max_uses?: number;
+    /* eslint-disable @typescript-eslint/no-explicit-any */
+    hooks?: ToolHooks<any, any>;
+    /* eslint-enable @typescript-eslint/no-explicit-any */
   }
 >;
 
@@ -128,10 +142,10 @@ export type AppendToolResultFn = (config: ToolResultConfig) => Promise<void>;
  * Contains the content for the tool message and the result to return from processToolCalls.
  */
 export interface ToolHandlerResponse<TResult> {
-  /** Content for the tool message added to the thread */
-  content: ToolMessageContent;
-  /** Result returned from processToolCalls */
-  result: TResult;
+  /** Content sent back to the LLM as the tool call response */
+  toolResponse: ToolMessageContent;
+  /** Data returned to the workflow and hooks for further processing */
+  data: TResult | null;
 }
 
 /**
@@ -225,7 +239,7 @@ export interface ToolCallResult<
 > {
   toolCallId: string;
   name: TName;
-  result: TResult;
+  data: TResult | null;
 }
 
 /**
@@ -424,9 +438,36 @@ export function createToolRouter<T extends ToolMap>(
   }
 
   if (options.subagents) {
+    // Build per-subagent hook dispatcher keyed by subagent name
+    const subagentHooksMap = new Map<string, SubagentHooks>();
+    for (const s of options.subagents) {
+      if (s.hooks) subagentHooksMap.set(s.name, s.hooks);
+    }
+
+    const resolveSubagentName = (args: unknown): string =>
+      (args as GenericTaskToolSchemaType).subagent;
+
     toolMap.set("Task", {
       ...createTaskTool(options.subagents),
       handler: createTaskHandler(options.subagents),
+      ...(subagentHooksMap.size > 0 && {
+        hooks: {
+          onPreToolUse: async (ctx): Promise<PreToolUseHookResult> => {
+            const hooks = subagentHooksMap.get(resolveSubagentName(ctx.args));
+            return hooks?.onPreExecution?.(ctx) ?? {};
+          },
+          onPostToolUse: async (ctx): Promise<void> => {
+            const hooks = subagentHooksMap.get(resolveSubagentName(ctx.args));
+            await hooks?.onPostExecution?.(ctx);
+          },
+          onPostToolUseFailure: async (
+            ctx
+          ): Promise<PostToolUseFailureHookResult> => {
+            const hooks = subagentHooksMap.get(resolveSubagentName(ctx.args));
+            return hooks?.onExecutionFailure?.(ctx) ?? {};
+          },
+        } satisfies ToolHooks,
+      }),
     });
   }
 
@@ -455,8 +496,10 @@ export function createToolRouter<T extends ToolMap>(
     handlerContext?: ToolHandlerContext
   ): Promise<ToolCallResultUnion<TResults> | null> {
     const startTime = Date.now();
+    const tool = toolMap.get(toolCall.name);
+    const toolHooks = tool?.hooks;
 
-    // PreToolUse hook - can skip or modify args
+    // --- PreToolUse: global then per-tool ---
     let effectiveArgs: unknown = toolCall.args;
     if (options.hooks?.onPreToolUse) {
       const preResult = await options.hooks.onPreToolUse({
@@ -465,7 +508,6 @@ export function createToolRouter<T extends ToolMap>(
         turn,
       });
       if (preResult?.skip) {
-        // Skip this tool call - append a skip message and return null
         await appendToolResult({
           threadId: options.threadId,
           toolCallId: toolCall.id,
@@ -480,10 +522,31 @@ export function createToolRouter<T extends ToolMap>(
         effectiveArgs = preResult.modifiedArgs;
       }
     }
+    if (toolHooks?.onPreToolUse) {
+      const preResult = await toolHooks.onPreToolUse({
+        args: effectiveArgs,
+        threadId: options.threadId,
+        turn,
+      });
+      if (preResult?.skip) {
+        await appendToolResult({
+          threadId: options.threadId,
+          toolCallId: toolCall.id,
+          content: JSON.stringify({
+            skipped: true,
+            reason: "Skipped by tool PreToolUse hook",
+          }),
+        });
+        return null;
+      }
+      if (preResult?.modifiedArgs !== undefined) {
+        effectiveArgs = preResult.modifiedArgs;
+      }
+    }
 
-    const tool = toolMap.get(toolCall.name);
+    // --- Execute handler ---
     let result: unknown;
-    let content: ToolMessageContent;
+    let content!: ToolMessageContent;
 
     try {
       if (tool) {
@@ -491,31 +554,54 @@ export function createToolRouter<T extends ToolMap>(
           effectiveArgs as Parameters<typeof tool.handler>[0],
           (handlerContext ?? {}) as Parameters<typeof tool.handler>[1]
         );
-        result = response.result;
-        content = response.content;
+        result = response.data;
+        content = response.toolResponse;
       } else {
         result = { error: `Unknown tool: ${toolCall.name}` };
         content = JSON.stringify(result, null, 2);
       }
     } catch (error) {
-      // PostToolUseFailure hook - can recover from errors
-      if (options.hooks?.onPostToolUseFailure) {
+      // --- PostToolUseFailure: per-tool then global ---
+      const err = error instanceof Error ? error : new Error(String(error));
+      let recovered = false;
+
+      if (toolHooks?.onPostToolUseFailure) {
+        const failureResult = await toolHooks.onPostToolUseFailure({
+          args: effectiveArgs,
+          error: err,
+          threadId: options.threadId,
+          turn,
+        });
+        if (failureResult?.fallbackContent !== undefined) {
+          content = failureResult.fallbackContent;
+          result = { error: String(error), recovered: true };
+          recovered = true;
+        } else if (failureResult?.suppress) {
+          content = JSON.stringify({ error: String(error), suppressed: true });
+          result = { error: String(error), suppressed: true };
+          recovered = true;
+        }
+      }
+
+      if (!recovered && options.hooks?.onPostToolUseFailure) {
         const failureResult = await options.hooks.onPostToolUseFailure({
           toolCall,
-          error: error instanceof Error ? error : new Error(String(error)),
+          error: err,
           threadId: options.threadId,
           turn,
         });
         if (failureResult?.fallbackContent !== undefined) {
           content = failureResult.fallbackContent;
           result = { error: String(error), recovered: true };
+          recovered = true;
         } else if (failureResult?.suppress) {
           content = JSON.stringify({ error: String(error), suppressed: true });
           result = { error: String(error), suppressed: true };
-        } else {
-          throw error;
+          recovered = true;
         }
-      } else {
+      }
+
+      if (!recovered) {
         throw error;
       }
     }
@@ -530,12 +616,21 @@ export function createToolRouter<T extends ToolMap>(
     const toolResult = {
       toolCallId: toolCall.id,
       name: toolCall.name,
-      result,
+      data: result,
     } as ToolCallResultUnion<TResults>;
 
-    // PostToolUse hook - called after successful execution
+    // --- PostToolUse: per-tool then global ---
+    const durationMs = Date.now() - startTime;
+    if (toolHooks?.onPostToolUse) {
+      await toolHooks.onPostToolUse({
+        args: effectiveArgs,
+        result: result,
+        threadId: options.threadId,
+        turn,
+        durationMs,
+      });
+    }
     if (options.hooks?.onPostToolUse) {
-      const durationMs = Date.now() - startTime;
       await options.hooks.onPostToolUse({
         toolCall,
         result: toolResult,
@@ -654,13 +749,13 @@ export function createToolRouter<T extends ToolMap>(
         await appendToolResult({
           threadId: options.threadId,
           toolCallId: toolCall.id,
-          content: response.content,
+          content: response.toolResponse,
         });
 
         return {
           toolCallId: toolCall.id,
           name: toolCall.name as TName,
-          result: response.result,
+          data: response.data ?? null,
         };
       };
 
@@ -706,6 +801,93 @@ export function createToolRouter<T extends ToolMap>(
   };
 }
 
+/**
+ * Identity function that creates a generic inference context for a tool definition.
+ * TypeScript infers TResult from the handler and flows it to hooks automatically.
+ *
+ * @example
+ * ```typescript
+ * tools: {
+ *   AskUser: defineTool({
+ *     ...askUserTool,
+ *     handler: handleAskUser,
+ *     hooks: {
+ *       onPostToolUse: ({ result }) => {
+ *         // result is correctly typed as the handler's return data type
+ *       },
+ *     },
+ *   }),
+ * }
+ * ```
+ */
+export function defineTool<
+  TName extends string,
+  TSchema extends z.ZodType,
+  TResult,
+  TContext = ToolHandlerContext,
+>(
+  tool: ToolWithHandler<TName, TSchema, TResult, TContext>
+): ToolWithHandler<TName, TSchema, TResult, TContext> {
+  return tool;
+}
+
+/**
+ * Identity function that provides full type inference for subagent configurations.
+ * Verifies the workflow function's input parameters match the configured context,
+ * and properly types the lifecycle hooks with Task tool args and inferred result type.
+ *
+ * @example
+ * ```ts
+ * // With typed context — workflow must accept { prompt, context }
+ * const researcher = defineSubagent({
+ *   name: "researcher",
+ *   description: "Researches topics",
+ *   workflow: researcherWorkflow, // (input: { prompt: string; context: { apiKey: string } }) => Promise<...>
+ *   context: { apiKey: "..." },
+ *   resultSchema: z.object({ findings: z.string() }),
+ *   hooks: {
+ *     onPostExecution: ({ result }) => {
+ *       // result is typed as { findings: string }
+ *     },
+ *   },
+ * });
+ *
+ * // Without context — workflow only needs { prompt }
+ * const writer = defineSubagent({
+ *   name: "writer",
+ *   description: "Writes content",
+ *   workflow: writerWorkflow, // (input: { prompt: string }) => Promise<...>
+ *   resultSchema: z.object({ content: z.string() }),
+ * });
+ * ```
+ */
+// With context — verifies workflow accepts { prompt, context: TContext }
+export function defineSubagent<
+  TResult extends z.ZodType = z.ZodType,
+  TContext extends Record<string, unknown> = Record<string, unknown>,
+>(
+  config: Omit<SubagentConfig<TResult>, "hooks" | "workflow" | "context"> & {
+    workflow:
+      | string
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      | ((input: { prompt: string; context: TContext }) => Promise<any>);
+    context: TContext;
+    hooks?: SubagentHooks<GenericTaskToolSchemaType, z.infer<TResult>>;
+  }
+): SubagentConfig<TResult>;
+// Without context — verifies workflow accepts { prompt }
+export function defineSubagent<TResult extends z.ZodType = z.ZodType>(
+  config: Omit<SubagentConfig<TResult>, "hooks" | "workflow"> & {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    workflow: string | ((input: { prompt: string }) => Promise<any>);
+    hooks?: SubagentHooks<GenericTaskToolSchemaType, z.infer<TResult>>;
+  }
+): SubagentConfig<TResult>;
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function defineSubagent(config: any): SubagentConfig {
+  return config;
+}
+
 /**
  * Utility to check if there were no tool calls besides a specific one
  */

package/src/lib/types.ts

@@ -4,11 +4,11 @@ import type {
   InferToolResults,
   ParsedToolCallUnion,
   ToolCallResultUnion,
-  ToolDefinition,
   ToolMap,
 } from "./tool-router";
 
 import type { MessageContent, StoredMessage } from "@langchain/core/messages";
+import type { Workflow } from "@temporalio/workflow";
 import type { z } from "zod";
 
 /**
@@ -25,7 +25,7 @@ export type AgentStatus =
  * Base state that all agents must have
  */
 export interface BaseAgentState {
-  tools: ToolDefinition[];
+  tools: SerializableToolDefinition[];
   status: AgentStatus;
   version: number;
   turns: number;
@@ -103,6 +103,19 @@ export interface ZeitlichAgentConfig<T extends ToolMap> {
   };
 }
 
+/**
+ * A JSON-serializable tool definition for state storage.
+ * Uses a plain JSON Schema object instead of a live Zod instance,
+ * so it survives Temporal serialization without losing constraints (min, max, etc.).
+ */
+export interface SerializableToolDefinition {
+  name: string;
+  description: string;
+  schema: Record<string, unknown>;
+  strict?: boolean;
+  max_uses?: number;
+}
+
 /**
  * Configuration passed to runAgent activity
  */
@@ -142,12 +155,44 @@ export interface SubagentConfig<TResult extends z.ZodType = z.ZodType> {
   name: string;
   /** Description shown to the parent agent explaining what this subagent does */
   description: string;
-  /** Temporal workflow type name (used with executeChild) */
-  workflowType: string;
+  /** Temporal workflow function or type name (used with executeChild) */
+  workflow: string | Workflow;
   /** Optional task queue - defaults to parent's queue if not specified */
   taskQueue?: string;
   /** Optional Zod schema to validate the child workflow's result. If omitted, result is passed through as-is. */
   resultSchema?: TResult;
+  /** Optional static context passed to the subagent on every invocation */
+  context?: Record<string, unknown>;
+  /** Per-subagent lifecycle hooks */
+  hooks?: SubagentHooks;
+}
+
+/**
+ * Per-subagent lifecycle hooks - defined on a SubagentConfig.
+ * Runs in addition to global hooks (global pre → subagent pre → execute → subagent post → global post).
+ */
+export interface SubagentHooks<TArgs = unknown, TResult = unknown> {
+  /** Called before this subagent executes - can skip or modify args */
+  onPreExecution?: (ctx: {
+    args: TArgs;
+    threadId: string;
+    turn: number;
+  }) => PreToolUseHookResult | Promise<PreToolUseHookResult>;
+  /** Called after this subagent executes successfully */
+  onPostExecution?: (ctx: {
+    args: TArgs;
+    result: TResult;
+    threadId: string;
+    turn: number;
+    durationMs: number;
+  }) => void | Promise<void>;
+  /** Called when this subagent execution fails */
+  onExecutionFailure?: (ctx: {
+    args: TArgs;
+    error: Error;
+    threadId: string;
+    turn: number;
+  }) => PostToolUseFailureHookResult | Promise<PostToolUseFailureHookResult>;
 }
 
 /**
@@ -156,6 +201,8 @@ export interface SubagentConfig<TResult extends z.ZodType = z.ZodType> {
 export interface SubagentInput {
   /** The prompt/task from the parent agent */
   prompt: string;
+  /** Optional context parameters passed from the parent agent */
+  context?: Record<string, unknown>;
 }
 
 // ============================================================================
@@ -328,6 +375,34 @@ export type SessionEndHook = (
   ctx: SessionEndHookContext
 ) => void | Promise<void>;
 
+/**
+ * Per-tool lifecycle hooks - defined directly on a tool definition.
+ * Runs in addition to global hooks (global pre → tool pre → execute → tool post → global post).
+ */
+export interface ToolHooks<TArgs = unknown, TResult = unknown> {
+  /** Called before this tool executes - can skip or modify args */
+  onPreToolUse?: (ctx: {
+    args: TArgs;
+    threadId: string;
+    turn: number;
+  }) => PreToolUseHookResult | Promise<PreToolUseHookResult>;
+  /** Called after this tool executes successfully */
+  onPostToolUse?: (ctx: {
+    args: TArgs;
+    result: TResult | null;
+    threadId: string;
+    turn: number;
+    durationMs: number;
+  }) => void | Promise<void>;
+  /** Called when this tool execution fails */
+  onPostToolUseFailure?: (ctx: {
+    args: TArgs;
+    error: Error;
+    threadId: string;
+    turn: number;
+  }) => PostToolUseFailureHookResult | Promise<PostToolUseFailureHookResult>;
+}
+
 /**
  * Combined hooks interface for session lifecycle
  */

package/src/tools/ask-user-question/handler.ts

@@ -21,5 +21,5 @@ export const handleAskUserQuestionToolResult: ActivityToolHandler<
       }).toDict()
   );
 
-  return { content: "Question submitted", result: { chatMessages: messages } };
+  return { toolResponse: "Question submitted", data: { chatMessages: messages } };
 };