fix(cli): keep stream batches honest
Validate known events before delivery and save cursors only after consumption. Keep unknown kinds quiet and classify stream failures.
12345678910111213141516171819202122232425262728{ "malformed_batches": { "malformed JSON": "{\"cursor\":8,", "missing cursor": "{\"events\":[]}", "null cursor": "{\"cursor\":null,\"events\":[]}", "fractional cursor": "{\"cursor\":8.5,\"events\":[]}", "negative cursor": "{\"cursor\":-1,\"events\":[]}", "overflowing cursor": "{\"cursor\":9223372036854775808,\"events\":[]}", "regressing cursor": "{\"cursor\":6,\"events\":[]}", "missing events": "{\"cursor\":8}", "null events": "{\"cursor\":8,\"events\":null}", "missing event type": "{\"cursor\":8,\"events\":[{\"text\":\"not an assistant\"}]}", "empty event type": "{\"cursor\":8,\"events\":[{\"type\":\"\",\"text\":\"not an assistant\"}]}", "nonstring event type": "{\"cursor\":8,\"events\":[{\"type\":42}]}", "partial batch": "{\"cursor\":8,\"events\":[{\"type\":\"text\",\"text\":\"must not escape\"},{\"type\":\"message\",\"role\":\"assistant\",\"text\":false}]}", "missing known field": "{\"cursor\":8,\"events\":[{\"type\":\"message\",\"role\":\"assistant\"}]}", "missing identity": "{\"cursor\":8,\"events\":[{\"type\":\"turn_completed\",\"turnId\":\"\"}]}", "negative transcript offset": "{\"cursor\":8,\"events\":[{\"type\":\"committed\",\"seq\":-1}]}", "tool args nonstring": "{\"cursor\":8,\"events\":[{\"type\":\"tool\",\"callId\":\"c\",\"name\":\"python\",\"args\":false,\"result\":\"\"}]}", "missing user source": "{\"cursor\":8,\"events\":[{\"type\":\"user\",\"text\":\"hello\",\"triggeredAt\":\"\"}]}", "missing user trigger": "{\"cursor\":8,\"events\":[{\"type\":\"user\",\"text\":\"hello\",\"source\":\"chat\"}]}", "invalid usage count": "{\"cursor\":8,\"events\":[{\"type\":\"usage\",\"model\":\"test\",\"recordedAt\":1,\"promptTokens\":-1}]}", "missing progress": "{\"cursor\":8,\"events\":[{\"type\":\"tool_progress\"}]}", "invalid member identity": "{\"cursor\":8,\"events\":[{\"type\":\"turn_membership\",\"turnId\":\"turn\",\"submissionIds\":[\"\"]}]}", "invalid thinking elapsed": "{\"cursor\":8,\"events\":[{\"type\":\"thinking\",\"text\":\"thought\",\"elapsedMs\":1.5}]}", "compaction missing summary": "{\"cursor\":8,\"events\":[{\"type\":\"compacted\",\"evicted\":2}]}" }}3 unmodified lines4567891011127 unmodified lines202122232324252627281 unmodified line30313233333434353637383915 unmodified lines55565758596061626364636465666768696870717271737475767735 unmodified lines11311411511611711811912012125 unmodified lines3 unmodified lines// timing cannot guarantee, including a cleanup response after caller cancellation.import ( "context" "encoding/json" "errors" "fmt" "net" "net/http" "net/http/httptest"7 unmodified linesfunc promptTestConnection(server *httptest.Server) *daemon.Connection { return daemon.NewConnection(daemon.ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, "")}
func TestPromptRejectsAnOlderDaemonBeforeCreatingASession(t *testing.T) {func TestPromptRequiresSubmissionCancellationBeforeCreatingASession(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { if request.URL.Path != "/health" { t.Errorf("unsupported prompt issued %s", request.URL.Path) }1 unmodified line })) defer server.Close() service := Service{Connect: func(context.Context) (*daemon.Connection, error) { return promptTestConnection(server), nil }} _, err := service.RunPrompt(t.Context(), PromptOptions{Prompt: "hello"}) if err == nil || !strings.Contains(err.Error(), "needs an update") { t.Fatalf("old daemon result: %v", err) if _, ok := errors.AsType[*daemon.UpgradeRequiredError](err); !ok { t.Fatalf("missing required capability did not prevent submission: %v", err) }}
func TestPromptFollowsCombinedSubmissionThroughRetryAndCompletion(t *testing.T) {15 unmodified lines submitted <- payload.SubmissionID writer.WriteHeader(http.StatusAccepted) _ = json.NewEncoder(writer).Encode(map[string]any{"ok": true, "queued": true, "operationId": payload.OperationID}) case "/sessions/test/stream": writer.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprint(writer, "data: {\"cursor\":0,\"events\":[{\"type\":\"reset\"}]}\n\n") writer.(http.Flusher).Flush() id := <-submitted events := []map[string]any{ {"type": "turn_membership", "turnId": "combined", "submissionIds": []string{id, "other"}}, {"type": "user", "clientId": "another-client", "text": "other contribution"}, {"type": "message", "text": "before retry"}, {"type": "user", "clientId": "another-client", "source": "chat", "triggeredAt": "", "text": "other contribution"}, {"type": "message", "role": "assistant", "text": "before retry"}, {"type": "error", "text": "temporary failure"}, {"type": "retry"}, {"type": "compacted", "evicted": 1, "summary": "summary"}, {"type": "message", "text": "combined final answer"}, {"type": "message", "role": "assistant", "text": "combined final answer"}, {"type": "turn_completed", "turnId": "combined"}, {"type": "turn_membership", "turnId": "later", "submissionIds": []string{"other"}}, {"type": "message", "text": "later answer"}, {"type": "message", "role": "assistant", "text": "later answer"}, } page, _ := json.Marshal(map[string]any{"cursor": 1, "events": events}) _, _ = fmt.Fprintf(writer, "data: %s\n\n", page) writer.(http.Flusher).Flush()35 unmodified lines writer.WriteHeader(http.StatusAccepted) _ = json.NewEncoder(writer).Encode(map[string]any{"ok": true, "queued": true, "operationId": payload.OperationID}) close(submitted) case "/sessions/test/stream": writer.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprint(writer, "data: {\"cursor\":0,\"events\":[{\"type\":\"reset\"}]}\n\n") writer.(http.Flusher).Flush() <-request.Context().Done() case "/sessions/test/cancel-submission":25 unmodified lines26 unmodified lines2728293031313233333435353637383938 unmodified lines26 unmodified lines cursors = append(cursors, cursor) var body string switch cursor { case "": body = `{"events":[{"type":"user","text":"third"},{"type":"message","text":"third reply"}],"before":5,"more":true}` body = `{"events":[{"type":"user","source":"chat","triggeredAt":"","text":"third"},{"type":"message","role":"assistant","text":"third reply"}],"before":5,"more":true}` case "5": body = `{"events":[{"type":"user","text":"second"},{"type":"message","text":"second reply"}],"before":3,"more":true}` body = `{"events":[{"type":"user","source":"chat","triggeredAt":"","text":"second"},{"type":"message","role":"assistant","text":"second reply"}],"before":3,"more":true}` case "3": body = `{"events":[{"type":"user","text":"first"},{"type":"message","text":"first reply"}],"before":1,"more":false}` body = `{"events":[{"type":"user","source":"chat","triggeredAt":"","text":"first"},{"type":"message","role":"assistant","text":"first reply"}],"before":1,"more":false}` default: t.Errorf("unexpected cursor %s", cursor) } _, _ = w.Write([]byte(body))38 unmodified lines1 unmodified line2345667891 unmodified line11121314161516171824 unmodified lines4344454649505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989947484950514 unmodified lines1 unmodified line// through a scripted provider; these tests exercise the public transport boundaries.package daemon
import ( "bufio" "context" "encoding/json" "errors" "fmt"1 unmodified line "net/http" "net/http/httptest" "os" "path/filepath" "strings" "testing" "time")
24 unmodified lines }) }}
func TestMalformedBatchDoesNotAdvanceCursorOrDeliverPartialEvents(t *testing.T) { client := NewChatClient(NewConnection(ConnectionSnapshot{}, ""), "s") delivered := 0 scan := func(payload string) error { return client.readStream(context.Background(), bufio.NewScanner(strings.NewReader(payload)), func(event StreamEvent) error { if event.Type == EventText { delivered++ } return nil }) } if err := scan("data: {\"cursor\":7,\"events\":[{\"type\":\"text\",\"text\":\"ok\"}]}\n\n"); err != nil { t.Fatal(err) } if err := scan("data: {\"cursor\":8,\"events\":[{\"type\":\"text\",\"text\":\"partial\"},{\"type\":\"message\",\"text\":false}]}\n\n"); err == nil { t.Fatal("malformed known event accepted") } if client.afterSeq != 7 || delivered != 1 { t.Fatalf("malformed batch changed stream state: cursor=%d delivered=%d", client.afterSeq, delivered) } if err := scan("data: {\"cursor\":9,\"events\":[{\"type\":\"future\",\"text\":false}]}\n\n"); err != nil { t.Fatal(err) } if client.afterSeq != 9 || delivered != 1 { t.Fatal("unknown event should advance cursor without delivery") }}
func TestTypedEventFormatsAndValidation(t *testing.T) { for _, body := range []string{ `{"type":"turn_membership","turnId":"turn","submissionIds":["one","two"]}`, `{"type":"turn_completed","turnId":"turn"}`, `{"type":"tool","name":"python","args":"{\"code\":\"print(1)\"}","result":"1"}`, `{"type":"usage","promptTokens":10,"elapsedMs":1.5,"cacheFade":[{"at":1,"cached":2}]}`, `{"type":"tool_progress","progress":null}`, } { if _, err := decodeChatEvent(json.RawMessage(body)); err != nil { t.Fatalf("valid format %s: %v", body, err) } } for _, body := range []string{ `{"type":"turn_membership","turnId":"turn","submissionIds":[1]}`, `{"type":"usage","promptTokens":"10"}`, `{"type":"thinking","text":"thought","elapsedMs":1.5}`, `{"type":"committed","seq":0}`, `{"type":"compacted","evicted":2}`, } { if _, err := decodeChatEvent(json.RawMessage(body)); err == nil { t.Fatalf("malformed event accepted: %s", body) } }func TestAgentEventsValidateNullableSendersAndBooleanFields(t *testing.T) { event, err := decodeAgentEvent(json.RawMessage(`{"type":"mail","from":null,"fromName":"","to":"s","kind":"message","bytes":2}`)) if err != nil || event == nil || event.To != "s" { t.Fatalf("nullable bus sender rejected: %v", err) }4 unmodified lines6 unmodified lines7891011121314151639 unmodified lines565758596061626364656663676869707172207 unmodified linesExpand all28028128228328428528628728028828929029129214 unmodified lines3073083093103033113123053063073133093143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523 unmodified lines35635735835932836036136236336433033133233333436536636733533636833833936937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140235335435540335740440540640740840936341041141241341436841541641741841910 unmodified lines4304314324333873884344354363903913924374384394404414424434444454464474484494504512 unmodified lines45445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751810 unmodified lines52953053153243143253353453553653753853954054154254343654454554654744054854955055119 unmodified lines5715725735744685755765774715785795804745815824765835844784794804815855864835875885894865905915925935945955964935975985996006016024996036046056066076086096106116126136146156166176186196206216226236242 unmodified lines6276286296305105115126316326335166346356366376384 unmodified lines6436446456465295305315325336476486496506516521 unmodified line6546556566575436586596606616626636646656666676686696706716726735556746756766775595605615625636785655665675685695705715726796806816825776836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187196 unmodified lines "encoding/hex" "encoding/json" "errors" "fmt" "io" "net" "net/http" "net/url" "strings" "sync"39 unmodified lines OK bool `json:"ok"` Queued bool `json:"queued"`}
// ChatClient retains the cursor and argument state for one stream subscription.// Callers must serialize stream subscriptions and reset operations.type ChatClient struct { conn *Connection argumentsByCall map[string]*strings.Builder agentID string clientID string afterSeq int afterSeq int64 requireReset bool mu sync.Mutex}
// NewChatClient shares conn's HTTP transport. A nil connection panics.207 unmodified linesExpand all } if err := json.Unmarshal(data, &header); err != nil { return nil, fmt.Errorf("invalid event: %w", err) } if header.Type == "" { return nil, errors.New("invalid event: missing or empty type") } switch header.Type { case "", EventText, EventReset, EventRetry, EventError, EventInterrupted, EventMessage, EventUser, EventThinking, EventNote, EventCompacted, EventToolProgress, EventTool, EventUsage, EventCommitted, EventTurnMembership, EventTurnCompleted, "arguments_delta", "turn_started": case EventText, EventReset, EventRetry, EventError, EventInterrupted, EventMessage, EventUser, EventThinking, EventNote, EventCompacted, EventToolProgress, EventTool, EventUsage, EventCommitted, EventTurnMembership, EventTurnCompleted, "arguments_delta", "turn_started": default: return nil, nil } var event wireChatEvent14 unmodified lines return nil } var err error switch header.Type { case "", EventText, EventMessage, EventUser, EventThinking, EventNote, EventError: case EventText, EventMessage, EventUser, EventThinking, EventNote, EventError: err = require("text") if header.Type == "" { event.Type = EventText } if header.Type == EventMessage { event.Role = "assistant" err = require("text", "role") if event.Role != "assistant" { return nil, errors.New("invalid message event: role must be assistant") } } if header.Type == EventUser { err = require("text", "source", "triggeredAt") } case "arguments_delta": err = require("callId", "name", "text") if event.CallID == "" || event.ToolName == "" { return nil, errors.New("invalid argument delta identity") } case EventTurnMembership: err = require("turnId", "submissionIds") if event.TurnID == "" || event.SubmissionIDs == nil { return nil, errors.New("invalid turn membership identity") } for _, id := range event.SubmissionIDs { if id == "" { return nil, errors.New("invalid empty submission identity") } } case EventTurnCompleted: err = require("turnId") if event.TurnID == "" { return nil, errors.New("invalid empty turn identity") } case EventReset: if raw, ok := fields["before"]; ok && string(raw) == "null" { return nil, errors.New("invalid reset before") } if raw, ok := fields["more"]; ok && string(raw) == "null" { return nil, errors.New("invalid reset more") } case EventCommitted: err = require("seq") if event.Seq <= 0 { err = errors.New("invalid committed event: seq must be positive")3 unmodified lines if event.Evicted < 0 || event.Evicted > 1000000 || len(event.Summary) > 60000 { err = errors.New("invalid compacted event") } case EventTool: err = require("name", "result") err = require("callId", "name", "args", "result") if event.CallID == "" || event.ToolName == "" { return nil, errors.New("invalid tool identity") } if len(event.Args) != 0 && string(event.Args) != "null" { args := event.Args if args[0] == '"' { var encoded string if decodeErr := json.Unmarshal(args, &encoded); decodeErr != nil { return nil, decodeErr var encoded string if decodeErr := json.Unmarshal(event.Args, &encoded); decodeErr != nil { return nil, fmt.Errorf("invalid tool args: %w", decodeErr) } args = []byte(encoded) } // Invalid tool argument JSON is display-only, and has historically rendered empty. if json.Unmarshal(args, &event.ToolArgs) != nil { // Tool arguments are model-generated text, not protocol-owned JSON. if json.Unmarshal([]byte(encoded), &event.ToolArgs) != nil { event.ToolArgs = map[string]any{} } } case EventUsage: var usage Usage if decodeErr := json.Unmarshal(data, &usage); decodeErr != nil { return nil, fmt.Errorf("invalid usage event: %w", decodeErr) } if fieldErr := require("model", "recordedAt"); fieldErr != nil { return nil, fieldErr } if usage.RecordedAt == nil || *usage.RecordedAt < 0 { return nil, errors.New("invalid usage timestamp") } for _, count := range []*int{usage.PromptTokens, usage.CachedPromptTokens, usage.CacheWriteTokens, usage.CompletionTokens, usage.TotalTokens} { if count != nil && *count < 0 { return nil, errors.New("invalid usage token count") } } if usage.ElapsedMs != nil && *usage.ElapsedMs < 0 || usage.TokensPerSecond != nil && *usage.TokensPerSecond < 0 { return nil, errors.New("invalid usage rate or duration") } for _, step := range usage.CacheFade { if step.At < 0 || step.Cached != nil && *step.Cached < 0 { return nil, errors.New("invalid cache fade") } } event.Usage = &usage case EventToolProgress: if fieldErr := validateToolProgress(fields["progress"]); fieldErr != nil { return nil, fieldErr } if _, exists := fields["progress"]; !exists { err = errors.New("invalid tool progress event: missing progress") } if progress := event.Progress; progress != nil { if len(progress.CallID) > 200 || len(progress.Name) > 100 || (progress.Phase != "generating" && progress.Phase != "running") { if progress.CallID == "" || progress.Name == "" || len(progress.CallID) > 200 || len(progress.Name) > 100 { return nil, errors.New("invalid tool progress event") } progress.Name = cleanLabel(progress.Name) if intent := progress.Intent; intent != nil { if (intent.Kind != "write" && intent.Kind != "edit" && intent.Kind != "read" && intent.Kind != "run") || len(intent.Target) > 300 { progress.Intent = nil return nil, errors.New("invalid tool progress intent") } } if code := progress.Code; code != nil { if code.Offset < 0 || (progress.Phase == "generating" && len(code.Text) > 512) || len(code.Text) > 16000 { progress.Code = nil return nil, errors.New("invalid tool progress code") } else if progress.Phase == "running" { if progress.Intent == nil && progress.Name == "python" { progress.Intent = ParsePythonIntent(code.Text) }10 unmodified lines if err := json.Unmarshal(event.Elapsed, &event.ElapsedMs); err != nil { return nil, fmt.Errorf("invalid elapsedMs: %w", err) } } if event.Timestamp != nil && *event.Timestamp < 0 { event.Timestamp = nil if event.Timestamp != nil && *event.Timestamp < 0 || event.ElapsedMs < 0 || event.Before < 0 { return nil, errors.New("invalid event timestamp, elapsed time, or history cursor") } if event.ElapsedMs < 0 { event.ElapsedMs = 0 } if event.Image != nil { event.Image = ParseImageMetadata(event.Image) if event.Image == nil { return nil, errors.New("invalid image metadata") } } if event.ToolTrace != nil { event.ToolTrace = ParseToolTrace(event.ToolTrace) if event.ToolTrace == nil { return nil, errors.New("invalid tool trace") } } return &event, nil}
2 unmodified linesfunc (c *ChatClient) ResetStream() { c.mu.Lock() c.argumentsByCall = make(map[string]*strings.Builder) c.afterSeq = -1 c.requireReset = false c.mu.Unlock()}
// RequireStreamReset requests durable history before accepting further events.// Call it after the subscription has stopped writing.func (c *ChatClient) RequireStreamReset() { c.mu.Lock() c.argumentsByCall = make(map[string]*strings.Builder) c.afterSeq = -1 c.requireReset = true c.mu.Unlock()}
type streamBatch struct { Cursor *int64 `json:"cursor"` Events []json.RawMessage `json:"events"`}
func streamFailure(kind StreamFailureKind, cause error) error { return &StreamError{Kind: kind, Cause: cause}}
func classifyStreamFailure(err error) error { if err == nil { return nil } if _, ok := errors.AsType[*StreamError](err); ok { return err } if apiErr, ok := errors.AsType[*APIError](err); ok { kind := StreamTerminal if apiErr.StatusCode == http.StatusRequestTimeout || apiErr.StatusCode == http.StatusTooManyRequests || apiErr.StatusCode >= 500 { kind = StreamTransient } return streamFailure(kind, err) } if errors.Is(err, bufio.ErrTooLong) { return streamFailure(StreamProtocol, err) } if _, ok := errors.AsType[net.Error](err); ok { return streamFailure(StreamTransient, err) } if isConnectionError(err) || errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return streamFailure(StreamTransient, err) } return streamFailure(StreamTerminal, err)}
// Stream uses ctx for cancellation; onEvent must also honor ctx if it blocks.func (c *ChatClient) Stream(ctx context.Context, tail int, onEvent func(StreamEvent) error) error { return c.StreamWithProgress(ctx, tail, onEvent, nil)}
// StreamWithProgress reports fully consumed batches, including ignored additive// events, so subscribers can reset their reconnect backoff after progress.func (c *ChatClient) StreamWithProgress(ctx context.Context, tail int, onEvent func(StreamEvent) error, onBatchConsumed func()) error { defer func() { // Cancellation also clears state when the request fails before reading. // Transient failures and EOF retain the cursor and unfinished arguments. if ctx.Err() != nil {10 unmodified lines route += fmt.Sprintf("&tail=%d", tail) }
operation := Operation{Name: "stream session", Method: http.MethodGet, Path: c.agentPath(route), Policy: ReadRecovery} return scanEventStream(ctx, c.conn, operation, streamLimits{lineBytes: 10 * 1024 * 1024, errorBytes: 64 * 1024}, func(scanner *bufio.Scanner) error { return c.readStream(ctx, scanner, onEvent) err := scanEventStream(ctx, c.conn, operation, streamLimits{requireSSE: true, lineBytes: 10 * 1024 * 1024, errorBytes: 64 * 1024}, func(scanner *bufio.Scanner) error { return c.readStream(ctx, scanner, onBatchConsumed, func(event StreamEvent) error { if err := onEvent(event); err != nil { return streamFailure(StreamTerminal, err) } return nil }) }) return classifyStreamFailure(err)}
func (c *ChatClient) readStream(ctx context.Context, scanner *bufio.Scanner, onEvent func(StreamEvent) error) error {func (c *ChatClient) readStream(ctx context.Context, scanner *bufio.Scanner, onBatchConsumed func(), onEvent func(StreamEvent) error) error { reporter := NewToolProgressReporter(func(progress *ToolProgress) error { return onEvent(StreamEvent{Type: EventToolProgress, Progress: progress}) }) defer func() { _ = reporter.Report(nil, "running") }() var eventType string for scanner.Scan() { if err := ctx.Err(); err != nil { return err19 unmodified lines Error string `json:"error"` Text string `json:"text"` } if err := json.Unmarshal([]byte(payload), &failure); err != nil { return fmt.Errorf("invalid stream error: %w", err) return streamFailure(StreamProtocol, fmt.Errorf("invalid stream error: %w", err)) } if failure.Error != "" { return errors.New(failure.Error) return streamFailure(StreamTerminal, errors.New(failure.Error)) } if failure.Text != "" { return errors.New(failure.Text) return streamFailure(StreamTerminal, errors.New(failure.Text)) } return errors.New(payload) return streamFailure(StreamProtocol, errors.New("invalid stream error: missing error or text")) } var batch struct { Cursor *int `json:"cursor"` Events []json.RawMessage `json:"events"` } var batch streamBatch if err := json.Unmarshal([]byte(payload), &batch); err != nil { return fmt.Errorf("invalid stream batch: %w", err) return streamFailure(StreamProtocol, fmt.Errorf("invalid stream batch: %w", err)) } if batch.Cursor == nil || *batch.Cursor < 0 || batch.Events == nil { return errors.New("invalid stream batch: missing cursor or events") return streamFailure(StreamProtocol, errors.New("invalid stream batch: missing cursor or events")) } // Validate the whole batch before delivering events or changing preview state. events := make([]wireChatEvent, 0, len(batch.Events)) for _, raw := range batch.Events { event, err := decodeChatEvent(raw) if err != nil { return err return streamFailure(StreamProtocol, err) } if event != nil { events = append(events, *event) } } snapshot := false reset := len(events) > 0 && events[0].Type == EventReset c.mu.Lock() previous, requiresReset := c.afterSeq, c.requireReset c.mu.Unlock() if requiresReset && !reset { return streamFailure(StreamProtocol, errors.New("recovery stream must begin with reset")) } if *batch.Cursor < previous && !reset { return streamFailure(StreamProtocol, errors.New("stream cursor regressed without reset")) } for i, event := range events { if event.Type == EventReset && i != 0 { return streamFailure(StreamProtocol, errors.New("reset must begin its batch")) } } if err := c.validateArgumentBatch(events); err != nil { return streamFailure(StreamProtocol, err) } for _, event := range events { switch event.Type { case EventReset, EventRetry, "turn_started", EventMessage, EventInterrupted, EventError: c.mu.Lock()2 unmodified lines if err := reporter.Report(nil, "running"); err != nil { return err } } if event.Type == EventReset { snapshot = true } if event.Type == "turn_started" { continue } if event.Type == EventTool || (event.Type == EventToolProgress && event.Progress != nil && event.Progress.Phase == "running" && event.Progress.CallID != "") { if event.Type == EventTool || (event.Type == EventToolProgress && event.Progress != nil && event.Progress.Phase == "running") { callID := event.CallID if event.Progress != nil { callID = event.Progress.CallID }4 unmodified lines return err } } if event.Type == "arguments_delta" { if event.CallID != "" && event.Text != "" { arguments, err := c.appendArguments(event.CallID, event.Text) if err != nil { return err } if event.Text != "" { arguments := c.appendArguments(event.CallID, event.Text) call := &ToolCallAssembly{ID: event.CallID} call.Function.Name = event.ToolName call.Function.Arguments = arguments if err := reporter.Report(call, "generating"); err != nil {1 unmodified line } } continue } event.Replayed = snapshot && event.Type != EventReset event.Replayed = reset && event.Type != EventReset if err := onEvent(event.StreamEvent); err != nil { return err } } c.mu.Lock() c.afterSeq = *batch.Cursor c.requireReset = false c.mu.Unlock() if onBatchConsumed != nil { onBatchConsumed() } } return nil}
func (c *ChatClient) appendArguments(callID, text string) (string, error) {func (c *ChatClient) appendArguments(callID, text string) string { c.mu.Lock() defer c.mu.Unlock() arguments, exists := c.argumentsByCall[callID] retained := 0 for key, value := range c.argumentsByCall { retained += len(key) + value.Len() } extraCallID := 0 if !exists { extraCallID = len(callID) } if (!exists && len(c.argumentsByCall) >= 32) || retained+len(text)+extraCallID > 2000000 { c.argumentsByCall = make(map[string]*strings.Builder) c.afterSeq = -1 return "", errors.New("too many tool argument previews for this client to display") } if !exists { arguments = new(strings.Builder) c.argumentsByCall[callID] = arguments } arguments.WriteString(text) return arguments.String(), nil return arguments.String()}
// validateArgumentBatch checks preview limits before delivering any batch event.func (c *ChatClient) validateArgumentBatch(events []wireChatEvent) error { sizes := make(map[string]int) c.mu.Lock() for id, args := range c.argumentsByCall { sizes[id] = args.Len() } c.mu.Unlock() for _, event := range events { switch event.Type { case EventReset, EventRetry, "turn_started", EventMessage, EventInterrupted, EventError: clear(sizes) case EventTool: delete(sizes, event.CallID) case EventToolProgress: if event.Progress != nil && event.Progress.Phase == "running" { delete(sizes, event.Progress.CallID) } case "arguments_delta": if event.Text == "" { continue } sizes[event.CallID] += len(event.Text) retained := 0 for id, size := range sizes { retained += len(id) + size } if len(sizes) > 32 || retained > 2000000 { return errors.New("too many tool argument previews for this client to display") } } } return nil}69 unmodified lines7071727374757677787980818283848586878889909169 unmodified lines } return e.Operation + " may have been accepted; check its result before trying again"}func (e *UncertainOutcomeError) Unwrap() error { return e.Cause }
// StreamFailureKind tells subscribers whether reconnecting can make progress.type StreamFailureKind uint8
const ( StreamTransient StreamFailureKind = iota StreamTerminal StreamProtocol)
// StreamError preserves the cause while defining subscription recovery.type StreamError struct { Kind StreamFailureKind Cause error}
func (e *StreamError) Error() string { return e.Cause.Error() }func (e *StreamError) Unwrap() error { return e.Cause }6 unmodified lines789101112131415184 unmodified linesExpand all20020120220320420520620720837 unmodified lines24624724824925025125225325425525625725825926048 unmodified lines6 unmodified lines "encoding/json" "errors" "fmt" "io" "mime" "net" "net/http" "strings" "syscall"184 unmodified linesExpand all panic("unreachable retry budget")}
type streamLimits struct { requireSSE bool lineBytes int errorBytes int64}
37 unmodified lines err = errors.Join(err, canceled) } return err } if limits.requireSSE { mediaType, _, parseErr := mime.ParseMediaType(res.Header.Get("Content-Type")) if parseErr != nil || mediaType != "text/event-stream" { _ = res.Body.Close() return streamFailure(StreamProtocol, fmt.Errorf("invalid stream content type %q", res.Header.Get("Content-Type"))) } } defer res.Body.Close() scanner := bufio.NewScanner(res.Body) scanner.Buffer(make([]byte, 64*1024), limits.lineBytes) if err := consume(scanner); err != nil {48 unmodified lines123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283// A daemon cannot emit deliberately malformed wire data. Controlled HTTP peers// exercise batch atomicity and failure classification through the public stream API.package daemon
import ( "encoding/json" "errors" "fmt" "io" "net" "net/http" "net/http/httptest" "os" "sync/atomic" "testing")
func TestStreamRejectsMalformedBatchesWithoutLosingTheCursor(t *testing.T) { data, err := os.ReadFile("../../../test/fixtures/session_stream.json") if err != nil { t.Fatal(err) } var fixture struct { Malformed map[string]string `json:"malformed_batches"` } if err := json.Unmarshal(data, &fixture); err != nil { t.Fatal(err) } if len(fixture.Malformed) == 0 { t.Fatal("empty stream contract fixture") } for name, batch := range fixture.Malformed { t.Run(name, func(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") switch requests.Add(1) { case 1: _, _ = io.WriteString(w, "data: "+`{"cursor":7,"events":[{"type":"text","text":"accepted"}]}`+"\n\n") case 2: _, _ = fmt.Fprintf(w, "data: %s\n\n", batch) default: if got := r.URL.Query().Get("after_seq"); got != "7" { t.Errorf("failed batch lost committed cursor: %s", got) } _, _ = io.WriteString(w, "data: "+`{"cursor":8,"events":[{"type":"text","text":"replayed"}]}`+"\n\n") } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") var delivered []string consume := func(event StreamEvent) error { if event.Type != EventText { t.Errorf("unexpected event delivered: %+v", event) } delivered = append(delivered, event.Text) return nil } if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatal(err) } consumed := 0 err := client.StreamWithProgress(t.Context(), 0, consume, func() { consumed++ }) if consumed != 0 { t.Fatal("malformed batch reported consumption") } failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != StreamProtocol { t.Fatalf("malformed batch was not a protocol failure: %v", err) } if len(delivered) != 1 || delivered[0] != "accepted" { t.Fatalf("malformed batch partially delivered: %v", delivered) } if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatal(err) } if len(delivered) != 2 || delivered[1] != "replayed" || requests.Load() != 3 { t.Fatalf("replay lost events or retried internally: %v, requests=%d", delivered, requests.Load()) } }) }}
func TestStreamIgnoresUnknownKindsAndCommitsTheirCursor(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") if requests.Add(1) == 1 { _, _ = io.WriteString(w, "data: "+`{"cursor":12,"events":[{"type":"future_output","text":"not assistant output","args":false}]}`+"\n\n") } else if got := r.URL.Query().Get("after_seq"); got != "12" { t.Errorf("ignored additive batch lost cursor: %s", got) } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") consumed := 0 for range 2 { if err := client.StreamWithProgress(t.Context(), 0, func(event StreamEvent) error { t.Errorf("unknown event delivered: %+v", event) return nil }, func() { consumed++ }); err != nil { t.Fatal(err) } } if consumed != 1 { t.Fatalf("ignored-only batch did not report consumption: %d", consumed) }}
func TestStreamCallbackFailureRetainsCauseAndUncommittedCursor(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") n := requests.Add(1) if n == 1 { _, _ = io.WriteString(w, "data: "+`{"cursor":7,"events":[]}`+"\n\n") return } if n == 3 && r.URL.Query().Get("after_seq") != "7" { t.Error("callback failure committed the batch cursor") } _, _ = io.WriteString(w, "data: "+`{"cursor":8,"events":[{"type":"text","text":"first"},{"type":"text","text":"second"}]}`+"\n\n") })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") if err := client.Stream(t.Context(), 0, func(StreamEvent) error { return nil }); err != nil { t.Fatal(err) } cause := errors.New("consumer stopped") var callbacks int err := client.Stream(t.Context(), 0, func(event StreamEvent) error { callbacks++ if callbacks == 2 { return cause } return nil }) failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != StreamTerminal || !errors.Is(err, cause) { t.Fatalf("callback failure lost classification or cause: %v", err) } if requests.Load() != 2 || callbacks != 2 { t.Fatalf("callback failure changed delivery or retried: requests=%d, callbacks=%d", requests.Load(), callbacks) } if err := client.Stream(t.Context(), 0, func(StreamEvent) error { return nil }); err != nil { t.Fatal(err) }}
func TestStreamClassifiesHTTPRefusalsWithoutRetryingThem(t *testing.T) { for _, test := range []struct { status int kind StreamFailureKind }{ {http.StatusNotFound, StreamTerminal}, {http.StatusForbidden, StreamTerminal}, {http.StatusBadRequest, StreamTerminal}, {http.StatusRequestTimeout, StreamTransient}, {http.StatusTooManyRequests, StreamTransient}, {http.StatusServiceUnavailable, StreamTransient}, } { t.Run(http.StatusText(test.status), func(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requests.Add(1) w.WriteHeader(test.status) _, _ = io.WriteString(w, `{"error":"stream unavailable"}`) })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("HTTP refusal delivered an event") return nil }) failure, ok := errors.AsType[*StreamError](err) api, apiOK := errors.AsType[*APIError](err) if !ok || failure.Kind != test.kind || !apiOK || api.StatusCode != test.status || requests.Load() != 1 { t.Fatalf("wrong refusal classification, cause, or retry: %v, requests=%d", err, requests.Load()) } }) }}
func TestStreamRejectsNonSSESuccessResponses(t *testing.T) { for _, media := range []string{"", "text/html", "application/json"} { t.Run(media, func(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header()["Content-Type"] = nil if media != "" { w.Header().Set("Content-Type", media) } _, _ = io.WriteString(w, `{"error":"this is not a stream"}`) })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("non-SSE response delivered an event") return nil }) failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != StreamProtocol { t.Fatalf("wrong media type treated as recoverable EOF: %v", err) } }) }}
func TestStreamRecoveryRequiresResetBeforeDeliveringHistory(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") switch requests.Add(1) { case 1: _, _ = io.WriteString(w, "data: "+`{"cursor":20,"events":[]}`+"\n\n") case 2: _, _ = io.WriteString(w, "data: "+`{"cursor":1,"events":[{"type":"text","text":"not durable reset"}]}`+"\n\n") case 3: _, _ = io.WriteString(w, "data: "+`{"cursor":1,"events":[{"type":"reset"},{"type":"message","role":"assistant","text":"durable"}]}`+"\n\n") default: if r.URL.Query().Get("after_seq") != "1" { t.Error("validated reset failed to commit its cursor") } } })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") consume := func(StreamEvent) error { return nil } if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatal(err) } client.RequireStreamReset() err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("non-reset recovery data escaped") return nil }) failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != StreamProtocol { t.Fatalf("missing reset accepted: %v", err) } var replayed bool if err := client.Stream(t.Context(), 0, func(event StreamEvent) error { if event.Type == EventMessage { replayed = event.Replayed } return nil }); err != nil { t.Fatal(err) } if !replayed { t.Fatal("reset snapshot was not marked as replayed history") } if err := client.Stream(t.Context(), 0, consume); err != nil { t.Fatal(err) }}
func TestStreamClassifiesExplicitSSEFailures(t *testing.T) { for _, test := range []struct { name, payload string kind StreamFailureKind }{ {"reported failure", `{"error":"session removed"}`, StreamTerminal}, {"malformed JSON", `{"error":`, StreamProtocol}, {"malformed field", `{"error":false}`, StreamProtocol}, {"missing reason", `{}`, StreamProtocol}, } { t.Run(test.name, func(t *testing.T) { var requests atomic.Int32 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { requests.Add(1) w.Header().Set("Content-Type", "text/event-stream") _, _ = fmt.Fprintf(w, "event: error\ndata: %s\n\n", test.payload) })) defer server.Close() client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, ""), "session") err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("SSE failure became transcript output") return nil }) failure, ok := errors.AsType[*StreamError](err) if !ok || failure.Kind != test.kind || requests.Load() != 1 { t.Fatalf("wrong explicit failure or automatic retry: %v, requests=%d", err, requests.Load()) } }) }}