@tanstack/ai 0.6.3 → 0.8.0

package/dist/esm/activities/chat/index.js

@@ -1,9 +1,10 @@
-import { aiEventClient } from "../../event-client.js";
+import { devtoolsMiddleware } from "@tanstack/ai-event-client";
 import { streamToText } from "../../stream-to-response.js";
-import { ToolCallManager, executeToolCalls } from "./tools/tool-calls.js";
+import { ToolCallManager, MiddlewareAbortError, executeToolCalls } from "./tools/tool-calls.js";
 import { convertSchemaToJsonSchema, isStandardSchema, parseWithStandardSchema } from "./tools/schema-converter.js";
 import { maxIterations } from "./agent-loop-strategies.js";
 import { convertMessagesToModelMessages } from "./messages.js";
+import { MiddlewareRunner } from "./middleware/compose.js";
 const kind = "text";
 function createChatOptions(options) {
   return options;
@@ -17,10 +18,11 @@ class TextEngine {
     this.currentMessageId = null;
     this.accumulatedContent = "";
     this.finishedEvent = null;
-    this.shouldEmitStreamEnd = true;
     this.earlyTermination = false;
     this.toolPhase = "continue";
     this.cyclePhase = "processText";
+    this.deferredPromises = [];
+    this.terminalHookCalled = false;
     this.adapter = config.adapter;
     this.params = config.params;
     this.systemPrompts = config.params.systemPrompts || [];
@@ -40,6 +42,45 @@ class TextEngine {
     this.streamId = this.createId("stream");
     this.effectiveRequest = config.params.abortController ? { signal: config.params.abortController.signal } : void 0;
     this.effectiveSignal = config.params.abortController?.signal;
+    const allMiddleware = [devtoolsMiddleware(), ...config.middleware || []];
+    this.middlewareRunner = new MiddlewareRunner(allMiddleware);
+    this.middlewareAbortController = new AbortController();
+    this.middlewareCtx = {
+      requestId: this.requestId,
+      streamId: this.streamId,
+      conversationId: config.params.conversationId,
+      phase: "init",
+      iteration: 0,
+      chunkIndex: 0,
+      signal: this.effectiveSignal,
+      abort: (reason) => {
+        this.abortReason = reason;
+        this.middlewareAbortController?.abort(reason);
+      },
+      context: config.context,
+      defer: (promise) => {
+        this.deferredPromises.push(promise);
+      },
+      // Provider / adapter info
+      provider: config.adapter.name,
+      model: config.params.model,
+      source: "server",
+      streaming: true,
+      // Config-derived (updated in beforeRun and applyMiddlewareConfig)
+      systemPrompts: this.systemPrompts,
+      toolNames: void 0,
+      options: void 0,
+      modelOptions: config.params.modelOptions,
+      // Computed
+      messageCount: this.initialMessageCount,
+      hasTools: this.tools.length > 0,
+      // Mutable per-iteration
+      currentMessageId: null,
+      accumulatedContent: "",
+      // References
+      messages: this.messages,
+      createId: (prefix) => this.createId(prefix)
+    };
   }
   /** Get the accumulated content after the chat loop completes */
   getAccumulatedContent() {
@@ -52,24 +93,77 @@ class TextEngine {
   async *run() {
     this.beforeRun();
     try {
+      this.middlewareCtx.phase = "init";
+      const initialConfig = this.buildMiddlewareConfig();
+      const transformedConfig = await this.middlewareRunner.runOnConfig(
+        this.middlewareCtx,
+        initialConfig
+      );
+      this.applyMiddlewareConfig(transformedConfig);
+      await this.middlewareRunner.runOnStart(this.middlewareCtx);
       const pendingPhase = yield* this.checkForPendingToolCalls();
       if (pendingPhase === "wait") {
         return;
       }
       do {
-        if (this.earlyTermination || this.isAborted()) {
+        if (this.earlyTermination || this.isCancelled()) {
           return;
         }
-        this.beginCycle();
+        await this.beginCycle();
         if (this.cyclePhase === "processText") {
+          this.middlewareCtx.phase = "beforeModel";
+          this.middlewareCtx.iteration = this.iterationCount;
+          const iterConfig = this.buildMiddlewareConfig();
+          const transformedConfig2 = await this.middlewareRunner.runOnConfig(
+            this.middlewareCtx,
+            iterConfig
+          );
+          this.applyMiddlewareConfig(transformedConfig2);
           yield* this.streamModelResponse();
         } else {
           yield* this.processToolCalls();
         }
         this.endCycle();
       } while (this.shouldContinue());
+      if (!this.terminalHookCalled && this.toolPhase !== "wait") {
+        this.terminalHookCalled = true;
+        await this.middlewareRunner.runOnFinish(this.middlewareCtx, {
+          finishReason: this.lastFinishReason,
+          duration: Date.now() - this.streamStartTime,
+          content: this.accumulatedContent,
+          usage: this.finishedEvent?.usage
+        });
+      }
+    } catch (error) {
+      if (!this.terminalHookCalled) {
+        this.terminalHookCalled = true;
+        if (error instanceof MiddlewareAbortError) {
+          this.abortReason = error.message;
+          await this.middlewareRunner.runOnAbort(this.middlewareCtx, {
+            reason: error.message,
+            duration: Date.now() - this.streamStartTime
+          });
+        } else {
+          await this.middlewareRunner.runOnError(this.middlewareCtx, {
+            error,
+            duration: Date.now() - this.streamStartTime
+          });
+        }
+      }
+      if (!(error instanceof MiddlewareAbortError)) {
+        throw error;
+      }
     } finally {
-      this.afterRun();
+      if (!this.terminalHookCalled && this.isCancelled()) {
+        this.terminalHookCalled = true;
+        await this.middlewareRunner.runOnAbort(this.middlewareCtx, {
+          reason: this.abortReason,
+          duration: Date.now() - this.streamStartTime
+        });
+      }
+      if (this.deferredPromises.length > 0) {
+        await Promise.allSettled(this.deferredPromises);
+      }
     }
   }
   beforeRun() {
@@ -82,55 +176,12 @@ class TextEngine {
     if (metadata !== void 0) options.metadata = metadata;
     this.eventOptions = Object.keys(options).length > 0 ? options : void 0;
     this.eventToolNames = tools?.map((t) => t.name);
-    aiEventClient.emit("text:request:started", {
-      ...this.buildTextEventContext(),
-      timestamp: Date.now()
-    });
-    const messagesToEmit = this.params.conversationId ? this.messages.slice(-1).filter((m) => m.role === "user") : this.messages;
-    messagesToEmit.forEach((message, index) => {
-      const messageIndex = this.params.conversationId ? this.messages.length - 1 : index;
-      const messageId = this.createId("msg");
-      const baseContext = this.buildTextEventContext();
-      const content = this.getContentString(message.content);
-      aiEventClient.emit("text:message:created", {
-        ...baseContext,
-        messageId,
-        role: message.role,
-        content,
-        toolCalls: message.toolCalls,
-        messageIndex,
-        timestamp: Date.now()
-      });
-      if (message.role === "user") {
-        aiEventClient.emit("text:message:user", {
-          ...baseContext,
-          messageId,
-          role: "user",
-          content,
-          messageIndex,
-          timestamp: Date.now()
-        });
-      }
-    });
+    this.middlewareCtx.options = this.eventOptions;
+    this.middlewareCtx.toolNames = this.eventToolNames;
   }
-  afterRun() {
-    if (!this.shouldEmitStreamEnd) {
-      return;
-    }
-    const now = Date.now();
-    aiEventClient.emit("text:request:completed", {
-      ...this.buildTextEventContext(),
-      content: this.accumulatedContent,
-      messageId: this.currentMessageId || void 0,
-      finishReason: this.lastFinishReason || void 0,
-      usage: this.finishedEvent?.usage,
-      duration: now - this.streamStartTime,
-      timestamp: now
-    });
-  }
-  beginCycle() {
+  async beginCycle() {
     if (this.cyclePhase === "processText") {
-      this.beginIteration();
+      await this.beginIteration();
     }
   }
   endCycle() {
@@ -141,27 +192,26 @@ class TextEngine {
     this.cyclePhase = "processText";
     this.iterationCount++;
   }
-  beginIteration() {
+  async beginIteration() {
     this.currentMessageId = this.createId("msg");
     this.accumulatedContent = "";
     this.finishedEvent = null;
-    const baseContext = this.buildTextEventContext();
-    aiEventClient.emit("text:message:created", {
-      ...baseContext,
-      messageId: this.currentMessageId,
-      role: "assistant",
-      content: "",
-      timestamp: Date.now()
+    this.middlewareCtx.currentMessageId = this.currentMessageId;
+    this.middlewareCtx.accumulatedContent = "";
+    await this.middlewareRunner.runOnIteration(this.middlewareCtx, {
+      iteration: this.iterationCount,
+      messageId: this.currentMessageId
     });
   }
   async *streamModelResponse() {
     const { temperature, topP, maxTokens, metadata, modelOptions } = this.params;
-    const tools = this.params.tools;
-    const toolsWithJsonSchemas = tools?.map((tool) => ({
+    const tools = this.tools;
+    const toolsWithJsonSchemas = tools.map((tool) => ({
       ...tool,
       inputSchema: tool.inputSchema ? convertSchemaToJsonSchema(tool.inputSchema) : void 0,
       outputSchema: tool.outputSchema ? convertSchemaToJsonSchema(tool.outputSchema) : void 0
     }));
+    this.middlewareCtx.phase = "modelStream";
     for await (const chunk of this.adapter.chatStream({
       model: this.params.model,
       messages: this.messages,
@@ -174,12 +224,22 @@ class TextEngine {
       modelOptions,
       systemPrompts: this.systemPrompts
     })) {
-      if (this.isAborted()) {
+      if (this.isCancelled()) {
         break;
       }
       this.totalChunkCount++;
-      yield chunk;
-      this.handleStreamChunk(chunk);
+      const outputChunks = await this.middlewareRunner.runOnChunk(
+        this.middlewareCtx,
+        chunk
+      );
+      for (const outputChunk of outputChunks) {
+        yield outputChunk;
+        this.handleStreamChunk(outputChunk);
+        this.middlewareCtx.chunkIndex++;
+      }
+      if (chunk.type === "RUN_FINISHED" && chunk.usage) {
+        await this.middlewareRunner.runOnUsage(this.middlewareCtx, chunk.usage);
+      }
       if (this.earlyTermination) {
         break;
       }
@@ -220,87 +280,25 @@ class TextEngine {
     } else {
       this.accumulatedContent += chunk.delta;
     }
-    aiEventClient.emit("text:chunk:content", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      content: this.accumulatedContent,
-      delta: chunk.delta,
-      timestamp: Date.now()
-    });
+    this.middlewareCtx.accumulatedContent = this.accumulatedContent;
   }
   handleToolCallStartEvent(chunk) {
     this.toolCallManager.addToolCallStartEvent(chunk);
-    aiEventClient.emit("text:chunk:tool-call", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      toolCallId: chunk.toolCallId,
-      toolName: chunk.toolName,
-      index: chunk.index ?? 0,
-      arguments: "",
-      timestamp: Date.now()
-    });
   }
   handleToolCallArgsEvent(chunk) {
     this.toolCallManager.addToolCallArgsEvent(chunk);
-    aiEventClient.emit("text:chunk:tool-call", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      toolCallId: chunk.toolCallId,
-      toolName: "",
-      index: 0,
-      arguments: chunk.delta,
-      timestamp: Date.now()
-    });
   }
   handleToolCallEndEvent(chunk) {
     this.toolCallManager.completeToolCall(chunk);
-    aiEventClient.emit("text:chunk:tool-result", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      toolCallId: chunk.toolCallId,
-      result: chunk.result || "",
-      timestamp: Date.now()
-    });
   }
   handleRunFinishedEvent(chunk) {
-    aiEventClient.emit("text:chunk:done", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      finishReason: chunk.finishReason,
-      usage: chunk.usage,
-      timestamp: Date.now()
-    });
-    if (chunk.usage) {
-      aiEventClient.emit("text:usage", {
-        ...this.buildTextEventContext(),
-        messageId: this.currentMessageId || void 0,
-        usage: chunk.usage,
-        timestamp: Date.now()
-      });
-    }
     this.finishedEvent = chunk;
     this.lastFinishReason = chunk.finishReason;
   }
-  handleRunErrorEvent(chunk) {
-    aiEventClient.emit("text:chunk:error", {
-      ...this.buildTextEventContext(),
-      messageId: this.currentMessageId || void 0,
-      error: chunk.error.message,
-      timestamp: Date.now()
-    });
+  handleRunErrorEvent(_chunk) {
     this.earlyTermination = true;
-    this.shouldEmitStreamEnd = false;
-  }
-  handleStepFinishedEvent(chunk) {
-    if (chunk.content || chunk.delta) {
-      aiEventClient.emit("text:chunk:thinking", {
-        ...this.buildTextEventContext(),
-        messageId: this.currentMessageId || void 0,
-        content: chunk.content || "",
-        delta: chunk.delta,
-        timestamp: Date.now()
-      });
-    }
+  }
+  handleStepFinishedEvent(_chunk) {
   }
   async *checkForPendingToolCalls() {
     const pendingToolCalls = this.getPendingToolCallsFromMessages();
@@ -314,34 +312,65 @@ class TextEngine {
       this.tools,
       approvals,
       clientToolResults,
-      (eventName, data) => this.createCustomEventChunk(eventName, data)
+      (eventName, data) => this.createCustomEventChunk(eventName, data),
+      {
+        onBeforeToolCall: async (toolCall, tool, args) => {
+          const hookCtx = {
+            toolCall,
+            tool,
+            args,
+            toolName: toolCall.function.name,
+            toolCallId: toolCall.id
+          };
+          return this.middlewareRunner.runOnBeforeToolCall(
+            this.middlewareCtx,
+            hookCtx
+          );
+        },
+        onAfterToolCall: async (info) => {
+          await this.middlewareRunner.runOnAfterToolCall(
+            this.middlewareCtx,
+            info
+          );
+        }
+      }
     );
     const executionResult = yield* this.drainToolCallGenerator(generator);
+    if (this.isMiddlewareAborted()) {
+      this.setToolPhase("stop");
+      return "stop";
+    }
+    await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, {
+      toolCalls: pendingToolCalls,
+      results: executionResult.results,
+      needsApproval: executionResult.needsApproval,
+      needsClientExecution: executionResult.needsClientExecution
+    });
     if (executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0) {
       if (executionResult.results.length > 0) {
-        for (const chunk of this.emitToolResults(
+        for (const chunk of this.buildToolResultChunks(
           executionResult.results,
           finishEvent
         )) {
           yield chunk;
         }
       }
-      for (const chunk of this.emitApprovalRequests(
+      for (const chunk of this.buildApprovalChunks(
         executionResult.needsApproval,
         finishEvent
       )) {
         yield chunk;
       }
-      for (const chunk of this.emitClientToolInputs(
+      for (const chunk of this.buildClientToolChunks(
         executionResult.needsClientExecution,
         finishEvent
       )) {
         yield chunk;
       }
-      this.shouldEmitStreamEnd = false;
+      this.setToolPhase("wait");
       return "wait";
     }
-    const toolResultChunks = this.emitToolResults(
+    const toolResultChunks = this.buildToolResultChunks(
       executionResult.results,
       finishEvent
     );
@@ -362,31 +391,64 @@ class TextEngine {
       return;
     }
     this.addAssistantToolCallMessage(toolCalls);
+    this.middlewareCtx.phase = "beforeTools";
     const { approvals, clientToolResults } = this.collectClientState();
     const generator = executeToolCalls(
       toolCalls,
       this.tools,
       approvals,
       clientToolResults,
-      (eventName, data) => this.createCustomEventChunk(eventName, data)
+      (eventName, data) => this.createCustomEventChunk(eventName, data),
+      {
+        onBeforeToolCall: async (toolCall, tool, args) => {
+          const hookCtx = {
+            toolCall,
+            tool,
+            args,
+            toolName: toolCall.function.name,
+            toolCallId: toolCall.id
+          };
+          return this.middlewareRunner.runOnBeforeToolCall(
+            this.middlewareCtx,
+            hookCtx
+          );
+        },
+        onAfterToolCall: async (info) => {
+          await this.middlewareRunner.runOnAfterToolCall(
+            this.middlewareCtx,
+            info
+          );
+        }
+      }
     );
     const executionResult = yield* this.drainToolCallGenerator(generator);
+    this.middlewareCtx.phase = "afterTools";
+    if (this.isMiddlewareAborted()) {
+      this.setToolPhase("stop");
+      return;
+    }
+    await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, {
+      toolCalls,
+      results: executionResult.results,
+      needsApproval: executionResult.needsApproval,
+      needsClientExecution: executionResult.needsClientExecution
+    });
     if (executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0) {
       if (executionResult.results.length > 0) {
-        for (const chunk of this.emitToolResults(
+        for (const chunk of this.buildToolResultChunks(
           executionResult.results,
           finishEvent
         )) {
           yield chunk;
         }
       }
-      for (const chunk of this.emitApprovalRequests(
+      for (const chunk of this.buildApprovalChunks(
         executionResult.needsApproval,
         finishEvent
       )) {
         yield chunk;
       }
-      for (const chunk of this.emitClientToolInputs(
+      for (const chunk of this.buildClientToolChunks(
         executionResult.needsClientExecution,
         finishEvent
       )) {
@@ -395,7 +457,7 @@ class TextEngine {
       this.setToolPhase("wait");
       return;
     }
-    const toolResultChunks = this.emitToolResults(
+    const toolResultChunks = this.buildToolResultChunks(
       executionResult.results,
       finishEvent
     );
@@ -409,7 +471,6 @@ class TextEngine {
     return this.finishedEvent?.finishReason === "tool_calls" && this.tools.length > 0 && this.toolCallManager.hasToolCalls();
   }
   addAssistantToolCallMessage(toolCalls) {
-    const messageId = this.currentMessageId ?? this.createId("msg");
     this.messages = [
       ...this.messages,
       {
@@ -418,14 +479,6 @@ class TextEngine {
         toolCalls
       }
     ];
-    aiEventClient.emit("text:message:created", {
-      ...this.buildTextEventContext(),
-      messageId,
-      role: "assistant",
-      content: this.accumulatedContent || "",
-      toolCalls,
-      timestamp: Date.now()
-    });
   }
   /**
    * Extract client state (approvals and client tool results) from original messages.
@@ -470,18 +523,9 @@ class TextEngine {
     }
     return { approvals, clientToolResults };
   }
-  emitApprovalRequests(approvals, finishEvent) {
+  buildApprovalChunks(approvals, finishEvent) {
     const chunks = [];
     for (const approval of approvals) {
-      aiEventClient.emit("tools:approval:requested", {
-        ...this.buildTextEventContext(),
-        messageId: this.currentMessageId || void 0,
-        toolCallId: approval.toolCallId,
-        toolName: approval.toolName,
-        input: approval.input,
-        approvalId: approval.approvalId,
-        timestamp: Date.now()
-      });
       chunks.push({
         type: "CUSTOM",
         timestamp: Date.now(),
@@ -500,17 +544,9 @@ class TextEngine {
     }
     return chunks;
   }
-  emitClientToolInputs(clientRequests, finishEvent) {
+  buildClientToolChunks(clientRequests, finishEvent) {
     const chunks = [];
     for (const clientTool of clientRequests) {
-      aiEventClient.emit("tools:input:available", {
-        ...this.buildTextEventContext(),
-        messageId: this.currentMessageId || void 0,
-        toolCallId: clientTool.toolCallId,
-        toolName: clientTool.toolName,
-        input: clientTool.input,
-        timestamp: Date.now()
-      });
       chunks.push({
         type: "CUSTOM",
         timestamp: Date.now(),
@@ -525,18 +561,9 @@ class TextEngine {
     }
     return chunks;
   }
-  emitToolResults(results, finishEvent) {
+  buildToolResultChunks(results, finishEvent) {
     const chunks = [];
     for (const result of results) {
-      aiEventClient.emit("tools:call:completed", {
-        ...this.buildTextEventContext(),
-        messageId: this.currentMessageId || void 0,
-        toolCallId: result.toolCallId,
-        toolName: result.toolName,
-        result: result.result,
-        duration: result.duration ?? 0,
-        timestamp: Date.now()
-      });
       const content = JSON.stringify(result.result);
       chunks.push({
         type: "TOOL_CALL_END",
@@ -554,13 +581,6 @@ class TextEngine {
           toolCallId: result.toolCallId
         }
       ];
-      aiEventClient.emit("text:message:created", {
-        ...this.buildTextEventContext(),
-        messageId: this.createId("msg"),
-        role: "tool",
-        content,
-        timestamp: Date.now()
-      });
     }
     return chunks;
   }
@@ -617,33 +637,44 @@ class TextEngine {
   isAborted() {
     return !!this.effectiveSignal?.aborted;
   }
-  buildTextEventContext() {
+  isMiddlewareAborted() {
+    return !!this.middlewareAbortController?.signal.aborted;
+  }
+  isCancelled() {
+    return this.isAborted() || this.isMiddlewareAborted();
+  }
+  buildMiddlewareConfig() {
     return {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      provider: this.adapter.name,
-      model: this.params.model,
-      clientId: this.params.conversationId,
-      source: "server",
-      systemPrompts: this.systemPrompts.length > 0 ? this.systemPrompts : void 0,
-      toolNames: this.eventToolNames,
-      options: this.eventOptions,
-      modelOptions: this.params.modelOptions,
-      messageCount: this.initialMessageCount,
-      hasTools: this.tools.length > 0,
-      streaming: true
+      messages: this.messages,
+      systemPrompts: [...this.systemPrompts],
+      tools: [...this.tools],
+      temperature: this.params.temperature,
+      topP: this.params.topP,
+      maxTokens: this.params.maxTokens,
+      metadata: this.params.metadata,
+      modelOptions: this.params.modelOptions
     };
   }
-  getContentString(content) {
-    if (typeof content === "string") return content;
-    const text = content?.map((part) => part.type === "text" ? part.content : "").join("") || "";
-    return text;
+  applyMiddlewareConfig(config) {
+    this.messages = config.messages;
+    this.systemPrompts = config.systemPrompts;
+    this.tools = config.tools;
+    this.params = {
+      ...this.params,
+      temperature: config.temperature,
+      topP: config.topP,
+      maxTokens: config.maxTokens,
+      metadata: config.metadata,
+      modelOptions: config.modelOptions
+    };
+    this.middlewareCtx.messages = this.messages;
+    this.middlewareCtx.systemPrompts = this.systemPrompts;
+    this.middlewareCtx.hasTools = this.tools.length > 0;
+    this.middlewareCtx.toolNames = this.tools.map((t) => t.name);
+    this.middlewareCtx.modelOptions = config.modelOptions;
   }
   setToolPhase(phase) {
     this.toolPhase = phase;
-    if (phase === "wait") {
-      this.shouldEmitStreamEnd = false;
-    }
   }
   /**
    * Drain an executeToolCalls async generator, yielding any CustomEvent chunks
@@ -687,11 +718,13 @@ function chat(options) {
   );
 }
 async function* runStreamingText(options) {
-  const { adapter, ...textOptions } = options;
+  const { adapter, middleware, context, ...textOptions } = options;
   const model = adapter.model;
   const engine = new TextEngine({
     adapter,
-    params: { ...textOptions, model }
+    params: { ...textOptions, model },
+    middleware,
+    context
   });
   for await (const chunk of engine.run()) {
     yield chunk;
@@ -704,14 +737,16 @@ function runNonStreamingText(options) {
   return streamToText(stream);
 }
 async function runAgenticStructuredOutput(options) {
-  const { adapter, outputSchema, ...textOptions } = options;
+  const { adapter, outputSchema, middleware, context, ...textOptions } = options;
   const model = adapter.model;
   if (!outputSchema) {
     throw new Error("outputSchema is required for structured output");
   }
   const engine = new TextEngine({
     adapter,
-    params: { ...textOptions, model }
+    params: { ...textOptions, model },
+    middleware,
+    context
   });
   for await (const _chunk of engine.run()) {
   }