zeitlich 0.2.7 → 0.2.9

package/src/lib/session.ts

@@ -24,7 +24,11 @@ import type { MessageContent } from "@langchain/core/messages";
 export interface ZeitlichSession<M = unknown> {
   runSession<T extends JsonSerializable<T>>(args: {
     stateManager: AgentStateManager<T>;
-  }): Promise<M | null>;
+  }): Promise<{
+    finalMessage: M | null;
+    exitReason: SessionExitReason;
+    usage: ReturnType<AgentStateManager<T>["getTotalUsage"]>;
+  }>;
 }
 
 /**
@@ -40,6 +44,7 @@ export interface SessionLifecycleHooks {
 export const createSession = async <T extends ToolMap, M = unknown>({
   threadId,
   agentName,
+  description,
   maxTurns = 50,
   metadata = {},
   runAgent,
@@ -86,7 +91,13 @@ export const createSession = async <T extends ToolMap, M = unknown>({
   };
 
   return {
-    runSession: async ({ stateManager }): Promise<M | null> => {
+    runSession: async ({
+      stateManager,
+    }): Promise<{
+      finalMessage: M | null;
+      exitReason: SessionExitReason;
+      usage: ReturnType<typeof stateManager.getTotalUsage>;
+    }> => {
       setHandler(
         defineUpdate<unknown, [MessageContent]>(`add${agentName}Message`),
         async (message: MessageContent) => {
@@ -115,8 +126,6 @@ export const createSession = async <T extends ToolMap, M = unknown>({
         });
       }
 
-      stateManager.setTools(toolRouter.getToolDefinitions());
-
       await initializeThread(threadId);
       if (appendSystemPrompt && systemPrompt && systemPrompt.trim() !== "") {
         await appendSystemMessage(threadId, systemPrompt);
@@ -134,17 +143,29 @@ export const createSession = async <T extends ToolMap, M = unknown>({
           stateManager.incrementTurns();
           const currentTurn = stateManager.getTurns();
 
-          const { message, rawToolCalls } = await runAgent({
+          stateManager.setTools(toolRouter.getToolDefinitions());
+
+          const { message, rawToolCalls, usage } = await runAgent({
             threadId,
             agentName,
             metadata,
+            systemPrompt,
+            description,
           });
 
+          if (usage) {
+            stateManager.updateUsage(usage);
+          }
+
           // No tools configured - treat any non-end_turn as completed
           if (!toolRouter.hasTools() || rawToolCalls.length === 0) {
             stateManager.complete();
             exitReason = "completed";
-            return message;
+            return {
+              finalMessage: message,
+              exitReason,
+              usage: stateManager.getTotalUsage(),
+            };
           }
 
           // Parse all tool calls uniformly through the router
@@ -165,9 +186,18 @@ export const createSession = async <T extends ToolMap, M = unknown>({
           }
 
           // Hooks can call stateManager.waitForInput() to pause the session
-          await toolRouter.processToolCalls(parsedToolCalls, {
-            turn: currentTurn,
-          });
+          const toolCallResults = await toolRouter.processToolCalls(
+            parsedToolCalls,
+            {
+              turn: currentTurn,
+            }
+          );
+
+          for (const result of toolCallResults) {
+            if (result.usage) {
+              stateManager.updateUsage(result.usage);
+            }
+          }
 
           if (stateManager.getStatus() === "WAITING_FOR_INPUT") {
             const conditionMet = await condition(
@@ -195,7 +225,11 @@ export const createSession = async <T extends ToolMap, M = unknown>({
         await callSessionEnd(exitReason, stateManager.getTurns());
       }
 
-      return null;
+      return {
+        finalMessage: null,
+        exitReason,
+        usage: stateManager.getTotalUsage(),
+      };
     },
   };
 };

package/src/lib/state-manager.ts

@@ -114,6 +114,25 @@ export interface AgentStateManager<TCustom extends JsonSerializable<TCustom>> {
 
   /** Set the tools (converts Zod schemas to JSON Schema for serialization) */
   setTools(newTools: ToolDefinition[]): void;
