@witqq/agent-sdk 0.1.1 → 0.3.0

package/dist/backends/claude.js

@@ -52,7 +52,9 @@ var BaseAgent = class {
     this.state = "running";
     try {
       const messages = [{ role: "user", content: prompt }];
-      return await this.executeRun(messages, options, ac.signal);
+      const result = await this.executeRun(messages, options, ac.signal);
+      this.enrichAndNotifyUsage(result);
+      return result;
     } finally {
       this.state = "idle";
       this.abortController = null;
@@ -64,7 +66,9 @@ var BaseAgent = class {
     const ac = this.createAbortController(options?.signal);
     this.state = "running";
     try {
-      return await this.executeRun(messages, options, ac.signal);
+      const result = await this.executeRun(messages, options, ac.signal);
+      this.enrichAndNotifyUsage(result);
+      return result;
     } finally {
       this.state = "idle";
       this.abortController = null;
@@ -77,12 +81,14 @@ var BaseAgent = class {
     this.state = "running";
     try {
       const messages = [{ role: "user", content: prompt }];
-      return await this.executeRunStructured(
+      const result = await this.executeRunStructured(
         messages,
         schema,
         options,
         ac.signal
       );
+      this.enrichAndNotifyUsage(result);
+      return result;
     } finally {
       this.state = "idle";
       this.abortController = null;
@@ -95,7 +101,21 @@ var BaseAgent = class {
     this.state = "streaming";
     try {
       const messages = [{ role: "user", content: prompt }];
-      yield* this.executeStream(messages, options, ac.signal);
+      const enriched = this.enrichStream(this.executeStream(messages, options, ac.signal));
+      yield* this.heartbeatStream(enriched);
+    } finally {
+      this.state = "idle";
+      this.abortController = null;
+    }
+  }
+  async *streamWithContext(messages, options) {
+    this.guardReentrancy();
+    this.guardDisposed();
+    const ac = this.createAbortController(options?.signal);
+    this.state = "streaming";
+    try {
+      const enriched = this.enrichStream(this.executeStream(messages, options, ac.signal));
+      yield* this.heartbeatStream(enriched);
     } finally {
       this.state = "idle";
       this.abortController = null;
@@ -117,6 +137,95 @@ var BaseAgent = class {
     this.abort();
     this.state = "disposed";
   }
+  // ─── Usage Enrichment ───────────────────────────────────────────
+  /** Enrich result usage with model/backend and fire onUsage callback */
+  enrichAndNotifyUsage(result) {
+    if (result.usage) {
+      result.usage = {
+        ...result.usage,
+        model: this.config.model,
+        backend: this.backendName
+      };
+      this.callOnUsage(result.usage);
+    }
+  }
+  /** Wrap a stream to enrich usage_update events and fire onUsage callback */
+  async *enrichStream(source) {
+    for await (const event of source) {
+      if (event.type === "usage_update") {
+        const usage = {
+          promptTokens: event.promptTokens,
+          completionTokens: event.completionTokens,
+          model: this.config.model,
+          backend: this.backendName
+        };
+        this.callOnUsage(usage);
+        yield { type: "usage_update", ...usage };
+      } else {
+        yield event;
+      }
+    }
+  }
+  /** Fire onUsage callback (fire-and-forget: errors logged, not propagated) */
+  callOnUsage(usage) {
+    if (!this.config.onUsage) return;
+    try {
+      this.config.onUsage(usage);
+    } catch (e) {
+      console.warn(
+        "[agent-sdk] onUsage callback error:",
+        e instanceof Error ? e.message : String(e)
+      );
+    }
+  }
+  // ─── Heartbeat ───────────────────────────────────────────────
+  /** Wrap a stream to emit heartbeat events at configured intervals.
+   *  When heartbeatInterval is not set, passes through directly. */
+  async *heartbeatStream(source) {
+    const interval = this.config.heartbeatInterval;
+    if (!interval || interval <= 0) {
+      yield* source;
+      return;
+    }
+    const iterator = source[Symbol.asyncIterator]();
+    let pendingEvent = null;
+    let heartbeatResolve = null;
+    const timer = setInterval(() => {
+      if (heartbeatResolve) {
+        const resolve = heartbeatResolve;
+        heartbeatResolve = null;
+        resolve();
+      }
+    }, interval);
+    try {
+      while (true) {
+        if (!pendingEvent) {
+          pendingEvent = iterator.next();
+        }
+        const heartbeatPromise = new Promise((resolve) => {
+          heartbeatResolve = resolve;
+        });
+        const eventDone = pendingEvent.then(
+          (r) => ({ kind: "event", result: r })
+        );
+        const heartbeatDone = heartbeatPromise.then(
+          () => ({ kind: "heartbeat" })
+        );
+        const winner = await Promise.race([eventDone, heartbeatDone]);
+        if (winner.kind === "heartbeat") {
+          yield { type: "heartbeat" };
+        } else {
+          pendingEvent = null;
+          heartbeatResolve = null;
+          if (winner.result.done) break;
+          yield winner.result.value;
+        }
+      }
+    } finally {
+      clearInterval(timer);
+      heartbeatResolve = null;
+    }
+  }
   // ─── Guards ───────────────────────────────────────────────────
   guardReentrancy() {
     if (this.state === "running" || this.state === "streaming") {
@@ -154,6 +263,9 @@ var BaseAgent = class {
 // src/utils/schema.ts
 function zodToJsonSchema(schema) {
   const schemaAny = schema;
+  if ("toJSONSchema" in schema && typeof schemaAny.toJSONSchema === "function") {
+    return schemaAny.toJSONSchema();
+  }
   if ("jsonSchema" in schema && typeof schemaAny.jsonSchema === "function") {
     return schemaAny.jsonSchema();
   }
@@ -342,7 +454,31 @@ function aggregateUsage(modelUsage) {
   }
   return { promptTokens, completionTokens };
 }
-function mapSDKMessage(msg, thinkingBlockIndices) {
+var ClaudeToolCallTracker = class {
+  queues = /* @__PURE__ */ new Map();
+  trackStart(toolCallId, toolName) {
+    if (!this.queues.has(toolName)) {
+      this.queues.set(toolName, []);
+    }
+    this.queues.get(toolName).push(toolCallId);
+  }
+  /** Peek at the current tool call ID for a tool name (does not consume) */
+  peekToolCallId(toolName) {
+    const queue = this.queues.get(toolName);
+    if (!queue || queue.length === 0) return "";
+    return queue[0];
+  }
+  /** Consume and return the first tool call ID for a tool name */
+  consumeToolCallId(toolName) {
+    const queue = this.queues.get(toolName);
+    if (!queue || queue.length === 0) return "";
+    return queue.shift();
+  }
+  clear() {
+    this.queues.clear();
+  }
+};
+function mapSDKMessage(msg, thinkingBlockIndices, toolCallTracker) {
   switch (msg.type) {
     case "assistant": {
       const betaMessage = msg.message;
@@ -354,9 +490,15 @@ function mapSDKMessage(msg, thinkingBlockIndices) {
       }
       for (const block of betaMessage.content) {
         if (block.type === "tool_use") {
+          const toolCallId = String(block.id ?? "");
+          const toolName = block.name ?? "unknown";
+          if (toolCallTracker) {
+            toolCallTracker.trackStart(toolCallId, toolName);
+          }
           events.push({
             type: "tool_call_start",
-            toolName: block.name ?? "unknown",
+            toolCallId,
+            toolName,
             args: block.input ?? {}
           });
         }
@@ -373,9 +515,18 @@ function mapSDKMessage(msg, thinkingBlockIndices) {
     case "tool_use_summary": {
       const summary = msg.summary;
       const toolName = msg.tool_name ?? "unknown";
+      const precedingIds = msg.preceding_tool_use_ids;
+      let toolCallId = "";
+      if (precedingIds && precedingIds.length > 0) {
+        toolCallId = precedingIds[0];
+        if (toolCallTracker) toolCallTracker.consumeToolCallId(toolName);
+      } else if (toolCallTracker) {
+        toolCallId = toolCallTracker.consumeToolCallId(toolName);
+      }
       if (summary) {
         return {
           type: "tool_call_end",
+          toolCallId,
           toolName,
           result: summary
         };
@@ -402,7 +553,9 @@ function mapSDKMessage(msg, thinkingBlockIndices) {
     }
     case "tool_progress": {
       const toolName = msg.tool_name;
-      return toolName ? { type: "tool_call_start", toolName, args: {} } : null;
+      if (!toolName) return null;
+      const toolCallId = toolCallTracker?.peekToolCallId(toolName) ?? "";
+      return { type: "tool_call_start", toolCallId, toolName, args: {} };
     }
     case "result": {
       if (msg.subtype === "success") {
@@ -427,6 +580,7 @@ function mapSDKMessage(msg, thinkingBlockIndices) {
   }
 }
 var ClaudeAgent = class extends BaseAgent {
+  backendName = "claude";
   options;
   tools;
   canUseTool;
@@ -607,10 +761,11 @@ var ClaudeAgent = class extends BaseAgent {
     opts = await this.buildMcpConfig(opts);
     const q = sdk.query({ prompt, options: opts });
     const thinkingBlockIndices = /* @__PURE__ */ new Set();
+    const toolCallTracker = new ClaudeToolCallTracker();
     try {
       for await (const msg of q) {
         if (signal.aborted) throw new AbortError();
-        const event = mapSDKMessage(msg, thinkingBlockIndices);
+        const event = mapSDKMessage(msg, thinkingBlockIndices, toolCallTracker);
         if (event) {
           if (Array.isArray(event)) {
             for (const e of event) yield e;