diff --git a/runtime/llm/sse.go b/runtime/llm/sse.go index ab2bff1..7ae4c35 100644 --- a/runtime/llm/sse.go +++ b/runtime/llm/sse.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "io" "strings" ) @@ -42,6 +43,7 @@ type usageWire struct { const ( maxObservedToolFragments = 10000 maxObservedToolArgumentBytes = 4*1024*1024 + 1 + maxObservedTailEvents = 10000 ) // ParseSSE reads an OpenAI-compatible server-sent event body and emits @@ -68,7 +70,14 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { lastFrameJSONComplete := false toolFragments := 0 toolArgumentBytes := 0 + sc := bufio.NewScanner(r) + sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) + reader := &sseEventReader{sc: sc} observation := func() *TerminalObservation { + tailKind := reader.tailKind + if tailKind == "" { + tailKind = "unknown" + } argumentsValid := len(order) > 0 for _, idx := range order { var args map[string]json.RawMessage @@ -81,6 +90,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { ReadOutcome: readOutcome, LastFrameJSONComplete: lastFrameJSONComplete, UsagePresent: sawUsage, ToolFragments: toolFragments, ToolArgumentBytes: toolArgumentBytes, ToolArgumentsValid: argumentsValid, + TailEventKind: tailKind, TailEventCount: reader.tailCount, } } sendTerminalFailure := func(why string) { @@ -97,9 +107,6 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { case out <- failed: } } - sc := bufio.NewScanner(r) - sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) - reader := &sseEventReader{sc: sc} for { data, ok := reader.next() if !ok { @@ -111,12 +118,14 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { return } if data == "[DONE]" { + reader.noteTail("done_literal") sawDone = true readOutcome = "done_marker" break } var chunk sseChunk if err := json.Unmarshal([]byte(data), &chunk); err != nil { + reader.noteTail(badFrameKind(data, err)) readOutcome = "invalid_frame" lastFrameJSONComplete = false if sawUsage { @@ -126,6 +135,7 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) { emitError(out, content.String(), false, "malformed stream chunk: "+err.Error(), observation()) return } + reader.noteTail("json_frame") lastFrameJSONComplete = true if chunk.Error != nil && chunk.Error.Message != "" { readOutcome = "provider_error" @@ -300,9 +310,38 @@ func terminalStreamFailure(message Message, usage Usage, reason string, complete // sseEventReader 按空行切事件,并把同一事件里的多行 data: 用换行拼回。 // 已经是完整 JSON 的 data 行单独成事件,这样单行流不用空行也能解析。 +// tailKind/tailCount 只记最后完成的结构事件类别与饱和计数,不存帧内容。 type sseEventReader struct { - sc *bufio.Scanner - carry string + sc *bufio.Scanner + carry string + tailKind string + tailCount int +} + +func (r *sseEventReader) noteTail(kind string) { + if r.tailCount < maxObservedTailEvents { + r.tailCount++ + } + r.tailKind = kind +} + +// badFrameKind classifies a data event that did not decode as a chunk using +// only the decoder's error category and position, never its text or bytes. +func badFrameKind(data string, err error) string { + if strings.TrimSpace(data) == "" { + return "empty_data" + } + var syntax *json.SyntaxError + if !errors.As(err, &syntax) { + return "unknown" + } + if syntax.Offset == int64(len(data)) { + var probe json.RawMessage + if errors.Is(json.NewDecoder(strings.NewReader(data)).Decode(&probe), io.ErrUnexpectedEOF) { + return "partial_frame" + } + } + return "unknown_non_json" } func (r *sseEventReader) next() (string, bool) { @@ -331,6 +370,7 @@ func (r *sseEventReader) next() (string, bool) { } trimmed := strings.TrimSpace(line) if strings.HasPrefix(trimmed, ":") { + r.noteTail("comment") continue } if !strings.HasPrefix(trimmed, "data:") { diff --git a/runtime/llm/sse_test.go b/runtime/llm/sse_test.go index 1da4afb..22b94c0 100644 --- a/runtime/llm/sse_test.go +++ b/runtime/llm/sse_test.go @@ -541,3 +541,81 @@ func TestParseSSEReasoningDelta(t *testing.T) { t.Fatalf("end content %q must not include reasoning", end.Message.Content) } } + +// collectTerminal drains one stream and fails unless it carries exactly one +// terminal message_end or error event. +func collectTerminal(t *testing.T, ch <-chan Event) Event { + t.Helper() + var end Event + terminals := 0 + for ev := range ch { + if ev.Kind == KindMessageEnd || ev.Kind == KindError { + terminals++ + end = ev + } + } + if terminals != 1 { + t.Fatalf("stream emitted %d terminal events, want 1", terminals) + } + return end +} + +func TestSSETailStructure(t *testing.T) { + toolFrames := `data: {"choices":[{"delta":{"content":"PRIVATE_TAIL_R06","tool_calls":[{"index":0,"id":"call_1","function":{"name":"PRIVATE_TAIL_R06","arguments":"{}"}}]}}]}` + "\n\n" + + `data: {"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}` + "\n\n" + textFrame := `data: {"choices":[{"delta":{"content":"PRIVATE_TAIL_R06"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}` + "\n\n" + tests := []struct { + name string + body string + kind string + count int + outcome string + complete bool + tools int + usage int + wantError bool + }{ + {"done", toolFrames + "data: [DONE]\n\n", "done_literal", 3, "done_marker", true, 1, 17, false}, + {"direct-eof", toolFrames, "json_frame", 2, "clean_eof", false, 0, 17, true}, + {"comment-then-eof", toolFrames + ": PRIVATE_TAIL_R06\n\n", "comment", 3, "clean_eof", false, 0, 17, true}, + {"empty-data", toolFrames + "data:\n\n", "empty_data", 3, "invalid_frame", false, 0, 17, true}, + {"unknown-non-json", toolFrames + "data: NOT_DONE\n\n", "unknown_non_json", 3, "invalid_frame", false, 0, 17, true}, + {"partial-frame", toolFrames + "data: {", "partial_frame", 3, "invalid_frame", false, 0, 17, true}, + {"plain-text-eof", textFrame, "json_frame", 1, "clean_eof", false, 0, 5, false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + end := collectTerminal(t, ParseSSE(context.Background(), strings.NewReader(tc.body))) + got := end.TerminalObservation + if got == nil || got.TailEventKind != tc.kind || got.TailEventCount != tc.count || got.ReadOutcome != tc.outcome { + t.Fatalf("wrong tail observation: %+v", got) + } + if end.StreamComplete != tc.complete || end.Usage.TotalTokens != tc.usage { + t.Fatalf("tail changed completion or lost Usage: complete=%v usage=%+v", end.StreamComplete, end.Usage) + } + if len(end.Message.ToolCalls) != tc.tools { + t.Fatalf("released %d tool calls, want %d", len(end.Message.ToolCalls), tc.tools) + } + if (end.Err != nil) != tc.wantError { + t.Fatalf("terminal error = %v, want error %v", end.Err, tc.wantError) + } + raw, err := json.Marshal(got) + if err != nil || strings.Contains(string(raw), "PRIVATE_TAIL_R06") { + t.Fatalf("unsafe tail observation: %s err=%v", raw, err) + } + }) + } +} + +func TestSSETailCountSaturates(t *testing.T) { + body := `data: {"choices":[{"delta":{"content":"x"},"finish_reason":"stop"}]}` + "\n\n" + + strings.Repeat(":\n", 10050) + end := collectTerminal(t, ParseSSE(context.Background(), strings.NewReader(body))) + got := end.TerminalObservation + if got == nil || got.TailEventKind != "comment" || got.TailEventCount != 10000 || got.ReadOutcome != "clean_eof" { + t.Fatalf("tail count did not saturate: %+v", got) + } + if end.Err != nil || end.Message.Content != "x" || len(end.Message.ToolCalls) != 0 { + t.Fatalf("comments changed the parsed message: %+v", end) + } +} diff --git a/runtime/llm/types.go b/runtime/llm/types.go index 2f26c3f..1499deb 100644 --- a/runtime/llm/types.go +++ b/runtime/llm/types.go @@ -108,6 +108,12 @@ type TerminalObservation struct { ToolFragments int `json:"toolFragments"` ToolArgumentBytes int `json:"toolArgumentBytes"` ToolArgumentsValid bool `json:"toolArgumentsValid"` + // TailEventKind is the structural class of the last SSE event the parser + // completed: done_literal, json_frame, empty_data, comment, + // unknown_non_json, partial_frame or unknown. TailEventCount counts those + // events and saturates at 10000. Neither marks the stream complete. + TailEventKind string `json:"tailEventKind"` + TailEventCount int `json:"tailEventCount"` // BodyBoundary is set only by Client on a response it received; a bare // ParseSSE call leaves it nil. BodyBoundary *BodyBoundaryObservation `json:"bodyBoundary,omitempty"`