feat: stream tool-call arguments progressively instead of buffering

internal/api/chat.go

@@ -191,9 +191,12 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open
 	var text strings.Builder
 	var firstDelta time.Time
 
-	// When tools are requested the whole reply is buffered so the raw
-	// tool-call JSON is never streamed to the client twice.
-	buffered := len(req.Tools) > 0
+	// With tools the raw tool-call JSON is converted into incremental
+	// tool_calls deltas instead of being buffered until the reply finishes.
+	var args *openai.ArgStreamer
+	if len(req.Tools) > 0 {
+		args = openai.NewArgStreamer()
+	}
 
 	write := func(v any) error {
 		raw, err := json.Marshal(v)
@@ -217,6 +220,16 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open
 			Choices: []openai.StreamChoice{{Delta: delta, FinishReason: finish}},
 		}
 	}
+	writeToolDelta := func(d openai.ToolArgDelta) error {
+		tc := openai.ToolCallDelta{Index: d.Index}
+		if d.ID != "" {
+			tc.ID, tc.Type = d.ID, "function"
+		}
+		if d.Name != "" || d.Args != "" {
+			tc.Function = &openai.FunctionDelta{Name: d.Name, Arguments: d.Args}
+		}
+		return write(chunk(openai.ChoiceDelta{ToolCalls: []openai.ToolCallDelta{tc}}, nil))
+	}
 
 	_ = write(chunk(openai.ChoiceDelta{Role: "assistant"}, nil))
 
@@ -229,10 +242,15 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open
 		}
 	}, func(seg string) error {
 		text.WriteString(seg)
-		if buffered {
-			return nil
+		if args == nil {
+			return write(chunk(openai.ChoiceDelta{Content: seg}, nil))
 		}
-		return write(chunk(openai.ChoiceDelta{Content: seg}, nil))
+		for _, d := range args.Feed(text.String()) {
+			if err := writeToolDelta(d); err != nil {
+				return err
+			}
+		}
+		return nil
 	})
 
 	duration := time.Since(start)
@@ -277,16 +295,33 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open
 	}
 
 	usage := estimateUsage(prompt, body)
