Guard streamed tool terminals and retain Usage #46

Merged
weichen merged 9 commits from feature/fast-slow-u3-r04-stream-guard into main 2026-10-19 09:20:09 +00:00
10 changed files with 1149 additions and 54 deletions
+20 -1
View File
@@ -39,7 +39,7 @@ func vocabScript() llm.StreamFunc {
step++
ch <- llm.Event{Kind: llm.KindReasoningDelta, ReasoningDelta: "思考"}
ch <- llm.Event{Kind: llm.KindTextDelta, Text: "正文"}
ch <- llm.Event{Kind: llm.KindToolCallDelta, ToolName: "list_dir", ArgumentsDelta: `{"path":"."}`}
ch <- llm.Event{Kind: llm.KindToolCallDelta, ToolIndex: 2, ToolID: "c1", ToolName: "list_dir", ArgumentsDelta: `{"path":"."}`}
ch <- llm.Event{Kind: llm.KindMessageEnd, Message: llm.Message{
Role: llm.RoleAssistant, StopReason: llm.StopToolCalls,
ToolCalls: []llm.ToolCall{{ID: "c1", Name: "list_dir", Arguments: `{"path":"."}`}},
@@ -54,6 +54,25 @@ func vocabScript() llm.StreamFunc {
}
}
func TestLoopPreservesToolCallDeltaIdentity(t *testing.T) {
reg := tools.WorkspaceTools(t.TempDir())
var got Event
_, err := Run(context.Background(), Config{
Model: "fake", Tools: reg, Stream: vocabScript(), MaxTurns: 4,
OnEvent: func(ev Event) {
if ev.Type == "tool_call_delta" {
got = ev
}
},
}, llm.Message{Role: llm.RoleUser, Content: "run"})
if err != nil {
t.Fatal(err)
}
if got.ToolIndex != 2 || got.CallID != "c1" || got.Tool != "list_dir" || got.Detail != `{"path":"."}` {
t.Fatalf("tool delta lost its provider identity: %#v", got)
}
}
func TestLoopEmitsProcessVocabulary(t *testing.T) {
reg := tools.WorkspaceTools(t.TempDir())
var observed []string
+80
View File
@@ -0,0 +1,80 @@
package llm
import (
"io"
"net/http"
"slices"
)
const maxObservedBodyBytes = 32*1024*1024 + 1
// bodyObserver counts the bytes a reader returns from an HTTP response body
// and whether it reached io.EOF. It never copies or keeps the bytes.
type bodyObserver struct {
src io.Reader
protocol string
framing string
declared int64
uncompressed bool
bytes int64
saturated bool
eof bool
}
// newBodyObserver classifies only what net/http exposes on the response;
// anything it cannot see reliably is recorded as unknown.
func newBodyObserver(resp *http.Response) *bodyObserver {
o := &bodyObserver{src: resp.Body, protocol: "other", framing: "unknown",
declared: resp.ContentLength, uncompressed: resp.Uncompressed}
switch resp.ProtoMajor {
case 1:
o.protocol = "http1"
switch {
case slices.Contains(resp.TransferEncoding, "chunked"):
o.framing = "chunked"
case resp.Uncompressed:
// Transparent decompression hides the wire framing and length.
case resp.ContentLength >= 0:
o.framing = "known_length"
case resp.Close:
o.framing = "close_delimited"
}
case 2:
o.protocol = "http2"
o.framing = "http2_stream"
}
return o
}
func (o *bodyObserver) Read(p []byte) (int, error) {
n, err := o.src.Read(p)
if n > 0 && !o.saturated {
if int64(n) >= maxObservedBodyBytes-o.bytes {
o.bytes = maxObservedBodyBytes
o.saturated = true
} else {
o.bytes += int64(n)
}
}
if err == io.EOF {
o.eof = true
}
return n, err
}
// snapshot must be taken after the last Read has returned.
func (o *bodyObserver) snapshot() *BodyBoundaryObservation {
relation := "unknown"
if o.declared >= 0 && !o.uncompressed && !o.saturated && o.eof {
relation = "mismatch"
if o.bytes == o.declared {
relation = "match"
}
}
return &BodyBoundaryObservation{
ProtocolClass: o.protocol, FramingClass: o.framing,
DeclaredLengthKnown: o.declared >= 0, BytesRead: o.bytes,
BytesSaturated: o.saturated, EOFSeen: o.eof,
LengthRelation: relation, Uncompressed: o.uncompressed,
}
}
+12 -3
View File
@@ -100,10 +100,10 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
resp, err := hc.Do(req)
if err != nil {
if ctx.Err() != nil {
emitAbort(out, "", ctx.Err())
emitAbort(out, "", ctx.Err(), &TerminalObservation{ReadOutcome: "cancelled"})
return
}
emitError(out, "", true, err.Error())
emitError(out, "", true, err.Error(), &TerminalObservation{ReadOutcome: "transport_error"})
return
}
defer resp.Body.Close()
@@ -113,6 +113,7 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
ev := Event{
Kind: KindError,
HTTPStatus: resp.StatusCode,
TerminalObservation: &TerminalObservation{ReadOutcome: "provider_error"},
Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg),
ErrorMessage: msg,
Err: &StreamError{Message: msg, Retryable: StatusRetryable(resp.StatusCode) && !ContextOverflow(msg)},
@@ -125,7 +126,15 @@ func (c *Client) run(ctx context.Context, model string, messages []Message, tool
out <- ev
return
}
for ev := range ParseSSE(ctx, resp.Body) {
observed := newBodyObserver(resp)
for ev := range ParseSSE(ctx, observed) {
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
ev.HTTPStatus = resp.StatusCode
// ParseSSE sends its terminal event only after its last body read.
if ev.TerminalObservation != nil {
ev.TerminalObservation.BodyBoundary = observed.snapshot()
}
}
out <- ev
}
}
+398
View File
@@ -0,0 +1,398 @@
package llm
import (
"bytes"
"compress/gzip"
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"sync/atomic"
"testing"
)
type terminalRoundTripFunc func(*http.Request) (*http.Response, error)
func (f terminalRoundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
func TestClientTerminalObservationHTTPStatus(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Errorf("path = %q", r.URL.Path)
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = w.Write([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":2,\"completion_tokens\":1,\"total_tokens\":3}}\n\ndata: [DONE]\n\n"))
}))
defer server.Close()
client := &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}
var end Event
for ev := range client.Stream(context.Background(), "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.HTTPStatus != http.StatusOK || end.TerminalObservation == nil ||
end.TerminalObservation.ReadOutcome != "done_marker" || !end.StreamComplete || end.Usage.TotalTokens != 3 {
t.Fatalf("successful HTTP terminal lost transport evidence: %#v", end)
}
}
func TestClientTerminalObservationProviderAndTransportFailure(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, "PRIVATE_PROVIDER_BODY_R05", http.StatusTooManyRequests)
}))
defer server.Close()
tests := []struct {
name string
client *Client
status int
outcome string
}{
{"provider-http", &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}, http.StatusTooManyRequests, "provider_error"},
{"transport", &Client{BaseURL: "http://test.invalid", APIKey: "fixture", HTTPClient: &http.Client{Transport: terminalRoundTripFunc(func(*http.Request) (*http.Response, error) {
return nil, errors.New("PRIVATE_TRANSPORT_ERROR_R05")
})}}, 0, "transport_error"},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var end Event
for ev := range tc.client.Stream(context.Background(), "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
end = ev
}
if end.Kind != KindError || end.HTTPStatus != tc.status || end.TerminalObservation == nil ||
end.TerminalObservation.ReadOutcome != tc.outcome {
t.Fatalf("provider/transport failure was not classified: %#v", end)
}
})
}
}
// bodyBoundaryToolFrames is a complete tool call plus Usage without [DONE].
const bodyBoundaryToolFrames = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_PRIVATE_R06\",\"function\":{\"name\":\"read_PRIVATE_R06\",\"arguments\":\"{\\\"path\\\":\\\"PRIVATE_R06.txt\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":2,\"total_tokens\":7}}\n\n"
const bodyBoundaryDone = "data: [DONE]\n\n"
func streamBodyBoundaryEnd(t *testing.T, ctx context.Context, client *Client) Event {
t.Helper()
var end Event
terminals := 0
for ev := range client.Stream(ctx, "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
end = ev
terminals++
} else if ev.TerminalObservation != nil {
t.Fatalf("non-terminal event carried an observation: %#v", ev)
}
}
if terminals != 1 {
t.Fatalf("terminal events = %d, want 1", terminals)
}
if end.TerminalObservation != nil {
raw, err := json.Marshal(end.TerminalObservation)
if err != nil || strings.Contains(string(raw), "PRIVATE") || strings.Contains(string(raw), "127.0.0.1") {
t.Fatalf("unsafe terminal observation: %s err=%v", raw, err)
}
}
return end
}
type bodyBoundaryWant struct {
protocol, framing, relation string
declared, eof, uncompressed bool
bytes int64
}
func requireBodyBoundary(t *testing.T, end Event, want bodyBoundaryWant) {
t.Helper()
if end.TerminalObservation == nil || end.TerminalObservation.BodyBoundary == nil {
t.Fatalf("missing body boundary observation: %#v", end)
}
body := end.TerminalObservation.BodyBoundary
if body.ProtocolClass != want.protocol || body.FramingClass != want.framing ||
body.LengthRelation != want.relation || body.DeclaredLengthKnown != want.declared ||
body.EOFSeen != want.eof || body.Uncompressed != want.uncompressed || body.BytesSaturated {
t.Fatalf("wrong body boundary: %+v want %+v", body, want)
}
if want.bytes >= 0 && body.BytesRead != want.bytes {
t.Fatalf("bytes read = %d, want %d", body.BytesRead, want.bytes)
}
}
func requireUnreleasedTools(t *testing.T, end Event) {
t.Helper()
if end.Kind != KindMessageEnd || end.Err == nil || end.StreamComplete || len(end.Message.ToolCalls) != 0 ||
end.Usage.TotalTokens != 7 || end.HTTPStatus != http.StatusOK {
t.Fatalf("unfinished tool stream was released or lost usage: %#v", end)
}
}
func fixtureResponseClient(resp func() *http.Response) *Client {
return &Client{BaseURL: "http://fixture.invalid", APIKey: "fixture", HTTPClient: &http.Client{
Transport: terminalRoundTripFunc(func(*http.Request) (*http.Response, error) { return resp(), nil }),
}}
}
func TestClientBodyBoundaryShortRead(t *testing.T) {
t.Run("declared-longer-than-body", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Content-Length", strconv.Itoa(len(bodyBoundaryToolFrames)+64))
_, _ = io.WriteString(w, bodyBoundaryToolFrames)
}))
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
requireUnreleasedTools(t, end)
if end.TerminalObservation.ReadOutcome != "read_error" {
t.Fatalf("short body was not a read error: %+v", end.TerminalObservation)
}
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "unknown",
declared: true, bytes: int64(len(bodyBoundaryToolFrames))})
})
t.Run("declared-length-mismatch", func(t *testing.T) {
client := fixtureResponseClient(func() *http.Response {
return &http.Response{StatusCode: http.StatusOK, ProtoMajor: 1, ProtoMinor: 1,
ContentLength: int64(len(bodyBoundaryToolFrames) + 64), Header: http.Header{},
Body: io.NopCloser(strings.NewReader(bodyBoundaryToolFrames))}
})
end := streamBodyBoundaryEnd(t, context.Background(), client)
requireUnreleasedTools(t, end)
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "mismatch",
declared: true, eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
})
}
func TestClientBodyBoundaryUncompressed(t *testing.T) {
var compressed bytes.Buffer
zw := gzip.NewWriter(&compressed)
_, _ = io.WriteString(zw, bodyBoundaryToolFrames)
if err := zw.Close(); err != nil {
t.Fatal(err)
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Content-Encoding", "gzip")
w.Header().Set("Content-Length", strconv.Itoa(compressed.Len()))
_, _ = w.Write(compressed.Bytes())
}))
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
requireUnreleasedTools(t, end)
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "unknown", relation: "unknown",
eof: true, uncompressed: true, bytes: int64(len(bodyBoundaryToolFrames))})
}
// doneThenErrorBody returns the whole stream without io.EOF; any later read fails.
type doneThenErrorBody struct {
data *strings.Reader
reads int
}
func (b *doneThenErrorBody) Read(p []byte) (int, error) {
if b.data.Len() > 0 {
return b.data.Read(p)
}
b.reads++
return 0, errors.New("PRIVATE_READ_AFTER_DONE_R06")
}
func (b *doneThenErrorBody) Close() error { return nil }
func TestClientDoneDoesNotNeedBodyEOF(t *testing.T) {
stream := bodyBoundaryToolFrames + bodyBoundaryDone
body := &doneThenErrorBody{data: strings.NewReader(stream)}
client := fixtureResponseClient(func() *http.Response {
return &http.Response{StatusCode: http.StatusOK, ProtoMajor: 1, ProtoMinor: 1, ContentLength: -1,
TransferEncoding: []string{"chunked"}, Header: http.Header{}, Body: body}
})
end := streamBodyBoundaryEnd(t, context.Background(), client)
if end.Kind != KindMessageEnd || end.Err != nil || !end.StreamComplete || len(end.Message.ToolCalls) != 1 ||
end.Message.StopReason != StopToolCalls || end.Usage.TotalTokens != 7 || end.HTTPStatus != http.StatusOK ||
end.TerminalObservation.ReadOutcome != "done_marker" {
t.Fatalf("legal [DONE] tool terminal was not kept: %#v", end)
}
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "chunked", relation: "unknown",
bytes: int64(len(stream))})
if body.reads != 0 {
t.Fatalf("parser read past [DONE] %d times", body.reads)
}
var standalone Event
for ev := range ParseSSE(context.Background(), strings.NewReader(stream)) {
standalone = ev
}
if standalone.TerminalObservation == nil || standalone.TerminalObservation.BodyBoundary != nil ||
len(standalone.Message.ToolCalls) != 1 || !standalone.StreamComplete {
t.Fatalf("standalone ParseSSE changed: %#v", standalone)
}
}
func TestClientBodyBoundaryChunkedAndHTTP2(t *testing.T) {
flushed := func(stream string) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, stream)
w.(http.Flusher).Flush()
}
}
tests := []struct {
name string
server func() *httptest.Server
wantDone bool
want bodyBoundaryWant
}{
{"known-length", func() *httptest.Server {
stream := bodyBoundaryToolFrames + bodyBoundaryDone
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Content-Length", strconv.Itoa(len(stream)))
_, _ = io.WriteString(w, stream)
}))
}, true, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "match", declared: true, eof: true,
bytes: int64(len(bodyBoundaryToolFrames + bodyBoundaryDone))}},
{"chunked", func() *httptest.Server {
return httptest.NewServer(flushed(bodyBoundaryToolFrames))
}, false, bodyBoundaryWant{protocol: "http1", framing: "chunked", relation: "unknown", eof: true,
bytes: int64(len(bodyBoundaryToolFrames))}},
{"http2", func() *httptest.Server {
server := httptest.NewUnstartedServer(flushed(bodyBoundaryToolFrames))
server.EnableHTTP2 = true
server.StartTLS()
return server
}, false, bodyBoundaryWant{protocol: "http2", framing: "http2_stream", relation: "unknown", eof: true,
bytes: int64(len(bodyBoundaryToolFrames))}},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
server := tc.server()
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
requireBodyBoundary(t, end, tc.want)
if !tc.wantDone {
requireUnreleasedTools(t, end)
return
}
if end.Err != nil || !end.StreamComplete || len(end.Message.ToolCalls) != 1 || end.Usage.TotalTokens != 7 {
t.Fatalf("legal [DONE] tool terminal was not kept: %#v", end)
}
})
}
}
func TestClientBodyBoundaryCloseDelimitedUnknown(t *testing.T) {
t.Run("close-delimited", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Transfer-Encoding", "identity")
w.WriteHeader(http.StatusOK)
_, _ = io.WriteString(w, bodyBoundaryToolFrames)
w.(http.Flusher).Flush()
}))
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
requireUnreleasedTools(t, end)
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "close_delimited", relation: "unknown",
eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
})
t.Run("undeterminable", func(t *testing.T) {
client := fixtureResponseClient(func() *http.Response {
return &http.Response{StatusCode: http.StatusOK, ContentLength: -1, Header: http.Header{},
Body: io.NopCloser(strings.NewReader(bodyBoundaryToolFrames))}
})
end := streamBodyBoundaryEnd(t, context.Background(), client)
requireUnreleasedTools(t, end)
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "other", framing: "unknown", relation: "unknown",
eof: true, bytes: int64(len(bodyBoundaryToolFrames))})
})
}
func TestClientBodyBoundaryFailuresDoNotClaimEOF(t *testing.T) {
t.Run("cancelled", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"PRIVATE_R06\"}}]}\n\n")
w.(http.Flusher).Flush()
<-r.Context().Done()
}))
defer server.Close()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
client := &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()}
var end Event
for ev := range client.Stream(ctx, "fixture", []Message{{Role: RoleUser, Content: "hello"}}, nil, StreamOptions{}) {
if ev.Kind == KindTextDelta {
cancel()
}
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
end = ev
}
}
if end.Kind != KindError || end.Message.StopReason != StopAborted || end.TerminalObservation == nil ||
end.TerminalObservation.ReadOutcome != "cancelled" {
t.Fatalf("cancelled stream was not aborted: %#v", end)
}
if body := end.TerminalObservation.BodyBoundary; body == nil || body.EOFSeen || body.LengthRelation != "unknown" {
t.Fatalf("cancelled stream claimed body EOF: %+v", body)
}
})
t.Run("provider-4xx", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
http.Error(w, "PRIVATE_PROVIDER_BODY_R06", http.StatusBadRequest)
}))
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
if end.Kind != KindError || end.HTTPStatus != http.StatusBadRequest || end.TerminalObservation == nil ||
end.TerminalObservation.ReadOutcome != "provider_error" || end.TerminalObservation.BodyBoundary != nil {
t.Fatalf("provider error carried a body boundary: %#v", end)
}
})
t.Run("redirect-then-short-read", func(t *testing.T) {
var finalHits atomic.Int32
mux := http.NewServeMux()
mux.HandleFunc("/chat/completions", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Location", "/final/chat/completions")
w.WriteHeader(http.StatusTemporaryRedirect)
_, _ = io.WriteString(w, "PRIVATE_REDIRECT_BODY_R06")
})
mux.HandleFunc("/final/chat/completions", func(w http.ResponseWriter, r *http.Request) {
finalHits.Add(1)
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Content-Length", strconv.Itoa(len(bodyBoundaryToolFrames)+64))
_, _ = io.WriteString(w, bodyBoundaryToolFrames)
})
server := httptest.NewServer(mux)
defer server.Close()
end := streamBodyBoundaryEnd(t, context.Background(), &Client{BaseURL: server.URL, APIKey: "fixture", HTTPClient: server.Client()})
if hits := finalHits.Load(); hits != 1 {
t.Fatalf("redirect target hits = %d, want 1", hits)
}
requireUnreleasedTools(t, end)
requireBodyBoundary(t, end, bodyBoundaryWant{protocol: "http1", framing: "known_length", relation: "unknown",
declared: true, bytes: int64(len(bodyBoundaryToolFrames))})
})
}
func TestClientBodyBoundarySaturatesCount(t *testing.T) {
size := int64(maxObservedBodyBytes + 4096)
observer := newBodyObserver(&http.Response{ProtoMajor: 1, ContentLength: size,
Body: io.NopCloser(io.LimitReader(zeroReader{}, size))})
if n, err := io.Copy(io.Discard, observer); err != nil || n != size {
t.Fatalf("observer changed the body: n=%d err=%v", n, err)
}
body := observer.snapshot()
if body.BytesRead != maxObservedBodyBytes || !body.BytesSaturated || !body.EOFSeen ||
body.LengthRelation != "unknown" || !body.DeclaredLengthKnown {
t.Fatalf("saturated body boundary: %+v", body)
}
}
type zeroReader struct{}
func (zeroReader) Read(p []byte) (int, error) {
clear(p)
return len(p), nil
}
+181 -23
View File
@@ -4,6 +4,7 @@ import (
"bufio"
"context"
"encoding/json"
"errors"
"io"
"strings"
)
@@ -39,6 +40,12 @@ type usageWire struct {
} `json:"prompt_tokens_details"`
}
const (
maxObservedToolFragments = 10000
maxObservedToolArgumentBytes = 4*1024*1024 + 1
maxObservedTailEvents = 10000
)
// ParseSSE reads an OpenAI-compatible server-sent event body and emits
// deltas plus one terminal message_end or error event.
func ParseSSE(ctx context.Context, r io.Reader) <-chan Event {
@@ -58,35 +65,92 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
var usage Usage
finished := ""
sawDone := false
sawUsage := false
readOutcome := "clean_eof"
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
if err := json.Unmarshal([]byte(calls[idx].Arguments), &args); err != nil || args == nil {
argumentsValid = false
break
}
}
return &TerminalObservation{
ReadOutcome: readOutcome, LastFrameJSONComplete: lastFrameJSONComplete,
UsagePresent: sawUsage, ToolFragments: toolFragments,
ToolArgumentBytes: toolArgumentBytes, ToolArgumentsValid: argumentsValid,
TailEventKind: tailKind, TailEventCount: reader.tailCount,
}
}
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)
failed.TerminalObservation = observation()
select {
case <-ctx.Done():
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
case out <- failed:
}
}
for {
data, ok := reader.next()
if !ok {
break
}
if ctx.Err() != nil {
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
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 {
emitError(out, content.String(), false, "malformed stream chunk: "+err.Error())
reader.noteTail(badFrameKind(data, err))
readOutcome = "invalid_frame"
lastFrameJSONComplete = false
if sawUsage {
sendTerminalFailure("malformed stream chunk after reported usage")
return
}
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"
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)
emitError(out, content.String(), retryable, msg, observation())
return
}
if chunk.Usage != nil {
usage = chunk.Usage.Usage
sawUsage = true
if chunk.Usage.Details.Cached > 0 {
usage.CacheReadTokens = chunk.Usage.Details.Cached
}
@@ -102,7 +166,8 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
content.WriteString(choice.Delta.Content)
select {
case <-ctx.Done():
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
return
case out <- Event{Kind: KindTextDelta, Text: choice.Delta.Content}:
}
@@ -117,12 +182,23 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
reasoning.WriteString(choice.Delta.Reasoning)
select {
case <-ctx.Done():
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
return
case out <- Event{Kind: KindReasoningDelta, ReasoningDelta: choice.Delta.Reasoning}:
}
}
for _, tc := range choice.Delta.ToolCalls {
if toolFragments < maxObservedToolFragments {
toolFragments++
}
if n := len(tc.Function.Arguments); n > 0 && toolArgumentBytes < maxObservedToolArgumentBytes {
if n > maxObservedToolArgumentBytes-toolArgumentBytes {
toolArgumentBytes = maxObservedToolArgumentBytes
} else {
toolArgumentBytes += n
}
}
cur, ok := calls[tc.Index]
if !ok {
cur = &ToolCall{}
@@ -140,7 +216,8 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
}
select {
case <-ctx.Done():
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
return
case out <- Event{
Kind: KindToolCallDelta,
@@ -153,22 +230,58 @@ func emitParsed(ctx context.Context, r io.Reader, out chan<- Event) {
}
}
if err := sc.Err(); err != nil {
readOutcome = "read_error"
if ctx.Err() != nil {
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
return
}
emitError(out, content.String(), true, err.Error())
if sawUsage {
sendTerminalFailure("stream read failed after reported usage")
return
}
emitError(out, content.String(), true, err.Error(), observation())
return
}
// 取消必须压过正常结束。select 在两个分支都就绪时随机挑选,
// 竞态下会把 message_end 发给已经取消的调用方。
if ctx.Err() != nil {
emitAbort(out, content.String(), ctx.Err())
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
return
}
// 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 (finished == "tool_calls" || finished == "function_call") && len(order) == 0 {
sendTerminalFailure("tool call finish reason without a tool call")
return
}
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
}
// length、空结束原因,或没有 [DONE] 的中断,都不把半截工具调用交给执行层。
if !toolFinishExecutable(finished, sawDone) {
order = nil
}
msg := Message{Role: RoleAssistant, Content: content.String(), StopReason: mapFinish(finished, len(order) > 0)}
if reasoning.Len() > 0 {
@@ -182,16 +295,53 @@ 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}:
readOutcome = "cancelled"
emitAbort(out, content.String(), ctx.Err(), observation())
case out <- Event{Kind: KindMessageEnd, Message: msg, Usage: usage, FinishReason: finished, StreamComplete: sawDone, TerminalObservation: observation()}:
}
}
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 行单独成事件,这样单行流不用空行也能解析。
// tailKind/tailCount 只记最后完成的结构事件类别与饱和计数,不存帧内容。
type sseEventReader struct {
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) {
@@ -220,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:") {
@@ -241,17 +392,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
}
}
@@ -276,8 +426,8 @@ func mapFinish(reason string, hasTools bool) StopReason {
}
}
func emitError(out chan<- Event, partial string, retryable bool, message string) {
out <- Event{
func emitError(out chan<- Event, partial string, retryable bool, message string, observed ...*TerminalObservation) {
event := Event{
Kind: KindError,
Retryable: retryable,
ErrorMessage: message,
@@ -288,14 +438,18 @@ func emitError(out chan<- Event, partial string, retryable bool, message string)
StopReason: StopError,
},
}
if len(observed) > 0 {
event.TerminalObservation = observed[0]
}
out <- event
}
func emitAbort(out chan<- Event, partial string, err error) {
func emitAbort(out chan<- Event, partial string, err error, observed ...*TerminalObservation) {
msg := "aborted"
if err != nil {
msg = err.Error()
}
out <- Event{
event := Event{
Kind: KindError,
Retryable: false,
ErrorMessage: msg,
@@ -306,6 +460,10 @@ func emitAbort(out chan<- Event, partial string, err error) {
StopReason: StopAborted,
},
}
if len(observed) > 0 {
event.TerminalObservation = observed[0]
}
out <- event
}
// StreamError is a provider or transport failure encoded on the event stream.
+335 -6
View File
@@ -2,6 +2,7 @@ package llm
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
@@ -87,13 +88,39 @@ func TestParseSSEToolCallIncrements(t *testing.T) {
if end.Message.StopReason != StopToolCalls || len(end.Message.ToolCalls) != 1 {
t.Fatalf("end %#v", end.Message)
}
if end.TerminalObservation == nil || end.TerminalObservation.ReadOutcome != "done_marker" ||
!end.TerminalObservation.ToolArgumentsValid || end.TerminalObservation.ToolFragments != 3 {
t.Fatalf("valid multi-fragment terminal was not classified: %#v", end.TerminalObservation)
}
tc := end.Message.ToolCalls[0]
if tc.ID != "call_1" || tc.Name != "read_file" || tc.Arguments != `{"path":"a.md"}` {
t.Fatalf("tool call %#v", tc)
}
}
func TestParseSSELengthDropsUnfinishedToolCall(t *testing.T) {
func TestParseSSEToolFinishWithoutCallFails(t *testing.T) {
for _, reason := range []string{"tool_calls", "function_call"} {
t.Run(reason, func(t *testing.T) {
body := strings.Join([]string{
`data: {"choices":[{"delta":{},"finish_reason":"` + reason + `"}],"usage":{"prompt_tokens":12,"completion_tokens":5,"total_tokens":17}}`,
`data: [DONE]`,
"",
}, "\n")
var end Event
for ev := range ParseSSE(context.Background(), strings.NewReader(body)) {
if ev.Kind == KindMessageEnd {
end = ev
}
}
if end.Err == nil || end.Message.StopReason != StopError || end.Usage.TotalTokens != 17 ||
!end.StreamComplete || len(end.Message.ToolCalls) != 0 {
t.Fatalf("empty tool finish must fail and retain Usage: %#v", end)
}
})
}
}
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 +133,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 +149,172 @@ 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 TestTerminalObservationCleanEOFWithoutDone(t *testing.T) {
body := strings.Join([]string{
`data: {"choices":[{"delta":{"content":"PRIVATE_PROMPT_R05","tool_calls":[{"index":0,"id":"call_1","function":{"name":"PRIVATE_TOOL_R05","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.Err == nil || end.StreamComplete || len(end.Message.ToolCalls) != 0 || end.Usage.TotalTokens != 17 {
t.Fatalf("unfinished tool stream was released or unmetered: %#v", end)
}
meta := end.TerminalObservation
if meta == nil || meta.ReadOutcome != "clean_eof" || !meta.LastFrameJSONComplete ||
!meta.UsagePresent || meta.ToolFragments != 1 || meta.ToolArgumentBytes != 2 || !meta.ToolArgumentsValid {
t.Fatalf("missing content-free clean EOF evidence: %#v", meta)
}
raw, err := json.Marshal(meta)
if err != nil {
t.Fatal(err)
}
if strings.Contains(string(raw), "PRIVATE_PROMPT_R05") || strings.Contains(string(raw), "PRIVATE_TOOL_R05") {
t.Fatalf("terminal observation retained model content: %s", raw)
}
}
func TestTerminalObservationMalformedAndReadError(t *testing.T) {
usageAndTool := `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}}` + "\n\n"
tests := []struct {
name string
ctx context.Context
reader io.Reader
outcome string
usage int
frameValid bool
}{
{"invalid-frame", context.Background(), strings.NewReader(usageAndTool + "data: {not-json}\n\n"), "invalid_frame", 17, false},
{"read-error", context.Background(), &errorAfterDataReader{data: strings.NewReader(usageAndTool)}, "read_error", 17, true},
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
tests = append(tests, struct {
name string
ctx context.Context
reader io.Reader
outcome string
usage int
frameValid bool
}{"cancelled", ctx, strings.NewReader(usageAndTool), "cancelled", 0, false})
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
var end Event
for ev := range ParseSSE(tc.ctx, tc.reader) {
if ev.Kind == KindMessageEnd || ev.Kind == KindError {
end = ev
}
}
if end.Err == nil || len(end.Message.ToolCalls) != 0 || end.Usage.TotalTokens != tc.usage {
t.Fatalf("failure lost Usage or released tools: %#v", end)
}
meta := end.TerminalObservation
if meta == nil || meta.ReadOutcome != tc.outcome || meta.LastFrameJSONComplete != tc.frameValid {
t.Fatalf("wrong terminal classification: %#v", meta)
}
})
}
}
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)
}
if end.TerminalObservation == nil || end.TerminalObservation.ToolArgumentsValid ||
end.TerminalObservation.ReadOutcome != "done_marker" {
t.Fatalf("invalid arguments incorrectly classified: %#v", end.TerminalObservation)
}
}
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 +340,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())
@@ -290,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)
}
}
+45 -1
View File
@@ -98,6 +98,44 @@ type Usage struct {
LastTotalTokens int `json:"last_total_tokens,omitempty"`
}
// TerminalObservation records how a stream ended without retaining provider
// content, tool identity, arguments, request headers, or error text. Counts
// saturate at fixed limits; this evidence never authorizes a tool call.
type TerminalObservation struct {
ReadOutcome string `json:"readOutcome"`
LastFrameJSONComplete bool `json:"lastFrameJSONComplete"`
UsagePresent bool `json:"usagePresent"`
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"`
}
// BodyBoundaryObservation records how the client read one HTTP response body.
// It holds only fixed enums, booleans and a saturated byte count. Reaching
// EOF proves only that the body ended, not that the SSE stream completed.
type BodyBoundaryObservation struct {
// ProtocolClass is http1, http2 or other.
ProtocolClass string `json:"protocolClass"`
// FramingClass is known_length, chunked, http2_stream, close_delimited or unknown.
FramingClass string `json:"framingClass"`
DeclaredLengthKnown bool `json:"declaredLengthKnown"`
BytesRead int64 `json:"bytesRead"`
BytesSaturated bool `json:"bytesSaturated"`
EOFSeen bool `json:"eofSeen"`
// LengthRelation is match, mismatch or unknown.
LengthRelation string `json:"lengthRelation"`
Uncompressed bool `json:"uncompressed"`
}
// Event is one streamed update. Failures are KindError plus a final assistant
// message whose StopReason is error or aborted. Stream implementations do not panic.
type Event struct {
@@ -113,10 +151,16 @@ type Event struct {
// FinishReason 是供应商原始结束原因(stop、length、tool_calls 等)。
// 空表示这条事件不携带结束原因。StopReason 的映射不受此字段影响。
FinishReason string
// StreamComplete records whether the OpenAI-compatible SSE stream supplied
// its explicit [DONE] marker. Tool calls require this before execution.
StreamComplete bool
// TerminalObservation is present only on a terminal/error event. It does
// not change StreamComplete or the tool-execution safety gate.
TerminalObservation *TerminalObservation
Err error
Retryable bool
ErrorMessage string
// HTTPStatus 是提供方的 HTTP 状态。0 表示这条错误不是 HTTP 响应。
// HTTPStatus 是提供方的 HTTP 状态;0 表示尚无 HTTP 响应。
HTTPStatus int
// RetryAfter is the provider's requested wait. HasRetryAfter distinguishes
// a real zero from a missing or unparseable header.
+11 -2
View File
@@ -38,8 +38,10 @@ type Event struct {
Message llm.Message
Tool string
Text string
// ToolIndex preserves the provider's call slot for tool_call_delta.
ToolIndex int
// CallID 在 tool_execution_start / tool_execution_end 时携带工具调用的
// ID:并行执行下同回合多个调用并发进出,消费方按 ID 配对而非按顺序。
// ID;tool_call_delta 也携带同一 ID。消费方按 ID 配对而非按顺序。
CallID string
// Detail carries tool arguments/results or reasoning deltas for
// process-visualization consumers; log writers may ignore it.
@@ -354,6 +356,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
}
@@ -529,7 +537,8 @@ func consume(ctx context.Context, cfg Config, messages []llm.Message) (llm.Event
case llm.KindReasoningDelta:
emit(cfg, Event{Type: "reasoning_delta", Text: ev.ReasoningDelta})
case llm.KindToolCallDelta:
emit(cfg, Event{Type: "tool_call_delta", Tool: ev.ToolName, Detail: ev.ArgumentsDelta})
emit(cfg, Event{Type: "tool_call_delta", ToolIndex: ev.ToolIndex,
CallID: ev.ToolID, Tool: ev.ToolName, Detail: ev.ArgumentsDelta})
case llm.KindError:
return ev, ev.Err
case llm.KindMessageEnd:
+49
View File
@@ -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