Merge pull request #59 from patlux/fix/tool-call-streaming

fix(stream): forward incremental tool-call arguments

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Stream incremental tool-call arguments from the `/alpha/generate` transport instead of waiting for the final complete tool-call event.
 - Add a daily GitHub Actions synchronization job that opens or updates a pull request for CLI version, image capability, reasoning, effort, and output-limit changes in the latest published Command Code catalog.
 - Refresh static model capabilities from `command-code@1.32.2`, separating reasoning support from selectable effort levels and honoring model-specific output limits.
 - Reject truncated, aborted, and network-failed generate streams instead of reporting partial responses as successful.

README.md

@@ -7,15 +7,6 @@ A custom provider for [pi](https://github.com/earendil-works/pi) that connects t
 
 > **Disclaimer:** This is an unofficial, community-maintained integration. It is not affiliated with, endorsed by, or supported by Command Code. You need your own Command Code account, API key, and a plan with Provider API access. Command Code's terms, availability, and pricing apply.
 
-The extension uses one provider and automatically selects the transport supported by the authenticated account:
-
-- `GET /provider/v1/models` for model discovery
-- `POST /provider/v1/chat/completions` for non-Claude models with Provider API access
-- `POST /provider/v1/messages` for Claude models with Provider API access
-- `/alpha/generate` after the Provider API explicitly returns `403 upgrade_required`, which currently identifies Go-plan accounts
-
-The detected transport is remembered only for the running process and is re-evaluated when the credential changes. Other authentication, permission, rate-limit, network, and server errors never trigger the fallback.
-
 ## Install
 
 ```sh

src/core.ts

@@ -288,6 +288,10 @@ export function createStreamCommandCode(deps: CoreDependencies) {
       let textBlock: TextContent | undefined
       let currentTextIdx = -1
       let thinkingIdx = -1
+      const streamingToolCalls = new Map<
+        string,
+        { contentIndex: number; toolCall: ToolCallContent; partialArgs: string }
+      >()
       let finished = false
 
       const abortUpstream = () => {
@@ -400,25 +404,81 @@ export function createStreamCommandCode(deps: CoreDependencies) {
             break
           }
 
+          case "tool-input-start": {
+            endTextBlock()
+            endThinking()
+            const id = stringValue(event.id)
+            if (!id || streamingToolCalls.has(id)) break
+
+            const toolCall: ToolCallContent = {
+              type: "toolCall",
+              id,
+              name: stringValue(event.toolName) ?? "",
+              arguments: {},
+            }
+            output.content.push(toolCall)
+            const contentIndex = output.content.length - 1
+            streamingToolCalls.set(id, { contentIndex, toolCall, partialArgs: "" })
+            stream.push({
+              type: "toolcall_start",
+              contentIndex,
+              partial: output,
+            })
+            break
+          }
+
+          case "tool-input-delta": {
+            const id = stringValue(event.id)
+            const delta = stringValue(event.delta)
+            if (!id || delta === undefined) break
+            const active = streamingToolCalls.get(id)
+            if (!active) break
+
+            active.partialArgs += delta
+            active.toolCall.arguments = recordOrEmpty(active.partialArgs)
+            stream.push({
+              type: "toolcall_delta",
+              contentIndex: active.contentIndex,
+              delta,
+              partial: output,
+            })
+            break
+          }
+
+          case "tool-input-end": {
+            break
+          }
+
           case "tool-call": {
             endTextBlock()
             endThinking()
-            const toolCall: ToolCallContent = {
+            const id = stringValue(event.toolCallId) ?? ""
+            const active = streamingToolCalls.get(id)
+            const toolCall: ToolCallContent = active?.toolCall ?? {
               type: "toolCall",
-              id: stringValue(event.toolCallId) ?? "",
+              id,
               name: stringValue(event.toolName) ?? "",
-              arguments: recordOrEmpty(event.input ?? event.args ?? event.arguments),
+              arguments: {},
+            }
+            toolCall.name = stringValue(event.toolName) ?? toolCall.name
+            toolCall.arguments = recordOrEmpty(event.input ?? event.args ?? event.arguments)
+
+            let contentIndex: number
+            if (active) {
+              contentIndex = active.contentIndex
+              streamingToolCalls.delete(id)
+            } else {
+              output.content.push(toolCall)
+              contentIndex = output.content.length - 1
+              stream.push({
+                type: "toolcall_start",
+                contentIndex,
+                partial: output,
+              })
             }
-            output.content.push(toolCall)
-            const idx = output.content.length - 1
-            stream.push({
-              type: "toolcall_start",
-              contentIndex: idx,
-              partial: output,
-            })
             stream.push({
               type: "toolcall_end",
-              contentIndex: idx,
+              contentIndex,
               toolCall,
               partial: output,
             })

src/types.ts

@@ -174,6 +174,12 @@ export type AssistantMessageEvent =
       contentIndex: number
       partial: AssistantMessageLike
     }
+  | {
+      type: "toolcall_delta"
+      contentIndex: number
+      delta: string
+      partial: AssistantMessageLike
+    }
   | {
       type: "toolcall_end"
       contentIndex: number

tests/test-stream.ts

@@ -425,6 +425,104 @@ describe("streamCommandCode — successful streams", () => {
     assert.equal(toolCall?.type === "toolCall" ? toolCall.name : "", "read_file")
   })
 
+  it("streams incremental tool-call arguments from generate events", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [
+        JSON.stringify({
+          type: "tool-input-start",
+          id: "call_1",
+          toolName: "read_file",
+        }),
+        JSON.stringify({ type: "tool-input-delta", id: "call_1", delta: '{"path":"' }),
+        JSON.stringify({ type: "tool-input-delta", id: "call_1", delta: '/tmp/x"}' }),
+        JSON.stringify({ type: "tool-input-end", id: "call_1" }),
+        JSON.stringify({
+          type: "tool-call",
+          toolCallId: "call_1",
+          toolName: "read_file",
+          input: { path: "/tmp/x" },
+        }),
+        JSON.stringify({ type: "finish", finishReason: "tool-calls" }),
+      ],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    assert.deepEqual(eventTypes(events), [
+      "start",
+      "toolcall_start",
+      "toolcall_delta",
+      "toolcall_delta",
+      "toolcall_end",
+      "done",
+    ])
+    const deltas = events.flatMap((event) => (event.type === "toolcall_delta" ? [event.delta] : []))
+    assert.deepEqual(deltas, ['{"path":"', '/tmp/x"}'])
+
+    const done = events.at(-1)
+    if (done?.type !== "done") throw new Error("expected done")
+    assert.equal(done.reason, "toolUse")
+    const toolCall = done.message.content[0]
+    assert.equal(toolCall?.type, "toolCall")
+    if (toolCall?.type !== "toolCall") throw new Error("expected tool call")
+    assert.equal(toolCall.id, "call_1")
+    assert.equal(toolCall.name, "read_file")
+    assert.deepEqual(toolCall.arguments, { path: "/tmp/x" })
+  })
+
+  it("keeps concurrent incremental tool calls separate", async () => {
+    server.mockResponse({
+      type: "success",
+      events: [
+        JSON.stringify({ type: "tool-input-start", id: "call_1", toolName: "read_file" }),
+        JSON.stringify({ type: "tool-input-start", id: "call_2", toolName: "read_file" }),
+        JSON.stringify({ type: "tool-input-delta", id: "call_1", delta: '{"path":"/a"}' }),
+        JSON.stringify({ type: "tool-input-delta", id: "call_2", delta: '{"path":"/b"}' }),
+        JSON.stringify({
+          type: "tool-call",
+          toolCallId: "call_2",
+          toolName: "read_file",
+          input: { path: "/b" },
+        }),
+        JSON.stringify({
+          type: "tool-call",
+          toolCallId: "call_1",
+          toolName: "read_file",
+          input: { path: "/a" },
+        }),
+        JSON.stringify({ type: "finish", finishReason: "tool-calls" }),
+      ],
+    })
+    const { streamCommandCode } = createTestDeps({ apiBase: server.baseUrl() })
+
+    const events = await collectEvents(
+      streamCommandCode(makeModel(), makeContext(), { apiKey: "mock-key" }),
+    )
+
+    const starts = events.flatMap((event) =>
+      event.type === "toolcall_start" ? [event.contentIndex] : [],
+    )
+    const deltas = events.flatMap((event) =>
+      event.type === "toolcall_delta" ? [[event.contentIndex, event.delta] as const] : [],
+    )
+    const ends = events.flatMap((event) =>
+      event.type === "toolcall_end" ? [[event.contentIndex, event.toolCall.id] as const] : [],
+    )
+    assert.deepEqual(starts, [0, 1])
+    assert.deepEqual(deltas, [
+      [0, '{"path":"/a"}'],
+      [1, '{"path":"/b"}'],
+    ])
+    assert.deepEqual(ends, [
+      [1, "call_2"],
+      [0, "call_1"],
+    ])
+  })
+
   it("flushes reasoning if finish arrives without reasoning-end", async () => {
     server.mockResponse({
       type: "success",