feat: implement tool call tracking for API-based tools

Author: NoOne7135

agent/middleware/apiBasedTools.ts

@@ -1,7 +1,11 @@
 import { ToolMessage } from "@langchain/core/messages";
 import { createMiddleware } from "langchain";
 import { logger } from "adminforth";
-import type { ApiBasedTool } from "../../apiBasedTools.js";
+import { type ApiBasedTool } from "../../apiBasedTools.js";
+import {
+  createToolCallTracker,
+  type ToolCallEventSink,
+} from "../toolCallEvents.js";
 import { ALWAYS_AVAILABLE_API_TOOL_NAMES } from "../tools/index.js";
 import { createApiTool } from "../tools/apiTool.js";
 
@@ -67,37 +71,54 @@ export function createApiBasedToolsMiddleware(
     async wrapToolCall(request, handler) {
       const startedAt = Date.now();
       const toolInput = JSON.stringify(request.toolCall.args ?? {});
+      const { emitToolCallEvent } = request.runtime.context as {
+        emitToolCallEvent: ToolCallEventSink;
+      };
+      const toolCallTracker = createToolCallTracker({
+        emit: emitToolCallEvent,
+        toolCallId: request.toolCall.id,
+        toolName: request.toolCall.name,
+        input: (request.toolCall.args ?? {}) as Record<string, unknown>,
+        startedAt,
+      });
+      toolCallTracker.start();
       logger.info(
         `Invoking tool "${request.toolCall.name}" with input: ${toolInput}`,
       );
 
       try {
-        if (request.tool) {
-          return await handler(request);
-        }
+        let result;
 
-        const enabledApiToolNames = getEnabledApiToolNames(request.state.messages);
+        if (request.tool) {
+          result = await handler(request);
+        } else {
+          const enabledApiToolNames = getEnabledApiToolNames(request.state.messages);
 
-        if (enabledApiToolNames.has(request.toolCall.name)) {
-          return await handler({
-            ...request,
-            tool: dynamicTools[request.toolCall.name],
-          });
+          if (enabledApiToolNames.has(request.toolCall.name)) {
+            result = await handler({
+              ...request,
+              tool: dynamicTools[request.toolCall.name],
+            });
+          } else {
+            result = new ToolMessage({
+              content: `Tool "${request.toolCall.name}" is not loaded. Call fetch_tool_schema first.`,
+              tool_call_id: request.toolCall.id ?? "",
+              name: request.toolCall.name,
+              status: "error",
+            });
+          }
         }
 
-        return new ToolMessage({
-          content: `Tool "${request.toolCall.name}" is not loaded. Call fetch_tool_schema first.`,
-          tool_call_id: request.toolCall.id ?? "",
-          name: request.toolCall.name,
-          status: "error",
-        });
+        toolCallTracker.finishSuccess(result);
+        return result;
       } catch (error) {
         const errorDetails =
           error instanceof Error ? error.stack ?? error.message : String(error);
 
         logger.error(
           `Tool "${request.toolCall.name}" failed after ${Date.now() - startedAt}ms with input: ${toolInput}\n${errorDetails}`,
         );
+        toolCallTracker.finishError(error);
         throw error;
       } finally {
         logger.info(

agent/toolCallEvents.ts

@@ -1,3 +1,7 @@
+import { randomUUID } from "crypto";
+import YAML from "yaml";
+import { serializeUnknownError } from "../apiBasedTools.js";
+
 export type ToolCallEvent =
   | {
       toolCallId: string;
@@ -15,3 +19,45 @@ export type ToolCallEvent =
     };
 
 export type ToolCallEventSink = (event: ToolCallEvent) => void;
+
+export function createToolCallTracker(params: {
+  emit: ToolCallEventSink;
+  toolCallId?: string;
+  toolName: string;
+  input?: Record<string, unknown>;
+  startedAt?: number;
+}) {
+  const toolCallId = params.toolCallId ?? randomUUID();
+  const startedAt = params.startedAt ?? Date.now();
+
+  return {
+    start() {
+      params.emit({
+        toolCallId,
+        toolName: params.toolName,
+        phase: "start",
+        input: YAML.stringify(params.input ?? {}),
+      });
+    },
+    finishSuccess(output: unknown) {
+      params.emit({
+        toolCallId,
+        toolName: params.toolName,
+        phase: "end",
+        durationMs: Date.now() - startedAt,
+        output: YAML.stringify(output).trimEnd(),
+        error: null,
+      });
+    },
+    finishError(error: unknown) {
+      params.emit({
+        toolCallId,
+        toolName: params.toolName,
+        phase: "end",
+        durationMs: Date.now() - startedAt,
+        output: null,
+        error: YAML.stringify(serializeUnknownError(error)).trimEnd(),
+      });
+    },
+  };
+}

agent/tools/apiTool.ts

@@ -1,8 +1,5 @@
 import { tool } from "langchain";
-import { randomUUID } from "crypto";
-import YAML from "yaml";
 import type { ApiBasedTool } from "../../apiBasedTools.js";
-import { serializeUnknownError } from "../../apiBasedTools.js";
 
 const emptyToolSchema = {
   type: "object",
@@ -52,43 +49,11 @@ export function createApiTool(toolName: string, apiBasedTool: ApiBasedTool) {
   return tool(
     async (input, runtime) => {
       const normalizedInput = (input ?? {}) as Record<string, unknown>;
-      const toolCallId = randomUUID();
-      const startedAt = Date.now();
-      runtime.context.emitToolCallEvent({
-        toolCallId,
-        toolName,
-        phase: "start",
-        input: YAML.stringify(normalizedInput),
+      return apiBasedTool.call({
+        adminUser: runtime.context.adminUser,
+        inputs: normalizedInput,
+        userTimeZone: runtime.context.userTimeZone,
       });
-
-      try {
-        const output = await apiBasedTool.call({
-          adminUser: runtime.context.adminUser,
-          inputs: normalizedInput,
-          userTimeZone: runtime.context.userTimeZone,
-        });
-
-        runtime.context.emitToolCallEvent({
-          toolCallId,
-          toolName,
-          phase: "end",
-          durationMs: Date.now() - startedAt,
-          output,
-          error: null,
-        });
-
-        return output;
-      } catch (error) {
-        runtime.context.emitToolCallEvent({
-          toolCallId,
-          toolName,
-          phase: "end",
-          durationMs: Date.now() - startedAt,
-          output: null,
-          error: YAML.stringify(serializeUnknownError(error)),
-        });
-        throw error;
-      }
     },
     {
       name: toolName,