// Controlled streams exercise preview assembly across EOF, cancellation, and
// reset boundaries that provider timing cannot reliably produce.
package daemon
import (
"context"
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
)
func formatPage(cursor int, events []any) string {
batch, _ := json.Marshal(map[string]any{"generation": "generation-a", "cursor": cursor, "events": events})
return fmt.Sprintf("data: %s\n\n", batch)
}
func TestToolPreviewContinuesAcrossEOFAndTransientFailure(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
writer.Header().Set("Content-Type", "text/event-stream")
switch requests.Add(1) {
case 1:
_, _ = writer.Write([]byte(formatPage(1, []any{
map[string]any{"type": "reset"},
map[string]any{"type": "arguments_delta", "name": "python", "callId": "read", "text": `{"code":"read('first.py')"}`},
map[string]any{"type": "tool", "callId": "read", "name": "python", "args": `{"code":"read('first.py')"}`, "result": "done"},
map[string]any{"type": "arguments_delta", "name": "python", "callId": "edit", "text": `{"code":"edit('second`},
})))
case 2:
writer.WriteHeader(http.StatusServiceUnavailable)
case 3:
if cursor := request.URL.Query().Get("after_seq"); cursor != "1" {
t.Errorf("transient failure lost the consumed cursor: %s", cursor)
}
_, _ = writer.Write([]byte(formatPage(2, []any{
map[string]any{"type": "arguments_delta", "name": "python", "callId": "edit", "text": `.py')"}`},
map[string]any{"type": "tool_progress", "progress": map[string]any{"callId": "edit", "name": "python", "phase": "running"}},
map[string]any{"type": "tool", "callId": "edit", "name": "python", "args": `{"code":"edit('second.py')"}`, "result": "done"},
})))
default:
t.Errorf("unexpected automatic stream request: %s", request.URL)
}
}))
defer server.Close()
client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session")
var previews []string
var completed []string
consume := func(event StreamEvent) error {
if event.Type == EventToolProgress && event.Progress != nil && event.Progress.Code != nil {
previews = append(previews, event.Progress.Code.Text)
}
if event.Type == EventTool {
completed = append(completed, event.ToolResult)
}
return nil
}
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatal(err)
}
err := client.Stream(t.Context(), 0, consume)
failure, ok := errors.AsType[*StreamError](err)
if !ok || failure.Kind != StreamTransient {
t.Fatalf("server failure was not transient: %v", err)
}
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatal(err)
}
if len(previews) != 3 || previews[0] != "read('first.py')" || previews[1] != "edit('second" || previews[2] != "edit('second.py')" {
t.Fatalf("reconnect lost or mixed tool argument fragments: %q", previews)
}
if len(completed) != 2 {
t.Fatalf("lost completed tool results: %v", completed)
}
}
func TestStreamRejectsExcessUnfinishedCallsBeforeDeliveringTheBatch(t *testing.T) {
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
writer.Header().Set("Content-Type", "text/event-stream")
count := requests.Add(1)
switch {
case count <= 32:
events := []any{map[string]any{"type": "arguments_delta", "name": "python", "callId": fmt.Sprintf("call-%d", count), "text": `{"code":"`}}
if count == 1 {
events = append([]any{map[string]any{"type": "reset"}}, events...)
}
_, _ = writer.Write([]byte(formatPage(int(count), events)))
case count == 33:
_, _ = writer.Write([]byte(formatPage(33, []any{
map[string]any{"type": "text", "text": "must not escape"},
map[string]any{"type": "arguments_delta", "name": "python", "callId": "excess-call", "text": `{"code":"`},
})))
case count == 34:
if cursor := request.URL.Query().Get("after_seq"); cursor != "32" {
t.Errorf("rejected batch advanced the cursor: %s", cursor)
}
_, _ = writer.Write([]byte(formatPage(32, []any{})))
case count == 35:
if cursor := request.URL.Query().Get("after_seq"); cursor != "" || request.URL.Query().Has("after_generation") {
t.Errorf("recovery did not request a fresh snapshot: %s", cursor)
}
_, _ = writer.Write([]byte(formatPage(35, []any{
map[string]any{"type": "reset"},
map[string]any{"type": "message", "role": "assistant", "text": "recovered"},
})))
default:
t.Error("unexpected stream request")
}
}))
defer server.Close()
client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session")
consume := func(StreamEvent) error { return nil }
for range 32 {
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatalf("accepted call rejected: %v", err)
}
}
err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("oversized batch partially delivered"); return nil })
failure, ok := errors.AsType[*StreamError](err)
if !ok || failure.Kind != StreamProtocol {
t.Fatalf("excess unfinished call did not fail protocol validation: %v", err)
}
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatal(err)
}
client.ResetStream()
var recovered bool
if err := client.Stream(t.Context(), 0, func(event StreamEvent) error {
recovered = recovered || event.Type == EventMessage && event.Text == "recovered" && event.Replayed
return nil
}); err != nil {
t.Fatal(err)
}
if !recovered {
t.Fatal("recovery did not deliver durable history")
}
}
func TestStreamRejectsExcessArgumentBytesBeforeDeliveringTheBatch(t *testing.T) {
firstArguments := strings.Repeat("x", 999_999)
secondArguments := strings.Repeat("y", 2_000_000-len("first")-len("second")-len(firstArguments))
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
writer.Header().Set("Content-Type", "text/event-stream")
switch requests.Add(1) {
case 1:
_, _ = writer.Write([]byte(formatPage(1, []any{
map[string]any{"type": "reset"},
map[string]any{"type": "arguments_delta", "name": "other", "callId": "first", "text": firstArguments},
map[string]any{"type": "arguments_delta", "name": "other", "callId": "second", "text": secondArguments},
})))
case 2:
_, _ = writer.Write([]byte(formatPage(2, []any{
map[string]any{"type": "text", "text": "must not escape"},
map[string]any{"type": "arguments_delta", "name": "other", "callId": "second", "text": "!"},
})))
case 3:
if cursor := request.URL.Query().Get("after_seq"); cursor != "1" {
t.Errorf("rejected byte overflow advanced the cursor: %s", cursor)
}
_, _ = writer.Write([]byte(formatPage(1, []any{})))
case 4:
_, _ = writer.Write([]byte(formatPage(3, []any{map[string]any{"type": "reset"}})))
_, _ = writer.Write([]byte(formatPage(4, []any{
map[string]any{"type": "text", "text": "must not escape"},
map[string]any{"type": "arguments_delta", "name": "other", "callId": "first", "text": firstArguments},
map[string]any{"type": "arguments_delta", "name": "other", "callId": "second", "text": secondArguments + "!"},
})))
case 5:
if cursor := request.URL.Query().Get("after_seq"); cursor != "3" {
t.Errorf("overflow lost the preceding successful reset: %s", cursor)
}
_, _ = writer.Write([]byte(formatPage(3, []any{})))
default:
t.Error("unexpected stream request")
}
}))
defer server.Close()
client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session")
consume := func(StreamEvent) error { return nil }
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatalf("exact byte budget rejected: %v", err)
}
err := client.Stream(t.Context(), 0, func(StreamEvent) error { t.Error("overflowing batch partially delivered"); return nil })
failure, ok := errors.AsType[*StreamError](err)
if !ok || failure.Kind != StreamProtocol {
t.Fatalf("existing call overflow was not a protocol failure: %v", err)
}
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatal(err)
}
client.ResetStream()
var delivered []EventType
err = client.Stream(t.Context(), 0, func(event StreamEvent) error { delivered = append(delivered, event.Type); return nil })
failure, ok = errors.AsType[*StreamError](err)
if !ok || failure.Kind != StreamProtocol || len(delivered) != 1 || delivered[0] != EventReset {
t.Fatalf("fresh call overflow leaked a partial batch: error=%v, events=%v", err, delivered)
}
if err := client.Stream(t.Context(), 0, consume); err != nil {
t.Fatal(err)
}
}
func TestCancellationStartsTheNextSubscriptionWithoutOldArguments(t *testing.T) {
for _, beforeHeaders := range []bool{false, true} {
name := "during event delivery"
if beforeHeaders {
name = "before response headers"
}
t.Run(name, func(t *testing.T) {
var requests atomic.Int32
waiting := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
count := requests.Add(1)
if beforeHeaders && count == 2 {
close(waiting)
<-request.Context().Done()
return
}
writer.Header().Set("Content-Type", "text/event-stream")
arguments := `{"code":"old`
if count > 1 {
if cursor := request.URL.Query().Get("after_seq"); cursor != "" || request.URL.Query().Has("after_generation") {
t.Errorf("cancelled subscription retained its cursor: %s", cursor)
}
arguments = `{"code":"new()"}`
}
_, _ = writer.Write([]byte(formatPage(int(count), []any{
map[string]any{"type": "reset"},
map[string]any{"type": "arguments_delta", "name": "python", "callId": "same", "text": arguments},
})))
}))
defer server.Close()
client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session")
ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
defer cancel()
var err error
if beforeHeaders {
if err := client.Stream(t.Context(), 0, func(StreamEvent) error { return nil }); err != nil {
t.Fatal(err)
}
finished := make(chan error, 1)
go func() { finished <- client.Stream(ctx, 0, func(StreamEvent) error { return nil }) }()
select {
case <-waiting:
case <-ctx.Done():
t.Fatal("stream never requested response headers")
}
cancel()
err = <-finished
} else {
err = client.Stream(ctx, 0, func(event StreamEvent) error {
if event.Type == EventReset {
return nil
}
cancel()
return ctx.Err()
})
}
if !errors.Is(err, context.Canceled) {
t.Fatalf("cancellation cause lost: %v", err)
}
var preview string
if err := client.Stream(t.Context(), 0, func(event StreamEvent) error {
if event.Progress != nil && event.Progress.Code != nil {
preview = event.Progress.Code.Text
}
return nil
}); err != nil {
t.Fatal(err)
}
if preview != "new()" {
t.Fatalf("new subscription mixed old argument fragments: %q", preview)
}
})
}
}
func TestSnapshotEventsAreMarkedReplayed(t *testing.T) {
pages := formatPage(2, []any{
map[string]any{"type": "reset"},
map[string]any{"type": "thinking", "text": "earlier"},
}) + formatPage(3, []any{map[string]any{"type": "thinking", "text": "now"}})
server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
writer.Header().Set("Content-Type", "text/event-stream")
_, _ = writer.Write([]byte(pages))
}))
defer server.Close()
client := NewChatClient(NewConnection(ConnectionSnapshot{Port: server.Listener.Addr().(*net.TCPAddr).Port}, nil), "session")
replayed := map[string]bool{}
if err := client.Stream(t.Context(), 0, func(event StreamEvent) error {
if event.Type == EventThinking {
replayed[event.Text] = event.Replayed
}
return nil
}); err != nil {
t.Fatal(err)
}
if len(replayed) != 2 || !replayed["earlier"] || replayed["now"] {
t.Fatalf("only reset snapshot events should be replayed: %v", replayed)
}
}