+
+  /** Update the usage */
+  updateUsage(usage: {
+    inputTokens?: number;
+    outputTokens?: number;
+    cachedWriteTokens?: number;
+    cachedReadTokens?: number;
+    reasonTokens?: number;
+  }): void;
+
+  /** Get the total usage */
+  getTotalUsage(): {
+    totalInputTokens: number;
+    totalOutputTokens: number;
+    totalCachedWriteTokens: number;
+    totalCachedReadTokens: number;
+    totalReasonTokens: number;
+    turns: number;
+  };
 }
 
 /**
@@ -140,6 +159,11 @@ export function createAgentStateManager<
   let version = initialState?.version ?? 0;
   let turns = initialState?.turns ?? 0;
   let tools = initialState?.tools ?? [];
+  let totalInputTokens = 0;
+  let totalOutputTokens = 0;
+  let totalCachedWriteTokens = 0;
+  let totalCachedReadTokens = 0;
+  let totalReasonTokens = 0;
 
   // Tasks state
   const tasks = new Map<string, WorkflowTask>(initialState?.tasks);
@@ -283,6 +307,38 @@ export function createAgentStateManager<
       }
       return deleted;
     },
+
+    updateUsage(usage: {
+      inputTokens?: number;
+      outputTokens?: number;
+      cachedWriteTokens?: number;
+      cachedReadTokens?: number;
+      reasonTokens?: number;
+    }): void {
+      totalInputTokens += usage.inputTokens ?? 0;
+      totalOutputTokens += usage.outputTokens ?? 0;
+      totalCachedWriteTokens += usage.cachedWriteTokens ?? 0;
+      totalCachedReadTokens += usage.cachedReadTokens ?? 0;
+      totalReasonTokens += usage.reasonTokens ?? 0;
+    },
+
+    getTotalUsage(): {
+      totalInputTokens: number;
+      totalOutputTokens: number;
+      totalCachedWriteTokens: number;
+      totalCachedReadTokens: number;
+      totalReasonTokens: number;
+      turns: number;
+    } {
+      return {
+        totalInputTokens,
+        totalOutputTokens,
+        totalCachedWriteTokens,
+        totalCachedReadTokens,
+        totalReasonTokens,
+        turns,
+      };
+    },
   };
 }
 

package/src/lib/tool-router.ts

@@ -6,13 +6,14 @@ import type {
   PreToolUseHookResult,
   SubagentConfig,
   SubagentHooks,
+  TokenUsage,
   ToolHooks,
   ToolResultConfig,
 } from "./types";
 import type { SubagentArgs } from "../tools/subagent/tool";
 
 import type { z } from "zod";
-import { createSubagentTool } from "../tools/subagent/tool";
+import { createSubagentTool, SUBAGENT_TOOL_NAME } from "../tools/subagent/tool";
 import { createSubagentHandler } from "../tools/subagent/handler";
 
 export type { ToolMessageContent };
@@ -53,7 +54,7 @@ export interface ToolWithHandler<
   strict?: boolean;
   max_uses?: number;
   /** Whether this tool is available to the agent (default: true). Disabled tools are excluded from definitions and rejected at parse time. */
-  enabled?: boolean;
+  enabled?: () => boolean;
   /** Per-tool lifecycle hooks (run in addition to global hooks) */
   hooks?: ToolHooks<z.infer<TSchema>, TResult>;
 }
@@ -76,7 +77,7 @@ export type ToolMap = Record<
     handler: ToolHandler<any, any, any>;
     strict?: boolean;
     max_uses?: number;
-    enabled?: boolean;
+    enabled?: () => boolean;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     hooks?: ToolHooks<any, any>;
   }
@@ -153,6 +154,8 @@ export interface ToolHandlerResponse<TResult = null> {
    * payloads through Temporal's activity payload limit.
    */
   resultAppended?: boolean;
+  /** Token usage from the tool execution (e.g. child agent invocations) */
+  usage?: TokenUsage;
 }
 
 /**
@@ -228,6 +231,7 @@ export interface ToolCallResult<
   toolCallId: string;
   name: TName;
   data: TResult;
+  usage?: TokenUsage;
 }
 
 /**
@@ -412,25 +416,23 @@ export function createToolRouter<T extends ToolMap>(
   }
 
   /** Check if a tool is enabled (defaults to true when not specified) */
