From 18bf4ee723fcfee2ec59bb0ad9197b8ee0a93ece Mon Sep 17 00:00:00 2001 From: nyoungo Date: Thu, 8 Oct 2026 01:07:10 +0800 Subject: [PATCH] feat: stream tool-call arguments progressively instead of buffering --- internal/api/chat.go | 67 +++-- internal/api/messages.go | 189 ++++++++++--- internal/api/responses.go | 221 +++++++++++---- internal/openai/argstream.go | 432 ++++++++++++++++++++++++++++++ internal/openai/argstream_test.go | 259 ++++++++++++++++++ internal/openai/types.go | 21 +- 6 files changed, 1072 insertions(+), 117 deletions(-) create mode 100644 internal/openai/argstream.go create mode 100644 internal/openai/argstream_test.go diff --git a/internal/api/chat.go b/internal/api/chat.go index 0d2c038..7fa3a2b 100644 --- a/internal/api/chat.go +++ b/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{{ diff --git a/internal/api/messages.go b/internal/api/messages.go index 95fc29a..bacfe7f 100644 --- a/internal/api/messages.go +++ b/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", diff --git a/internal/api/responses.go b/internal/api/responses.go index f3e9836..8591b0f 100644 --- a/internal/api/responses.go +++ b/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", diff --git a/internal/openai/argstream.go b/internal/openai/argstream.go new file mode 100644 index 0000000..e015cec --- /dev/null +++ b/internal/openai/argstream.go @@ -0,0 +1,432 @@ +package openai + +import ( + "encoding/json" + "fmt" + "regexp" + "strings" +) + +// ToolArgDelta is one incremental tool_calls delta produced by ArgStreamer. +// Empty fields are omitted from the wire chunk. +type ToolArgDelta struct { + Index int + ID string + Name string + Args string +} + +// ArgStreamer turns raw tool-call JSON that grows over time into OpenAI-style +// incremental tool_calls deltas so tool requests stream instead of buffering. +// +// The raw document is re-scanned on every Feed; the decoded arguments text is +// a prefix-monotonic function of the document, so each delta is exactly the +// new suffix. Finalize reconciles the stream with the repaired parse result +// (balanceJSON / EscapeControlChars / forceCloseJSON) and reports any call +// that could not be streamed. +type ArgStreamer struct { + sent map[int]string // args text already written downstream + derived map[int]string // latest full decode (incl. hold-back tail) + idSent map[int]string + nameSent map[int]string + seen map[int]bool + nostream map[int]bool // arguments delivered as object/number — resend whole + diverge string // first divergence context, set by Finalize +} + +// NewArgStreamer creates a streamer for one request. +func NewArgStreamer() *ArgStreamer { + return &ArgStreamer{ + sent: map[int]string{}, + derived: map[int]string{}, + idSent: map[int]string{}, + nameSent: map[int]string{}, + seen: map[int]bool{}, + nostream: map[int]bool{}, + } +} + +// Sent reports whether any delta was written downstream. +func (a *ArgStreamer) Sent() bool { return len(a.seen) > 0 } + +// Diverge returns a short diagnostic of where Finalize found the repaired +// arguments diverging from the streamed text (empty when clean). +func (a *ArgStreamer) Diverge() string { return a.diverge } + +func (a *ArgStreamer) noteDiverge(idx int, target, emit, sent string) { + if a.diverge != "" { + return + } + at := firstDiffIndex(target, emit) + if d := firstDiffIndex(target, sent); d < at { + at = d + } + a.diverge = fmt.Sprintf("idx=%d tLen=%d eLen=%d sLen=%d diff@%d target=%q emit=%q", + idx, len(target), len(emit), len(sent), at, ctxWindow(target, at), ctxWindow(emit, at)) +} + +func firstDiffIndex(a, b string) int { + n := len(a) + if len(b) < n { + n = len(b) + } + for i := 0; i < n; i++ { + if a[i] != b[i] { + return i + } + } + return n +} + +func ctxWindow(s string, at int) string { + const win = 60 + start := at - win + if start < 0 { + start = 0 + } + end := at + win + if end > len(s) { + end = len(s) + } + return s[start:end] +} + +type argCallState struct { + derived string + id string + name string + nostream bool +} + +// Feed scans the full accumulated document and returns the deltas to send. +func (a *ArgStreamer) Feed(text string) []ToolArgDelta { + states := a.scan(text) + var out []ToolArgDelta + for i, st := range states { + if st.nostream { + a.nostream[i] = true + continue + } + if a.nostream[i] { + continue + } + emit := trimHold(st.derived) + prev := a.sent[i] + if !strings.HasPrefix(emit, prev) { + continue // defensive: decode regressed (should be monotonic) + } + d := ToolArgDelta{Index: i} + send := false + if a.idSent[i] == "" && st.id != "" { + d.ID = st.id + a.idSent[i] = st.id + send = true + } + if st.name != "" && a.nameSent[i] != st.name { + d.Name = st.name + a.nameSent[i] = st.name + send = true + } + if args := emit[len(prev):]; args != "" { + d.Args = args + send = true + } + a.sent[i] = emit + a.derived[i] = st.derived + if send { + a.seen[i] = true + out = append(out, d) + } + } + return out +} + +// Finalize reconciles the stream against the final parsed tool calls. +// deltas must be sent before the finish chunk; resend holds calls that were +// never streamed (send them in full, with their index). diverged signals that +// the repaired arguments are not an extension of what was already streamed — +// the client keeps the streamed text in that case. +func (a *ArgStreamer) Finalize(calls []ToolCall) (deltas []ToolArgDelta, resend []ToolCall, diverged bool) { + for i, c := range calls { + if a.nostream[i] || !a.seen[i] { + idx := i + c.Index = &idx + resend = append(resend, c) + continue + } + target := c.Function.Arguments + emit := trimHold(a.derived[i]) + sent := a.sent[i] + d := ToolArgDelta{Index: i} + flush := "" + switch { + case strings.HasPrefix(target, emit): + // append-only repair (closers etc.): send the rest incl. held tail + flush = target[len(sent):] + case strings.HasPrefix(emit, target): + // repair dropped the held tail (dangling comma/whitespace). If the + // client has not been sent beyond the new end, flush up to it; + // otherwise it keeps extra text and we can only report it. + if strings.HasPrefix(target, sent) { + flush = target[len(sent):] + } else { + diverged = true + a.noteDiverge(i, target, emit, sent) + } + case strings.HasPrefix(target, sent): + // repair only rewrote the unsent held region: catch the client up + flush = target[len(sent):] + default: + // mid-document cut (forceCloseJSON anchor rebuild): already sent + // text can differ from the repair — report and keep the stream. + diverged = true + a.noteDiverge(i, target, emit, sent) + } + if flush != "" { + d.Args = flush + } + if a.idSent[i] == "" && c.ID != "" { + d.ID = c.ID + } + if c.Function.Name != "" && a.nameSent[i] != c.Function.Name { + d.Name = c.Function.Name + } + if d.Args != "" || d.ID != "" || d.Name != "" { + deltas = append(deltas, d) + } + } + return deltas, resend, diverged +} + +var ( + argsKeyWord = []byte(`"arguments"`) + fenceRe = regexp.MustCompile("```") + tcNameRe = regexp.MustCompile(`"name"\s*:\s*"((?:[^"\\]|\\.)*)"`) + tcIDRe = regexp.MustCompile(`"id"\s*:\s*"((?:[^"\\]|\\.)*)"`) +) + +func (a *ArgStreamer) scan(text string) []argCallState { + var states []argCallState + if !extractableToolJSON(text) { + return states + } + // Only scan the payload: with prose in front the arguments keys before + // it must not count as tool calls. A fenced payload starts after its + // fence; fences later than the payload live inside a string value and + // must not move the base. + base := 0 + if !strings.HasPrefix(strings.TrimSpace(text), "{") { + if tc := strings.Index(text, `{"tool_calls"`); tc > 0 { + base = tc + } else if f := fenceRe.FindStringIndex(text); f != nil { + base = f[1] + } + } + prevEnd := base + off := base + for { + j := strings.Index(text[off:], string(argsKeyWord)) + if j < 0 { + break + } + pos := off + j + off = pos + len(argsKeyWord) + winStart := prevEnd + if pos-winStart > 8192 { + winStart = pos - 8192 + } + if winStart > pos { + winStart = pos + } + id, name := lookbackIDName(text[winStart:pos]) + k := skipSpace(text, pos+len(argsKeyWord)) + if k >= len(text) || text[k] != ':' { + continue // the word without key syntax + } + st := argCallState{id: id, name: name} + v := skipSpace(text, k+1) + if v >= len(text) { + prevEnd = len(text) + } else { + switch text[v] { + case '"': + decoded, end := decodeArgsValue(text, v+1) + st.derived = innerEscape(decoded) + prevEnd = end + case '{', '[': + st.nostream = true + prevEnd = v + 1 + default: + st.nostream = true + prevEnd = v + 1 + } + } + states = append(states, st) + } + return states +} + +// extractableToolJSON mirrors the candidate gate in TryParseToolCalls: the +// document must look like it could become a tool_calls payload — starting +// with the JSON, fenced, or with prose in front of the payload. +func extractableToolJSON(text string) bool { + t := strings.TrimSpace(text) + if strings.HasPrefix(t, "{") { + return strings.Contains(t, `"tool_calls"`) + } + return strings.Contains(t, `{"tool_calls"`) +} + +func skipSpace(s string, i int) int { + for i < len(s) { + switch s[i] { + case ' ', '\t', '\r', '\n': + i++ + default: + return i + } + } + return i +} + +// decodeArgsValue decodes the outer-JSON string value starting after the +// opening quote at start. It returns the decoded inner text and the index +// after the closing quote (or len(text) when the value is still open). An +// incomplete escape pair at the end of text is withheld for the next scan. +func decodeArgsValue(text string, start int) (string, int) { + var b strings.Builder + i := start + for i < len(text) { + c := text[i] + if c == '\\' { + if i+1 >= len(text) { + break + } + c2 := text[i+1] + switch c2 { + case '"', '\\', '/': + b.WriteByte(c2) + case 'b': + b.WriteByte('\b') + case 'f': + b.WriteByte('\f') + case 'n': + b.WriteByte('\n') + case 'r': + b.WriteByte('\r') + case 't': + b.WriteByte('\t') + case 'u': + r, n := decodeUnicodeEscape(text, i) + if n == 0 { + if i+6 > len(text) { + // incomplete at end: withhold for the next scan + return b.String(), i + } + // malformed but complete (\uZZZZ): emit raw — the + // inner heal doubles the backslash, like the parse path + b.WriteString(`\u`) + i += 2 + continue + } + b.WriteRune(r) + i += n + continue + default: + // invalid outer escape (…\d): keep the backslash as inner + // text so the inner heal matches the repaired parse + b.WriteByte('\\') + b.WriteByte(c2) + } + i += 2 + continue + } + if c == '"' { + return b.String(), i + 1 + } + b.WriteByte(c) + i++ + } + return b.String(), len(text) +} + +// decodeUnicodeEscape decodes \uXXXX (and surrogate pairs) at text[i]. +// It returns the rune and total sequence length, or (0, 0) when incomplete. +func decodeUnicodeEscape(text string, i int) (rune, int) { + if i+6 > len(text) || text[i+1] != 'u' || !isHex4(text[i+2:i+6]) { + return 0, 0 + } + var r rune + if _, err := fmt.Sscanf(text[i+2:i+6], "%x", &r); err != nil { + return 0, 0 + } + if r >= 0xD800 && r <= 0xDBFF { + // high surrogate — needs a low surrogate to follow + if i+12 <= len(text) && text[i+6] == '\\' && text[i+7] == 'u' && isHex4(text[i+8:i+12]) { + var lo rune + if _, err := fmt.Sscanf(text[i+8:i+12], "%x", &lo); err == nil && lo >= 0xDC00 && lo <= 0xDFFF { + return 0x10000 + (r-0xD800)<<10 + (lo - 0xDC00), 12 + } + } + return r, 6 // lone surrogate — emit as-is, clients tolerate + } + return r, 6 +} + +// innerEscape applies EscapeControlChars semantics to decoded inner text. +func innerEscape(decoded string) string { + if !needsEscape(decoded) { + return decoded + } + var st escapeState + var out strings.Builder + out.Grow(len(decoded) + 16) + st.feed(decoded, &out) + return out.String() +} + +// trimHold strips the suffix that balanceJSON would trim or drop when the +// channel cap stops the reply there: whitespace, a dangling comma/colon, or a +// lone backslash. Held text is released by a later Feed or Finalize. +func trimHold(s string) string { + i := len(s) + for i > 0 { + switch s[i-1] { + case ' ', '\t', '\r', '\n', ',', ':': + i-- + default: + goto backslash + } + } +backslash: + if i > 0 && s[i-1] == '\\' { + run := 0 + for j := i - 1; j >= 0 && s[j] == '\\'; j-- { + run++ + } + if run%2 == 1 { + i-- + } + } + return s[:i] +} + +// lookbackIDName finds the id / name of the tool call that owns the upcoming +// arguments key: the last occurrence before it within the local window. +func lookbackIDName(window string) (id, name string) { + if m := tcIDRe.FindAllStringSubmatch(window, -1); len(m) > 0 { + id = decodeJSONString(m[len(m)-1][1]) + } + if m := tcNameRe.FindAllStringSubmatch(window, -1); len(m) > 0 { + name = decodeJSONString(m[len(m)-1][1]) + } + return id, name +} + +func decodeJSONString(s string) string { + var out string + if err := json.Unmarshal([]byte(`"`+s+`"`), &out); err == nil { + return out + } + return s +} diff --git a/internal/openai/argstream_test.go b/internal/openai/argstream_test.go new file mode 100644 index 0000000..db58555 --- /dev/null +++ b/internal/openai/argstream_test.go @@ -0,0 +1,259 @@ +package openai + +import ( + "encoding/json" + "strings" + "testing" +) + +const ( + asDoc = `{"tool_calls":[{"id":"call_7f","type":"function","function":{"name":"write_file","arguments":"{\"path\":\"a.txt\",\"content\":\"hello world\"}"}}]}` + asInner = `{"path":"a.txt","content":"hello world"}` +) + +// feedAll streams doc in step-byte fragments and returns the concatenated +// argument deltas plus the first delta that carried id/name. +func feedAll(a *ArgStreamer, doc string, step int) (args string, head ToolArgDelta, headSet bool) { + var b strings.Builder + for i := step; i <= len(doc); i += step { + for _, d := range a.Feed(doc[:i]) { + if !headSet && (d.ID != "" || d.Name != "") { + head, headSet = d, true + } + b.WriteString(d.Args) + } + } + for _, d := range a.Feed(doc) { + if !headSet && (d.ID != "" || d.Name != "") { + head, headSet = d, true + } + b.WriteString(d.Args) + } + return b.String(), head, headSet +} + +func TestArgStreamerFragmentsReassemble(t *testing.T) { + a := NewArgStreamer() + got, head, headSet := feedAll(a, asDoc, 7) + if got != asInner { + t.Fatalf("args mismatch:\n got %q\nwant %q", got, asInner) + } + if !headSet || head.ID != "call_7f" || head.Name != "write_file" { + t.Fatalf("head delta missing id/name: %+v set=%v", head, headSet) + } + if !a.Sent() { + t.Fatal("Sent() = false after streaming") + } + calls := []ToolCall{{ID: "call_7f", Type: "function", Function: FunctionCall{Name: "write_file", Arguments: asInner}}} + deltas, resend, diverged := a.Finalize(calls) + if diverged || len(resend) != 0 || len(deltas) != 0 { + t.Fatalf("clean finalize: deltas=%+v resend=%+v diverged=%v", deltas, resend, diverged) + } +} + +func TestArgStreamerHoldsTrailingComma(t *testing.T) { + doc := `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"f","arguments":"{\"a\":1,` + a := NewArgStreamer() + got, _, _ := feedAll(a, doc, 3) + if got != `{"a":1` { + t.Fatalf("held comma leaked: %q", got) + } + // balanceJSON trims the dangling comma and closes the object + calls := []ToolCall{{ID: "c1", Type: "function", Function: FunctionCall{Name: "f", Arguments: `{"a":1}`}}} + deltas, resend, diverged := a.Finalize(calls) + if diverged || len(resend) != 0 { + t.Fatalf("diverged=%v resend=%+v", diverged, resend) + } + var flush strings.Builder + for _, d := range deltas { + flush.WriteString(d.Args) + } + if got+flush.String() != `{"a":1}` { + t.Fatalf("finalize flush wrong: streamed=%q flush=%q", got, flush.String()) + } +} + +func TestArgStreamerHoldsDanglingBackslash(t *testing.T) { + // inner content ends with a lone backslash (truncated escape) + doc := `{"tool_calls":[{"id":"c1","type":"function","function":{"name":"f","arguments":"{\"path\":\"ab\` + a := NewArgStreamer() + got, _, _ := feedAll(a, doc, 100) + if strings.HasSuffix(got, `\`) { + t.Fatalf("dangling backslash streamed: %q", got) + } + if got != `{"path":"ab` { + t.Fatalf("unexpected prefix: %q", got) + } +} + +func TestArgStreamerObjectArgumentsResend(t *testing.T) { + doc := `{"tool_calls":[{"id":"c2","type":"function","function":{"name":"f","arguments":{"a":1}}}]}` + a := NewArgStreamer() + got, _, _ := feedAll(a, doc, 5) + if got != "" { + t.Fatalf("object-form arguments must not stream: %q", got) + } + if a.Sent() { + t.Fatal("Sent() = true for object-form") + } + calls := []ToolCall{{ID: "c2", Type: "function", Function: FunctionCall{Name: "f", Arguments: `{"a":1}`}}} + deltas, resend, diverged := a.Finalize(calls) + if diverged || len(deltas) != 0 || len(resend) != 1 { + t.Fatalf("resend expected: deltas=%+v resend=%+v diverged=%v", deltas, resend, diverged) + } + if resend[0].Index == nil || *resend[0].Index != 0 || resend[0].Function.Arguments != `{"a":1}` { + t.Fatalf("resend entry wrong: %+v", resend[0]) + } +} + +func TestArgStreamerDivergedOnMidCut(t *testing.T) { + doc := `{"tool_calls":[{"id":"c3","type":"function","function":{"name":"f","arguments":"{\"a\":1}"}}]}` + a := NewArgStreamer() + feedAll(a, doc, 4) + calls := []ToolCall{{ID: "c3", Type: "function", Function: FunctionCall{Name: "f", Arguments: `{"z":9}`}}} + _, _, diverged := a.Finalize(calls) + if !diverged { + t.Fatal("mid-document rewrite not reported") + } +} + +func TestArgStreamerMultipleCalls(t *testing.T) { + doc := `{"tool_calls":[` + + `{"id":"c1","type":"function","function":{"name":"one","arguments":"{\"x\":1}"}},` + + `{"id":"c2","type":"function","function":{"name":"two","arguments":"{\"y\":2}"}}]}` + inner1, inner2 := `{"x":1}`, `{"y":2}` + a := NewArgStreamer() + var got0, got1 strings.Builder + var head0, head1 ToolArgDelta + for i := 1; i <= len(doc); i++ { + for _, d := range a.Feed(doc[:i]) { + if d.ID != "" { + if d.Index == 0 { + head0 = d + } else { + head1 = d + } + } + if d.Index == 0 { + got0.WriteString(d.Args) + } else { + got1.WriteString(d.Args) + } + } + } + if got0.String() != inner1 || got1.String() != inner2 { + t.Fatalf("args: idx0=%q idx1=%q", got0.String(), got1.String()) + } + if head0.ID != "c1" || head0.Name != "one" || head1.ID != "c2" || head1.Name != "two" { + t.Fatalf("heads: 0=%+v 1=%+v", head0, head1) + } + calls := []ToolCall{ + {ID: "c1", Type: "function", Function: FunctionCall{Name: "one", Arguments: inner1}}, + {ID: "c2", Type: "function", Function: FunctionCall{Name: "two", Arguments: inner2}}, + } + _, resend, diverged := a.Finalize(calls) + if diverged || len(resend) != 0 { + t.Fatalf("multi finalize: resend=%+v diverged=%v", resend, diverged) + } +} + +func TestArgStreamerEscapesControlCharsInInnerJSON(t *testing.T) { + // raw newline inside the inner JSON string (control char) must be escaped + // exactly like EscapeControlChars would do at parse time + doc := "{\"tool_calls\":[{\"id\":\"c4\",\"type\":\"function\",\"function\":{\"name\":\"f\",\"arguments\":\"{\\\"a\\\":\\\"line1\nline2\\\"}\"}}]}" + a := NewArgStreamer() + got, _, _ := feedAll(a, doc, 100) + want := "{\"a\":\"line1\\nline2\"}" + if got != want { + t.Fatalf("control char escaping:\n got %q\nwant %q", got, want) + } +} + +func TestArgStreamerIgnoresPlainProse(t *testing.T) { + a := NewArgStreamer() + if d := a.Feed("just chatting, no tool here"); len(d) != 0 { + t.Fatalf("plain prose produced deltas: %+v", d) + } + if a.Sent() { + t.Fatal("Sent() = true on prose") + } +} + +func TestArgStreamerProsePrefix(t *testing.T) { + // the model opens with a sentence (that mentions the arguments key as + // words) before the actual tool payload + prose := "我来写文件，参数放在 \"arguments\": \"值\" 这个字段里。\n\n" + doc := prose + asDoc + a := NewArgStreamer() + got, head, headSet := feedAll(a, doc, 7) + if got != asInner { + t.Fatalf("args mismatch:\n got %q\nwant %q", got, asInner) + } + if !headSet || head.ID != "call_7f" || head.Name != "write_file" { + t.Fatalf("head delta missing id/name: %+v set=%v", head, headSet) + } + calls := []ToolCall{{ID: "call_7f", Type: "function", Function: FunctionCall{Name: "write_file", Arguments: asInner}}} + deltas, resend, diverged := a.Finalize(calls) + if diverged || len(resend) != 0 || len(deltas) != 0 { + t.Fatalf("clean finalize: deltas=%+v resend=%+v diverged=%v", deltas, resend, diverged) + } +} + +// TestArgStreamerParityWithParseHeal feeds a payload whose inner JSON carries +// an invalid escape (\q) and asserts the streamed text equals the repaired +// parse result — the exact divergence the release fix targets. +func TestArgStreamerParityWithParseHeal(t *testing.T) { + doc := `{"tool_calls":[{"id":"c9","type":"function","function":{"name":"f","arguments":"{\"path\":\"a.txt\",\"content\":\"x\\qy\"}"}}]}` + _, calls, ok := TryParseToolCalls(doc) + if !ok || len(calls) != 1 { + t.Fatalf("parse failed: ok=%v calls=%+v", ok, calls) + } + target := calls[0].Function.Arguments + if !strings.Contains(target, `x\\qy`) || !json.Valid([]byte(target)) { + t.Fatalf("healed target invalid: %q", target) + } + a := NewArgStreamer() + got, _, _ := feedAll(a, doc, 3) + if got != target { + t.Fatalf("stream/parse divergence:\n stream=%q\n target=%q", got, target) + } + deltas, resend, diverged := a.Finalize(calls) + if diverged || len(resend) != 0 || len(deltas) != 0 { + t.Fatalf("clean finalize expected: deltas=%+v resend=%+v diverged=%v", deltas, resend, diverged) + } +} + +func TestTryParseProsePrefix(t *testing.T) { + prose := "好的，我这就创建文件。\n\n" + trailing := "\n\n完成。" + _, calls, ok := TryParseToolCalls(prose + asDoc + trailing) + if !ok || len(calls) != 1 || calls[0].Function.Arguments != asInner { + t.Fatalf("prose prefix parse: ok=%v calls=%+v", ok, calls) + } + remaining, _, ok := TryParseToolCalls(prose + asDoc) + if !ok || strings.TrimSpace(remaining) != strings.TrimSpace(prose) { + t.Fatalf("prose should survive as remaining: %q ok=%v", remaining, ok) + } + if _, _, ok := TryParseToolCalls("纯散文，没有工具调用。"); ok { + t.Fatal("plain prose must not parse as tool calls") + } +} + +func TestTrimHold(t *testing.T) { + cases := map[string]string{ + `abc,`: `abc`, + `abc: `: `abc`, + `abc\\`: `abc\\`, // even backslash run stays + `abc\`: `abc`, // lone dangling backslash held + `abc\n`: `abc\n`, // complete escape pair stays (backslash+n text) + `abc":`: `abc"`, // colon after closing quote held + `normal`: `normal`, + `abc\r\n`: `abc\r\n`, // literal \r\n escape pair text stays + "abc\r\n": "abc", // real CRLF bytes stripped + } + for in, want := range cases { + if got := trimHold(in); got != want { + t.Errorf("trimHold(%q) = %q, want %q", in, got, want) + } + } +} diff --git a/internal/openai/types.go b/internal/openai/types.go index 87adc2e..13543c1 100644 --- a/internal/openai/types.go +++ b/internal/openai/types.go @@ -97,11 +97,26 @@ type Function struct { } type ToolCall struct { + Index *int `json:"index,omitempty"` ID string `json:"id"` Type string `json:"type"` Function FunctionCall `json:"function"` } +// ToolCallDelta is one entry of a streamed tool_calls delta. Index is always +// present; the other fields are emitted only when they carry content. +type ToolCallDelta struct { + Index int `json:"index"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function *FunctionDelta `json:"function,omitempty"` +} + +type FunctionDelta struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` +} + type FunctionCall struct { Name string `json:"name"` Arguments string `json:"arguments"` @@ -172,9 +187,9 @@ type ChatCompletionResponse struct { } type ChoiceDelta struct { - Role string `json:"role,omitempty"` - Content string `json:"content,omitempty"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ToolCalls []ToolCallDelta `json:"tool_calls,omitempty"` } type StreamChoice struct {