From 94afd71f00bc919dae7baccd699abe2a11077648 Mon Sep 17 00:00:00 2001 From: weichen Date: Wed, 14 Oct 2026 02:09:44 +0000 Subject: [PATCH] Reject incomplete streamed tool calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Why: a model tool finish could execute partial JSON or an unterminated SSE stream, while post-Usage errors could lose their charge. Evidence: missing-DONE, malformed-JSON, identity, non-object and post-Usage RED→GREEN tests; GOWORK=off go vet ./... and go test -count=1 ./... passed. Impact: invalid tool terminals fail without execution or automatic retry and retain returned Usage; valid text and multi-part tool streams keep working. Refs: runtime/llm/sse.go; runtime/loop.go --- runtime/llm/sse.go | 75 ++++++++++++++++--- runtime/llm/sse_test.go | 161 ++++++++++++++++++++++++++++++++++++++-- runtime/llm/types.go | 9 ++- runtime/loop.go | 18 +++-- runtime/loop_test.go | 49 ++++++++++++ 5 files changed, 288 insertions(+), 24 deletions(-) diff --git a/runtime/llm/sse.go b/runtime/llm/sse.go index 82ee100..07e7db3 100644 --- a/runtime/llm/sse.go +++ b/runtime/llm/sse.go @@ -58,6 +58,19 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { var usage Usage finished := "" sawDone := false + sawUsage := false + sendTerminalFailure := func(why string) { + message := Message{Role: RoleAssistant, Content: content.String(), StopReason: StopError} + if reasoning.Len() > 0 { + message.Thinking = []ThinkingBlock{{Thinking: reasoning.String()}} + } + failed := terminalStreamFailure(message, usage, finished, sawDone, why) + select { + case <-ctx.Done(): + emitAbort(out, content.String(), ctx.Err()) + case out <- failed: + } + } sc := bufio.NewScanner(r) sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) reader := &sseEventReader{sc: sc} @@ -76,17 +89,26 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { } var chunk sseChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { + if sawUsage { + sendTerminalFailure("malformed stream chunk after reported usage") + return + } emitError(out, content.String(), false, "malformed stream chunk: "+err.Error()) return } if chunk.Error != nil && chunk.Error.Message != "" { msg := chunk.Error.Message + if sawUsage { + sendTerminalFailure("provider error after reported usage: " + msg) + return + } retryable := CapacityLimited(msg) && !ContextOverflow(msg) emitError(out, content.String(), retryable, msg) return } if chunk.Usage != nil { usage = chunk.Usage.Usage + sawUsage = true if chunk.Usage.Details.Cached > 0 { usage.CacheReadTokens = chunk.Usage.Details.Cached } @@ -157,6 +179,10 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { emitAbort(out, content.String(), ctx.Err()) return } + if sawUsage { + sendTerminalFailure("stream read failed after reported usage") + return + } emitError(out, content.String(), true, err.Error()) return } @@ -166,9 +192,34 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { emitAbort(out, content.String(), ctx.Err()) return } - // length、空结束原因，或没有 [DONE] 的中断，都不把半截工具调用交给执行层。 - if !toolFinishExecutable(finished, sawDone) { - order = nil + // A finish reason alone does not prove the tool arguments arrived intact. + // Preserve returned usage when rejecting a partial terminal, so the host + // can charge this call and stop rather than treating it as a free retry. + if len(order) > 0 { + problem := "" + switch { + case !sawDone: + problem = "tool call stream ended before [DONE]" + case !toolFinishExecutable(finished, sawDone): + problem = "tool call ended without a supported finish reason" + default: + for _, idx := range order { + call := calls[idx] + var args map[string]json.RawMessage + if call.ID == "" || call.Name == "" { + problem = "tool call identity is incomplete" + break + } + if err := json.Unmarshal([]byte(call.Arguments), &args); err != nil || args == nil { + problem = "tool call arguments are not a JSON object" + break + } + } + } + if problem != "" { + sendTerminalFailure(problem) + return + } } msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)} if reasoning.Len() > 0 { @@ -183,10 +234,17 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { select { case <-ctx.Done(): emitAbort(out, content.String(), ctx.Err()) - case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished}: + case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished, StreamComplete: sawDone}: } } +func terminalStreamFailure(message Message, usage Usage, reason string, complete bool, why string) Event { + err := &StreamError{Message: why, Retryable: false} + return Event{Kind: KindMessageEnd, Message: message, Usage: usage, + FinishReason: reason, StreamComplete: complete, Err: err, + ErrorMessage: why, Retryable: false} +} + // sseEventReader 按空行切事件，并把同一事件里的多行 data: 用换行拼回。 // 已经是完整 JSON 的 data 行单独成事件，这样单行流不用空行也能解析。 type sseEventReader struct { @@ -241,17 +299,16 @@ func (r *sseEventReader) next() (string, bool) { // toolFinishExecutable 只在供应商明确以工具调用或正常 stop 结束时执行。 // length、缺失的结束原因，以及没有 [DONE] 的断开，都视为未完成。 func toolFinishExecutable(reason string, sawDone bool) bool { + if !sawDone { + return false + } switch reason { case "tool_calls", "function_call": return true case "stop": return true - case "length": - return false - case "": - return false default: - return sawDone && reason != "error" && reason != "content_filter" + return false } } diff --git a/runtime/llm/sse_test.go b/runtime/llm/sse_test.go index 5b943bc..2dc3cb9 100644 --- a/runtime/llm/sse_test.go +++ b/runtime/llm/sse_test.go @@ -93,7 +93,7 @@ func TestParseSSEToolCallIncrements(t *testing.T) { } } -func TestParseSSELengthDropsUnfinishedToolCall(t *testing.T) { +func TestParseSSELengthFailsUnfinishedToolCall(t *testing.T) { body := strings.Join([]string{ `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}`, `data: {"choices":[{"delta":{},"finish_reason":"length"}]}`, @@ -106,15 +106,15 @@ func TestParseSSELengthDropsUnfinishedToolCall(t *testing.T) { end = ev } } - if end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 { - t.Fatalf("truncated tool call must not be executable: %#v", end.Message) + if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil { + t.Fatalf("truncated tool call must fail without execution: %#v", end) } if end.FinishReason != "length" { t.Fatalf("finish %q", end.FinishReason) } } -func TestParseSSEUncleanEndDropsToolCall(t *testing.T) { +func TestParseSSEUncleanEndFailsToolCall(t *testing.T) { body := `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"read_file","arguments":"{\"path\":"}}]}}]}` var end Event for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { @@ -122,8 +122,97 @@ func TestParseSSEUncleanEndDropsToolCall(t *testing.T) { end = ev } } - if end.Kind != KindMessageEnd || end.Message.StopReason != StopEnd || len(end.Message.ToolCalls) != 0 { - t.Fatalf("unclean stream must not execute tools: %#v", end) + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil { + t.Fatalf("unclean stream must fail without executing tools: %#v", end) + } +} + +func TestParseSSEToolFinishWithoutDone(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + if ev.Kind == KindMessageEnd { + end = ev + } + } + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || + len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable || + end.Usage.TotalTokens != 17 { + t.Fatalf("unsafe tool call from unfinished stream: %#v", end) + } +} + +func TestParseSSEInvalidToolArguments(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"publish_work_checkpoint","arguments":"{\"findings\":"}}]}}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + `data: [DONE]`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + if ev.Kind == KindMessageEnd { + end = ev + } + } + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || + len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable || + end.Usage.TotalTokens != 17 { + t.Fatalf("unsafe tool call with invalid JSON arguments: %#v", end) + } +} + +func TestParseSSERejectsMissingToolIdentity(t *testing.T) { + for _, tc := range []struct{ name, firstChunk string }{ + {"missing-id", `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"echo","arguments":"{}"}}]}}]}`}, + {"missing-name", `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"arguments":"{}"}}]}}]}`}, + } { + t.Run(tc.name, func(t *testing.T) { + body := strings.Join([]string{ + tc.firstChunk, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + `data: [DONE]`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + if ev.Kind == KindMessageEnd { + end = ev + } + } + if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil { + t.Fatalf("missing tool identity reached execution: %#v", end) + } + }) + } +} + +func TestParseSSERejectsNonObjectToolArguments(t *testing.T) { + for _, tc := range []struct{ name, arguments string }{ + {"array", "[]"}, + {"null", "null"}, + } { + t.Run(tc.name, func(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"` + tc.arguments + `"}}]}}]}`, + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + `data: [DONE]`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + if ev.Kind == KindMessageEnd { + end = ev + } + } + if end.Message.StopReason != StopError || len(end.Message.ToolCalls) != 0 || end.Err == nil { + t.Fatalf("non-object tool arguments reached execution: %#v", end) + } + }) } } @@ -149,6 +238,66 @@ func TestParseSSEMalformedIsError(t *testing.T) { } } +func TestParseSSEMalformedChunkAfterUsageKeepsUsage(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + `data: {not-json}`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + end = ev + } + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || + len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable || + end.Usage.TotalTokens != 17 { + t.Fatalf("malformed chunk dropped already returned usage: %#v", end) + } +} + +func TestParseSSEProviderErrorAfterUsageKeepsUsage(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + `data: {"error":{"message":"server is at capacity"}}`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), strings.NewReader(body)) { + end = ev + } + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || + len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable || + end.Usage.TotalTokens != 17 { + t.Fatalf("provider error after usage retried or lost its charge: %#v", end) + } +} + +type errorAfterDataReader struct{ data *strings.Reader } + +func (r *errorAfterDataReader) Read(p []byte) (int, error) { + n, err := r.data.Read(p) + if err == io.EOF { + return 0, io.ErrUnexpectedEOF + } + return n, err +} + +func TestParseSSEReaderErrorAfterUsageKeepsUsage(t *testing.T) { + body := strings.Join([]string{ + `data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"echo","arguments":"{}"}}]}}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`, + "", + }, "\n\n") + var end Event + for ev := range ParseSSE(context.Background(), &errorAfterDataReader{data: strings.NewReader(body)}) { + end = ev + } + if end.Kind != KindMessageEnd || end.Message.StopReason != StopError || + len(end.Message.ToolCalls) != 0 || end.Err == nil || end.Retryable || + end.Usage.TotalTokens != 17 { + t.Fatalf("reader error after usage retried or lost its charge: %#v", end) + } +} + func TestParseSSECancel(t *testing.T) { pr, pw := io.Pipe() ctx, cancel := context.WithCancel(context.Background()) diff --git a/runtime/llm/types.go b/runtime/llm/types.go index 8d5ea83..dbe1993 100644 --- a/runtime/llm/types.go +++ b/runtime/llm/types.go @@ -113,9 +113,12 @@ type Event struct { // FinishReason 是供应商原始结束原因（stop、length、tool_calls 等）。 // 空表示这条事件不携带结束原因。StopReason 的映射不受此字段影响。 FinishReason string - Err error - Retryable bool - ErrorMessage string + // StreamComplete records whether the OpenAI-compatible SSE stream supplied + // its explicit [DONE] marker. Tool calls require this before execution. + StreamComplete bool + Err error + Retryable bool + ErrorMessage string // HTTPStatus 是提供方的 HTTP 状态。0 表示这条错误不是 HTTP 响应。 HTTPStatus int // RetryAfter is the provider's requested wait. HasRetryAfter distinguishes diff --git a/runtime/loop.go b/runtime/loop.go index 980a52f..f4b45fc 100644 --- a/runtime/loop.go +++ b/runtime/loop.go @@ -72,8 +72,8 @@ type Config struct { // SerialTools 是并行执行下的顺序敏感工具名集合：这些调用按模型给出的 // 原顺序在调用方 goroutine 逐个执行（互斥且保序），其余调用并发执行。 SerialTools map[string]bool - GetSteering func() []llm.Message - GetFollowUp func() []llm.Message + GetSteering func() []llm.Message + GetFollowUp func() []llm.Message // OnEvent 观察循环事件。同一次 Run 内不会并发调用 // （与 BeforeTool、AfterTool 共用一把锁串行）。 OnEvent func(Event) @@ -227,9 +227,9 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { } cfg.traceCtx = newTraceContext() endRun := startSpan(cfg, "invoke_agent "+agentName(cfg), map[string]string{ - "gen_ai.system": providerName(cfg), - "gen_ai.agent.name": agentName(cfg), - "gen_ai.request.model": cfg.Model, + "gen_ai.system": providerName(cfg), + "gen_ai.agent.name": agentName(cfg), + "gen_ai.request.model": cfg.Model, }, true) defer endRun() var messages []llm.Message @@ -354,6 +354,12 @@ func Run(ctx context.Context, cfg Config, user llm.Message) (Result, error) { usage.LastPromptTokens = end.Usage.PromptTokens usage.LastCompletionTokens = end.Usage.CompletionTokens usage.LastTotalTokens = end.Usage.TotalTokens + if end.Err != nil { + cfg.logger().Error("runtime terminal stream failed", end.Err) + emitUsage(cfg, usage) + return Result{Text: last.Content, Messages: messages, StopReason: llm.StopError, + Usage: usage, Err: end.Err, Retryable: false}, nil + } if err := persist(cfg, last); err != nil { return Result{}, err } @@ -1156,7 +1162,7 @@ func drain(fn func() []llm.Message) []llm.Message { func emitUsage(cfg Config, usage llm.Usage) { emit(cfg, Event{ - Type: "usage", + Type: "usage", PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens, TotalTokens: usage.TotalTokens, LastPromptTokens: usage.LastPromptTokens, LastCompletionTokens: usage.LastCompletionTokens, LastTotalTokens: usage.LastTotalTokens, CacheReadTokens: usage.CacheReadTokens, CacheWriteTokens: usage.CacheWriteTokens, diff --git a/runtime/loop_test.go b/runtime/loop_test.go index e9bd9d9..ae8864c 100644 --- a/runtime/loop_test.go +++ b/runtime/loop_test.go @@ -2,6 +2,7 @@ package runtime import ( "context" + "errors" "path/filepath" "strings" "sync" @@ -12,6 +13,54 @@ import ( "agent-native/runtime/tools" ) +func TestRunTerminalToolFailureKeepsUsage(t *testing.T) { + hits, calls := 0, 0 + reg := tools.NewRegistry(t.TempDir()) + reg.Register(hitTool{hits: &hits}) + stream := func(context.Context, string, []llm.Message, []llm.ToolDecl, llm.StreamOptions) <-chan llm.Event { + calls++ + out := make(chan llm.Event, 1) + out <- llm.Event{ + Kind: llm.KindMessageEnd, + Message: llm.Message{Role: llm.RoleAssistant, StopReason: llm.StopError}, + Usage: llm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17}, + Err: errors.New("incomplete tool call"), + } + close(out) + return out + } + res, err := Run(context.Background(), Config{Model: "fixture", Tools: reg, Stream: stream}, + llm.Message{Role: llm.RoleUser, Content: "organize materials"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || res.StopReason != llm.StopError || res.Usage.TotalTokens != 17 || + hits != 0 || calls != 1 { + t.Fatalf("terminal failure retried, executed a tool, or lost usage: result=%#v hits=%d calls=%d", res, hits, calls) + } +} + +func TestRunTerminalErrorCannotMasqueradeAsSuccess(t *testing.T) { + stream := func(context.Context, string, []llm.Message, []llm.ToolDecl, llm.StreamOptions) <-chan llm.Event { + out := make(chan llm.Event, 1) + out <- llm.Event{Kind: llm.KindMessageEnd, + Message: llm.Message{Role: llm.RoleAssistant, Content: "incomplete", StopReason: llm.StopEnd}, + Usage: llm.Usage{PromptTokens: 12, CompletionTokens: 5, TotalTokens: 17}, + Err: errors.New("terminal failed"), + } + close(out) + return out + } + res, err := Run(context.Background(), Config{Model: "fixture", Stream: stream}, + llm.Message{Role: llm.RoleUser, Content: "go"}) + if err != nil { + t.Fatal(err) + } + if res.Err == nil || res.StopReason != llm.StopError || res.Usage.TotalTokens != 17 { + t.Fatalf("error-bearing terminal was accepted as success: %#v", res) + } +} + func script(turns ...llm.Message) llm.StreamFunc { var mu sync.Mutex i := 0