-  const isEnabled = (tool: ToolMap[string]): boolean => tool.enabled !== false;
+  const isEnabled = (tool: ToolMap[string] | SubagentConfig): boolean =>
+    tool.enabled?.() ?? true;
 
   if (options.subagents) {
-    const enabledSubagents = options.subagents.filter(
-      (s) => s.enabled !== false
-    );
-    if (enabledSubagents.length > 0) {
+    if (options.subagents.length > 0) {
       // Build per-subagent hook dispatcher keyed by subagent name
       const subagentHooksMap = new Map<string, SubagentHooks>();
-      for (const s of enabledSubagents) {
+      for (const s of options.subagents) {
         if (s.hooks) subagentHooksMap.set(s.agentName, s.hooks);
       }
 
       const resolveSubagentName = (args: unknown): string =>
         (args as SubagentArgs).subagent;
 
-      toolMap.set("Subagent", {
-        ...createSubagentTool(enabledSubagents),
-        handler: createSubagentHandler(enabledSubagents),
+      toolMap.set(SUBAGENT_TOOL_NAME, {
+        ...createSubagentTool(options.subagents),
+        handler: createSubagentHandler(options.subagents),
         ...(subagentHooksMap.size > 0 && {
           hooks: {
             onPreToolUse: async (ctx): Promise<PreToolUseHookResult> => {
@@ -655,15 +657,25 @@ export function createToolRouter<T extends ToolMap>(
     },
 
     getToolDefinitions(): ToolDefinition[] {
-      return Array.from(toolMap)
-        .filter(([, tool]) => isEnabled(tool))
-        .map(([name, tool]) => ({
-          name,
-          description: tool.description,
-          schema: tool.schema,
-          strict: tool.strict,
-          max_uses: tool.max_uses,
-        }));
+      const activeSubagents =
+        options.subagents?.filter((subagent) => isEnabled(subagent)) ?? [];
+
+      return [
+        ...Array.from(toolMap)
+          .filter(
+            ([, tool]) => isEnabled(tool) && tool.name !== SUBAGENT_TOOL_NAME
+          )
+          .map(([name, tool]) => ({
+            name,
+            description: tool.description,
+            schema: tool.schema,
+            strict: tool.strict,
+            max_uses: tool.max_uses,
+          })),
+        ...(activeSubagents.length > 0
+          ? [createSubagentTool(activeSubagents)]
+          : []),
+      ];
     },
 
     // --- Methods for processing tool calls ---

package/src/lib/types.ts

@@ -32,6 +32,10 @@ export interface BaseAgentState {
   turns: number;
   tasks: Map<string, WorkflowTask>;
   systemPrompt: string;
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  cachedWriteTokens: number;
+  cachedReadtTokens: number;
 }
 
 /**
@@ -50,17 +54,21 @@ export interface AgentFile {
   mimeType?: string;
 }
 
+export interface TokenUsage {
+  inputTokens?: number;
+  outputTokens?: number;
+  cachedWriteTokens?: number;
+  cachedReadTokens?: number;
+  reasonTokens?: number;
+}
+
 /**
  * Agent response from LLM invocation
  */
 export interface AgentResponse<M = StoredMessage> {
   message: M;
   rawToolCalls: RawToolCall[];
-  usage?: {
-    input_tokens?: number;
-    output_tokens?: number;
-    total_tokens?: number;
-  };
+  usage?: TokenUsage;
 }
 
 /**
@@ -143,9 +151,10 @@ export interface SerializableToolDefinition {
 /**
  * Configuration passed to runAgent activity
  */
-export interface RunAgentConfig {
+export interface RunAgentConfig extends AgentConfig {
+  /** The thread ID to use for the session */
   threadId: string;
-  agentName: string;
+  /** Metadata for the session */
   metadata?: Record<string, unknown>;
 }
 
@@ -155,6 +164,7 @@ export interface RunAgentConfig {
 export type RunAgentActivity<M = StoredMessage> = (
   config: RunAgentConfig
 ) => Promise<AgentResponse<M>>;
+
 /**
  * Configuration for appending a tool result
  */
@@ -186,7 +196,7 @@ export interface SubagentConfig<TResult extends z.ZodType = z.ZodType> {
   /** Description shown to the parent agent explaining what this subagent does */
   description: string;
   /** Whether this subagent is available (default: true). Disabled subagents are excluded from the Subagent tool. */
-  enabled?: boolean;
+  enabled?: () => boolean;
   /** Temporal workflow function or type name (used with executeChild) */
   workflow: string | Workflow;
   /** Optional task queue - defaults to parent's queue if not specified */

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

@@ -1,11 +1,11 @@
-import type { ToolHandler } from "../../lib/tool-router";
+import type { ActivityToolHandler } from "../../lib/tool-router";
 import type { AskUserQuestionArgs } from "./tool";
 
 /**
  * Creates handler for user interaction tool - creates AI messages for display.
  */
 export const createAskUserQuestionHandler =
-  (): ToolHandler<
+  (): ActivityToolHandler<
     AskUserQuestionArgs,
     {
       questions: {
@@ -16,7 +16,7 @@ export const createAskUserQuestionHandler =
       }[];
     }
   > =>
-  (args) => {
+  async (args) => {
     return {
       toolResponse: "Question submitted",
       data: { questions: args.questions },

package/src/tools/read-file/tool.ts

@@ -1,7 +1,7 @@
 import { z } from "zod";
 import type { ToolDefinition } from "../../lib/tool-router";
 
-export const readTool = {
+export const readFileTool = {
   name: "FileRead" as const,
   description: `Read file contents with optional pagination.
 
@@ -31,4 +31,4 @@ The tool returns the file content in an appropriate format:
   strict: true,
 } satisfies ToolDefinition;
 
-export type FileReadArgs = z.infer<typeof readTool.schema>;
+export type FileReadArgs = z.infer<typeof readFileTool.schema>;

package/src/tools/subagent/handler.ts

@@ -53,7 +53,7 @@ export function createSubagentHandler<
       taskQueue: config.taskQueue ?? parentTaskQueue,
     };
 
-    const { toolResponse, data } =
+    const { toolResponse, data, usage } =
       typeof config.workflow === "string"
         ? await executeChild(config.workflow, childOpts)
         : await executeChild(config.workflow, childOpts);
@@ -66,6 +66,7 @@ export function createSubagentHandler<
     return {
       toolResponse,
       data: validated,
+      ...(usage && { usage }),
     };
   };
 }

package/src/tools/subagent/tool.ts

@@ -1,7 +1,7 @@
 import z from "zod";
 import type { SubagentConfig } from "../../lib/types";
 
-const SUBAGENT_TOOL = "Subagent" as const;
+export const SUBAGENT_TOOL_NAME = "Subagent" as const;
 
 /**
  * Builds the tool description with available subagent information
@@ -13,13 +13,13 @@ function buildSubagentDescription(subagents: SubagentConfig[]): string {
 
   return `Launch a new agent to handle complex tasks autonomously.
 
-The ${SUBAGENT_TOOL} tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it.
+The ${SUBAGENT_TOOL_NAME} tool launches specialized agents (subprocesses) that autonomously handle complex tasks. Each agent type has specific capabilities and tools available to it.
 
 Available agent types:
 
 ${subagentList}
 
-When using the ${SUBAGENT_TOOL} tool, you must specify a subagent parameter to select which agent type to use.
+When using the ${SUBAGENT_TOOL_NAME} tool, you must specify a subagent parameter to select which agent type to use.
 
 Usage notes:
 
@@ -67,7 +67,7 @@ export function createSubagentTool<T extends SubagentConfig[]>(
   const names = subagents.map((s) => s.agentName);
 
   return {
-    name: SUBAGENT_TOOL,
+    name: SUBAGENT_TOOL_NAME,
     description: buildSubagentDescription(subagents),
     schema: z.object({
       subagent: z.enum(names).describe("The type of subagent to launch"),

package/src/tools/write-file/tool.ts

@@ -1,25 +1,24 @@
 import { z } from "zod";
 import type { ToolDefinition } from "../../lib/tool-router";
 
-export const writeTool = {
+export const writeFileTool = {
   name: "FileWrite" as const,
   description: `Create or overwrite a file with new content.
 
 Usage:
-- Provide the absolute path to the file
 - The file will be created if it doesn't exist
 - If the file exists, it will be completely overwritten
 
 IMPORTANT:
 - You must read the file first (in this session) before writing to it
 - This is an atomic write operation - the entire file is replaced
-- Path must be absolute (e.g., "/docs/readme.md", not "docs/readme.md")
+- Path must be relative to the root of the file system (e.g., "docs/readme.md", not "/docs/readme.md")
 `,
   schema: z.object({
-    file_path: z.string().describe("The absolute path to the file to write"),
+    file_path: z.string().describe("The path to the file to write"),
     content: z.string().describe("The content to write to the file"),
   }),
   strict: true,
 } satisfies ToolDefinition;
 
-export type FileWriteArgs = z.infer<typeof writeTool.schema>;
+export type FileWriteArgs = z.infer<typeof writeFileTool.schema>;

package/src/workflow.ts

@@ -113,9 +113,9 @@ export { globTool } from "./tools/glob/tool";
 export type { GlobArgs } from "./tools/glob/tool";
 export { grepTool } from "./tools/grep/tool";
 export type { GrepArgs } from "./tools/grep/tool";
-export { readTool } from "./tools/read-file/tool";
+export { readFileTool } from "./tools/read-file/tool";
 export type { FileReadArgs } from "./tools/read-file/tool";
-export { writeTool } from "./tools/write-file/tool";
+export { writeFileTool } from "./tools/write-file/tool";
 export type { FileWriteArgs } from "./tools/write-file/tool";
 export { editTool } from "./tools/edit/tool";
 export type { FileEditArgs } from "./tools/edit/tool";