-	var toolCalls []openai.ToolCall
-	content := body
-	if buffered {
-		remaining, calls, ok := openai.TryParseToolCalls(body)
-		if ok {
-			content = remaining
-			toolCalls = calls
+	// content is only what was NOT streamed: prose around a tool JSON, or the
+	// whole reply when a tools request did not parse as a tool call at all.
+	// Plain text was already streamed per delta above.
+	content := ""
+	stop := "stop"
+	var resend []openai.ToolCall
+	if args == nil {
+		// plain reply: content streamed live, final chunk carries no content
+	} else if remaining, calls, ok := openai.TryParseToolCalls(body); ok {
+		content = remaining
+		stop = "tool_calls"
+		deltas, miss, diverged := args.Finalize(calls)
+		if diverged {
+			s.log.Warn("streamed tool arguments diverged from repaired parse", "model", req.Model, "detail", args.Diverge())
+		}
+		for _, d := range deltas {
+			_ = writeToolDelta(d)
+		}
+		resend = miss
+	} else {
+		// tools were requested but the reply is not a tool call: it was never
+		// streamed, so deliver it as content in the final chunk.
+		content = body
+		if args.Sent() {
+			s.log.Warn("tool parse failed after streaming argument deltas", "model", req.Model)
 		}
 	}
-	stop := finishReason(toolCalls)
 
 	// Final chunk: content + tool_calls (or plain stop marker).
 	final := map[string]any{
@@ -304,8 +339,8 @@ func (s *Server) streamChat(ctx context.Context, w http.ResponseWriter, req open
 	if content != "" {
 		deltaPayload["content"] = content
 	}
-	if len(toolCalls) > 0 {
-		deltaPayload["tool_calls"] = toolCalls
+	if len(resend) > 0 {
+		deltaPayload["tool_calls"] = resend
 	}
 	if len(deltaPayload) > 0 {
 		final["choices"] = []map[string]any{{

internal/api/messages.go

@@ -140,17 +140,90 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req
 		},
 	})
 
-	// When tools are requested the reply is buffered so the raw tool-call
-	// JSON is parsed once and emitted as proper tool_use blocks instead of
-	// being streamed to the client as prose.
-	buffered := len(tools) > 0
-	if !buffered {
+	// With tools the raw tool-call JSON is converted into streamed tool_use
+	// blocks: content_block_start once id+name are known, then one
+	// input_json_delta per piece, instead of buffering the whole reply.
+	var args *openai.ArgStreamer
+	if len(tools) > 0 {
+		args = openai.NewArgStreamer()
+	}
+	if args == nil {
 		_ = writeEvent("content_block_start", map[string]any{
 			"type": "content_block_start", "index": 0,
 			"content_block": map[string]any{"type": "text", "text": ""},
 		})
 	}
 
+	// Anthropic content blocks are strictly sequential: close the open
+	// tool_use block before the next one starts. Deltas that arrive for an
+	// older block after a newer one opened are dropped — models emit calls
+	// in order.
+	var (
+		blockPos = 0  // next content block index
+		curIdx   = -1 // streamer index in focus
+		openPos  = -1 // content index of the open tool_use block, -1 = none
+		ids      = map[int]string{}
+		names    = map[int]string{}
+		pending  = map[int][]string{} // args held until the block can open
+	)
+	onToolDelta := func(d openai.ToolArgDelta) error {
+		if d.Index < curIdx {
+			return nil
+		}
+		if d.ID != "" {
+			ids[d.Index] = d.ID
+		}
+		if d.Name != "" {
+			names[d.Index] = d.Name
+		}
+		if d.Index > curIdx {
+			if openPos >= 0 {
+				if err := writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": openPos}); err != nil {
+					return err
+				}
+				openPos = -1
+				blockPos++
+			}
+			curIdx = d.Index
+		}
+		if openPos < 0 && names[d.Index] != "" {
+			id := ids[d.Index]
+			if id == "" {
+				id = anthropic.NewID("toolu_")
+			}
+			openPos = blockPos
+			if err := writeEvent("content_block_start", map[string]any{
+				"type": "content_block_start", "index": openPos,
+				"content_block": map[string]any{
+					"type": "tool_use", "id": id, "name": names[d.Index],
+					"input": json.RawMessage("{}"),
+				},
+			}); err != nil {
+				return err
+			}
+			for _, frag := range pending[d.Index] {
+				if err := writeEvent("content_block_delta", map[string]any{
+					"type": "content_block_delta", "index": openPos,
+					"delta": map[string]any{"type": "input_json_delta", "partial_json": frag},
+				}); err != nil {
+					return err
+				}
+			}
+			delete(pending, d.Index)
+		}
+		if d.Args == "" {
+			return nil
+		}
+		if openPos < 0 {
+			pending[d.Index] = append(pending[d.Index], d.Args)
+			return nil
+		}
+		return writeEvent("content_block_delta", map[string]any{
+			"type": "content_block_delta", "index": openPos,
+			"delta": map[string]any{"type": "input_json_delta", "partial_json": d.Args},
+		})
+	}
+
 	start := time.Now()
 	info := &upstream.ServeInfo{}
 	var text strings.Builder
@@ -162,13 +235,18 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req
 		}
 	}, func(seg string) error {
 		text.WriteString(seg)
-		if buffered {
-			return nil
+		if args == nil {
+			return writeEvent("content_block_delta", map[string]any{
+				"type": "content_block_delta", "index": 0,
+				"delta": map[string]any{"type": "text_delta", "text": seg},
+			})
 		}
-		return writeEvent("content_block_delta", map[string]any{
-			"type": "content_block_delta", "index": 0,
-			"delta": map[string]any{"type": "text_delta", "text": seg},
-		})
+		for _, d := range args.Feed(text.String()) {
+			if err := onToolDelta(d); err != nil {
+				return err
+			}
+		}
+		return nil
 	})
 
 	duration := time.Since(start)
@@ -201,40 +279,63 @@ func (s *Server) streamMessages(ctx context.Context, w http.ResponseWriter, req
 	usage := estimateUsage(prompt, body)
 	outText, calls := splitToolOutput(body, tools)
 
-	index := 0
-	if !buffered {
-		_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": 0})
-		index = 1
-	} else if outText != "" {
-		_ = writeEvent("content_block_start", map[string]any{
-			"type": "content_block_start", "index": 0,
-			"content_block": map[string]any{"type": "text", "text": ""},
-		})
-		_ = writeEvent("content_block_delta", map[string]any{
-			"type": "content_block_delta", "index": 0,
-			"delta": map[string]any{"type": "text_delta", "text": outText},
-		})
-		_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": 0})
-		index = 1
-	}
-	for _, c := range calls {
-		id := c.ID
-		if id == "" {
-			id = anthropic.NewID("toolu_")
+	if args != nil {
+		if len(calls) > 0 {
+			deltas, resend, diverged := args.Finalize(calls)
+			if diverged {
+				s.log.Warn("streamed tool arguments diverged from repaired parse", "model", req.Model, "detail", args.Diverge())
+			}
+			for _, d := range deltas {
+				_ = onToolDelta(d)
+			}
+			if openPos >= 0 {
+				_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": openPos})
+				openPos = -1
+				blockPos++
+			}
+			for _, c := range resend {
+				id := c.ID
+				if id == "" {
+					id = anthropic.NewID("toolu_")
+				}
+				_ = writeEvent("content_block_start", map[string]any{
+					"type": "content_block_start", "index": blockPos,
+					"content_block": map[string]any{
+						"type": "tool_use", "id": id, "name": c.Function.Name,
+						"input": json.RawMessage("{}"),
+					},
+				})
+				_ = writeEvent("content_block_delta", map[string]any{
+					"type": "content_block_delta", "index": blockPos,
+					"delta": map[string]any{"type": "input_json_delta", "partial_json": c.Function.Arguments},
+				})
+				_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": blockPos})
+				blockPos++
+			}
+		} else {
+			if args.Sent() {
+				s.log.Warn("tool parse failed after streaming argument deltas", "model", req.Model)
+			}
+			if openPos >= 0 {
+				_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": openPos})
+				openPos = -1
+				blockPos++
+			}
 		}
-		_ = writeEvent("content_block_start", map[string]any{
-			"type": "content_block_start", "index": index,
-			"content_block": map[string]any{
-				"type": "tool_use", "id": id, "name": c.Function.Name,
-				"input": json.RawMessage("{}"),
-			},
-		})
-		_ = writeEvent("content_block_delta", map[string]any{
-			"type": "content_block_delta", "index": index,
-			"delta": map[string]any{"type": "input_json_delta", "partial_json": c.Function.Arguments},
-		})
-		_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": index})
-		index++
+		if outText != "" {
+			_ = writeEvent("content_block_start", map[string]any{
+				"type": "content_block_start", "index": blockPos,
+				"content_block": map[string]any{"type": "text", "text": ""},
+			})
+			_ = writeEvent("content_block_delta", map[string]any{
+				"type": "content_block_delta", "index": blockPos,
+				"delta": map[string]any{"type": "text_delta", "text": outText},
+			})
+			_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": blockPos})
+			blockPos++
+		}
+	} else {
+		_ = writeEvent("content_block_stop", map[string]any{"type": "content_block_stop", "index": 0})
 	}
 	_ = writeEvent("message_delta", map[string]any{
 		"type":  "message_delta",

internal/api/responses.go

@@ -149,9 +149,105 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req
 		},
 	})
 
-	// When tools are requested the reply is buffered so the raw tool-call
-	// JSON is parsed once and emitted as proper function_call output items.
-	buffered := len(tools) > 0
+	// With tools the raw tool-call JSON is converted into a live
+	// function_call item stream: output_item.added once id+name are known,
+	// one function_call_arguments.delta per piece, done events at finalize.
+	var args *openai.ArgStreamer
+	if len(tools) > 0 {
+		args = openai.NewArgStreamer()
+	}
+
+	var (
+		seq     = 0
+		outIdx  = 0
+		acc     = map[int]string{} // full args seen per streamer index
+		emitted = map[int]int{}    // bytes of acc written as deltas
+		itemIDs = map[int]string{}
+		callIDs = map[int]string{}
+		fnNames = map[int]string{}
+		openIdx = -1 // streamer index of the in-progress function_call item
+		openOut = -1
+		output  []map[string]any
+	)
+	closeOpen := func() error {
+		if openIdx < 0 {
+			return nil
+		}
+		itemID := itemIDs[openIdx]
+		if err := write(map[string]any{
+			"type": "response.function_call_arguments.done", "item_id": itemID,
+			"output_index": openOut, "arguments": acc[openIdx], "sequence_number": seq,
+		}); err != nil {
+			return err
+		}
+		seq++
+		item := map[string]any{
+			"id": itemID, "type": "function_call", "status": "completed",
+			"call_id": callIDs[openIdx], "name": fnNames[openIdx], "arguments": acc[openIdx],
+		}
+		if err := write(map[string]any{
+			"type": "response.output_item.done", "output_index": openOut,
+			"sequence_number": seq, "item": item,
+		}); err != nil {
+			return err
+		}
+		seq++
+		output = append(output, item)
+		openIdx, openOut = -1, -1
+		return nil
+	}
+	onToolDelta := func(d openai.ToolArgDelta) error {
+		if d.Index < openIdx {
+			return nil // models emit calls in order — drop stale index
+		}
+		if d.ID != "" {
+			callIDs[d.Index] = d.ID
+		}
+		if d.Name != "" {
+			fnNames[d.Index] = d.Name
+		}
+		if d.Index != openIdx {
+			if err := closeOpen(); err != nil {
+				return err
+			}
+			if callIDs[d.Index] == "" || fnNames[d.Index] == "" {
+				// cannot open the item yet — hold args until id+name arrive
+				if d.Args != "" {
+					acc[d.Index] += d.Args
+				}
+				return nil
+			}
+			itemID := "fc_" + randHex(12)
+			itemIDs[d.Index] = itemID
+			openOut = outIdx
+			openIdx = d.Index
+			outIdx++
+			if err := write(map[string]any{
+				"type": "response.output_item.added", "output_index": openOut, "sequence_number": seq,
+				"item": map[string]any{
+					"id": itemID, "type": "function_call", "status": "in_progress",
+					"call_id": callIDs[d.Index], "name": fnNames[d.Index], "arguments": "",
+				},
+			}); err != nil {
+				return err
+			}
+			seq++
+		}
+		if d.Args != "" {
+			acc[d.Index] += d.Args
+		}
+		if openIdx != d.Index {
+			return nil
+		}
+		if fresh := acc[d.Index][emitted[d.Index]:]; fresh != "" {
+			emitted[d.Index] = len(acc[d.Index])
+			return write(map[string]any{
+				"type": "response.function_call_arguments.delta", "item_id": itemIDs[d.Index],
+				"output_index": openOut, "delta": fresh, "sequence_number": seq,
+			})
+		}
+		return nil
+	}
 
 	start := time.Now()
 	info := &upstream.ServeInfo{}
@@ -164,10 +260,15 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req
 		}
 	}, func(seg string) error {
 		text.WriteString(seg)
-		if buffered {
-			return nil
+		if args == nil {
+			return write(map[string]any{"type": "response.output_text.delta", "delta": seg})
 		}
-		return write(map[string]any{"type": "response.output_text.delta", "delta": seg})
+		for _, d := range args.Feed(text.String()) {
+			if err := onToolDelta(d); err != nil {
+				return err
+			}
+		}
+		return nil
 	})
 
 	duration := time.Since(start)
@@ -219,19 +320,67 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req
 	usage := estimateUsage(prompt, body)
 	outText, calls := splitToolOutput(body, tools)
 
-	stopReason := "stop"
+	var stopReason = "stop"
 	if len(calls) > 0 {
 		stopReason = "tool_use"
 	}
 
-	var output []map[string]any
-	if buffered {
-		// Replay the buffered reply as a proper item stream: message block
-		// first (when there is prose), then one function_call item per call.
-		// The final completed.response.output reuses the same item ids.
-		output = make([]map[string]any, 0, len(calls)+1)
-		seq := 0
-		outIdx := 0
+	if args == nil {
+		output = responseOutputItems(outText, calls)
+		_ = write(map[string]any{
+			"type": "response.output_text.done",
+			"text": body,
+		})
+	} else {
+		if len(calls) > 0 {
+			deltas, resend, diverged := args.Finalize(calls)
+			if diverged {
+				s.log.Warn("streamed tool arguments diverged from repaired parse", "model", req.Model, "detail", args.Diverge())
+			}
+			for _, d := range deltas {
+				_ = onToolDelta(d)
+			}
+			// close the open item with the accumulated args
+			_ = closeOpen()
+			// calls that were never streamed become full items now
+			for _, c := range resend {
+				itemID := "fc_" + randHex(12)
+				item := map[string]any{
+					"id": itemID, "type": "function_call", "status": "completed",
+					"call_id": c.ID, "name": c.Function.Name, "arguments": c.Function.Arguments,
+				}
+				_ = write(map[string]any{
+					"type": "response.output_item.added", "output_index": outIdx, "sequence_number": seq,
+					"item": map[string]any{
+						"id": itemID, "type": "function_call", "status": "in_progress",
+						"call_id": c.ID, "name": c.Function.Name, "arguments": "",
+					},
+				})
+				seq++
+				_ = write(map[string]any{
+					"type": "response.function_call_arguments.delta", "item_id": itemID,
+					"output_index": outIdx, "delta": c.Function.Arguments, "sequence_number": seq,
+				})
+				seq++
+				_ = write(map[string]any{
+					"type": "response.function_call_arguments.done", "item_id": itemID,
+					"output_index": outIdx, "arguments": c.Function.Arguments, "sequence_number": seq,
+				})
+				seq++
+				_ = write(map[string]any{
+					"type": "response.output_item.done", "output_index": outIdx,
+					"sequence_number": seq, "item": item,
+				})
+				seq++
+				outIdx++
+				output = append(output, item)
+			}
+		} else {
+			if args.Sent() {
+				s.log.Warn("tool parse failed after streaming argument deltas", "model", req.Model)
+			}
+			_ = closeOpen()
+		}
 		if outText != "" {
 			itemID := "msg_" + randHex(12)
 			item := map[string]any{
@@ -250,42 +399,12 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req
 			seq++
 			_ = write(map[string]any{
 				"type": "response.output_text.delta", "item_id": itemID,
-				"output_index": outIdx, "delta": outText, "content_index": 0,
+				"output_index": outIdx, "delta": outText, "content_index": 0, "sequence_number": seq,
 			})
+			seq++
 			_ = write(map[string]any{
 				"type": "response.output_text.done", "item_id": itemID,
-				"output_index": outIdx, "text": outText, "content_index": 0,
-			})
-			_ = write(map[string]any{
-				"type": "response.output_item.done", "output_index": outIdx,
-				"sequence_number": seq, "item": item,
-			})
-			seq++
-			outIdx++
-			output = append(output, item)
-		}
-		for _, c := range calls {
-			itemID := "fc_" + randHex(12)
-			item := map[string]any{
-				"id": itemID, "type": "function_call", "status": "completed",
-				"call_id": c.ID, "name": c.Function.Name, "arguments": c.Function.Arguments,
-			}
-			_ = write(map[string]any{
-				"type": "response.output_item.added", "output_index": outIdx, "sequence_number": seq,
-				"item": map[string]any{
-					"id": itemID, "type": "function_call", "status": "in_progress",
-					"call_id": c.ID, "name": c.Function.Name, "arguments": "",
-				},
-			})
-			seq++
-			_ = write(map[string]any{
-				"type": "response.function_call_arguments.delta", "item_id": itemID,
-				"output_index": outIdx, "delta": c.Function.Arguments, "sequence_number": seq,
-			})
-			seq++
-			_ = write(map[string]any{
-				"type": "response.function_call_arguments.done", "item_id": itemID,
-				"output_index": outIdx, "arguments": c.Function.Arguments, "sequence_number": seq,
+				"output_index": outIdx, "text": outText, "content_index": 0, "sequence_number": seq,
 			})
 			seq++
 			_ = write(map[string]any{
@@ -296,12 +415,6 @@ func (s *Server) streamResponses(ctx context.Context, w http.ResponseWriter, req
 			outIdx++
 			output = append(output, item)
 		}
-	} else {
-		output = responseOutputItems(outText, calls)
-		_ = write(map[string]any{
-			"type": "response.output_text.done",
-			"text": body,
-		})
 	}
 	_ = write(map[string]any{
 		"type": "response.completed",