refactor(daemon): give the client adapter protocol ownership
Keep routes, typed decoding, and recovery in domain files. Migrate callers and preserve confirmed results when a follow-up fails.
76 unmodified lines7778798081828384818283848516 unmodified lines76 unmodified lines
// switchModel moves an idle session to the chosen model without making it the// default for new sessions.func switchModel(ctx context.Context, conn *daemon.Connection, sessionID string, choice modelChoice) error { if err := daemon.CheckCapability(ctx, conn, "session_model", "to choose a model for an existing session"); err != nil { return err } _, err := daemon.SelectModel(ctx, conn, sessionID, map[string]string{"model": choice.model, "provider": choice.provider}) _, err := daemon.SelectModel(ctx, conn, sessionID, daemon.ModelSelectionRequest{Model: choice.model, Provider: choice.provider}) return err}
func (s *Service) Models(ctx context.Context) ([]string, error) {16 unmodified lines26 unmodified lines2728293031313233343510 unmodified lines4647484950505152525354555628 unmodified lines85868788898990919293112 unmodified linesExpand all26 unmodified lines conn, err := s.Connect(ctx) if err != nil { return PromptResult{}, err } if capabilityErr := daemon.CheckCapability(ctx, conn, "submission_cancellation", "prompt submission cancellation"); capabilityErr != nil { if capabilityErr := daemon.CheckPromptSupport(ctx, conn); capabilityErr != nil { return PromptResult{}, capabilityErr } var choice modelChoice if model != "" {10 unmodified lines return PromptResult{}, err } } if sessionID == "" { create := map[string]string{"workspace": cwd} create := daemon.CreateSessionRequest{Workspace: cwd} if model != "" { create["provider"], create["model"] = choice.provider, choice.model create.Provider, create.Model = choice.provider, choice.model } session, createErr := daemon.CreateSession(ctx, conn, create) if createErr != nil { return PromptResult{}, createErr28 unmodified lines handle, err := client.PrepareTurn(prompt, nil, false) if err != nil { return "", err } submissionID := handle.ID submissionID := handle.ID() defer func() { if failure != nil && ctx.Err() != nil { if _, uncertain := errors.AsType[*daemon.UncertainOutcomeError](failure); !uncertain { uncertain := &daemon.UncertainOutcomeError{Operation: "await turn", Handle: handle, Cause: failure}112 unmodified linesExpand all17 unmodified lines1819202122222324242526272849 unmodified lines17 unmodified lines var cursors []string server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/sessions": _, _ = w.Write([]byte(`[{"id":"s"}]`)) _, _ = w.Write([]byte(`[{"id":"s","title":"","workspace":"","provider":"fixture","model":"fixture","protocol":"responses","effort":null,"last_assistant_at":null}]`)) case "/sessions/s/status": _, _ = w.Write([]byte(`{"running":true,"idle":false}`)) _, _ = w.Write([]byte(`{"running":true,"idle":false,"phase":"preparing"}`)) case "/sessions/s/history": cursor := r.URL.Query().Get("before") cursors = append(cursors, cursor) var body string49 unmodified lines3 unmodified lines4567889101130 unmodified lines4243444547464748495019 unmodified lines7071727375747576777811 unmodified lines9091929395949596979827 unmodified lines1261271281291311301311321331348 unmodified lines14314414514614814714814915015141 unmodified lines3 unmodified linesimport ( "context" "errors" "fmt" "net/http" "path/filepath" "slices" "strings"
30 unmodified lines if err != nil { return PreparedOpen{}, err }
sessions, err := daemon.RequestOperation[[]daemon.Session](ctx, conn, daemon.Operation{Name: "list sessions", Method: http.MethodGet, Path: "/sessions", Policy: daemon.ReadRecovery}) sessions, err := daemon.ListSessions(ctx, conn) if err != nil { return PreparedOpen{}, err }
19 unmodified lines var initial *daemon.Session if selected != nil { initial = selected } else if configured && (fresh || len(sessions) == 0) { created, createErr := daemon.CreateSession(ctx, conn, map[string]string{"workspace": absWorkspace}) created, createErr := daemon.CreateSession(ctx, conn, daemon.CreateSessionRequest{Workspace: absWorkspace}) if createErr != nil { return PreparedOpen{}, createErr } initial = &created11 unmodified lines conn, err := s.Connect(ctx) if err != nil { return SessionList{}, err } sessions, err := daemon.RequestOperation[[]daemon.Session](ctx, conn, daemon.Operation{Name: "list sessions", Method: http.MethodGet, Path: "/sessions", Policy: daemon.ReadRecovery}) sessions, err := daemon.ListSessions(ctx, conn) if err != nil { return SessionList{}, err } settings, warning := daemon.GetSettings(ctx, conn)27 unmodified lines}
// resolveSession expands a shortened session ID against the daemon's sessions.func resolveSession(ctx context.Context, conn *daemon.Connection, id string) (string, error) { sessions, err := daemon.RequestOperation[[]daemon.Session](ctx, conn, daemon.Operation{Name: "resolve session", Method: http.MethodGet, Path: "/sessions", Policy: daemon.ReadRecovery}) sessions, err := daemon.ListSessions(ctx, conn) if err != nil { return "", err } session, err := matchSession(sessions, id)8 unmodified lines id, err = resolveSession(ctx, conn, id) if err != nil { return daemon.SendResult{}, err } return daemon.Submit(ctx, conn, id, map[string]any{"content": prompt}) return daemon.Submit(ctx, conn, id, daemon.SubmissionRequest{Content: &prompt})}
type InterruptionResult struct { Interrupted bool `json:"interrupted"`41 unmodified lines5 unmodified lines6789101112131416 unmodified lines31323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718848 unmodified lines237238239240909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321332412422432442452462472482492502512522532542552562572582592602612622632642651351361372662672685 unmodified lines "encoding/json" "errors" "fmt" "net/http" "net/url" "strings")
// AgentEvent is one event from the daemon's agent bus.16 unmodified lines Bytes int `json:"bytes"` Running bool `json:"running"`}
type Member struct { Session string `json:"session"` Parent string `json:"parent"` Name string `json:"name"` Depth int `json:"depth"` Closed bool `json:"closed"`}
type ChildResult struct { Session Session `json:"session"` Member Member `json:"member"`}
type AgentNode struct { Parent *string `json:"parent"` Address *string `json:"address"` Session Session `json:"session"` Name string `json:"name"` Depth int `json:"depth"` Running bool `json:"running"` Closed bool `json:"closed"`}
type AgentsSnapshot struct { Root string `json:"root"` Nodes []AgentNode `json:"nodes"`}
type ChildRequest struct { Name string `json:"name"` Task string `json:"task"`}
// StreamAgents delivers the first batch, including an empty readiness batch,// then nonempty batches until the stream ends or onBatch fails.// The caller controls cancellation, including any blocking work in onBatch.func StreamAgents(ctx context.Context, conn *Connection, onBatch func([]AgentEvent) error) error { operation := operation{Name: "stream agents", Method: http.MethodGet, Path: "/agents/stream", Policy: readRecovery} first := true err := scanEventStream(ctx, conn, operation, streamLimits{requireSSE: true, lineBytes: 8 * 1024 * 1024}, func(scanner *bufio.Scanner) error { for scanner.Scan() { line := scanner.Text() if !strings.HasPrefix(line, "data:") { continue } var batch struct { Events []json.RawMessage `json:"events"` } if err := json.Unmarshal([]byte(strings.TrimSpace(line[5:])), &batch); err != nil { return streamFailure(StreamProtocol, fmt.Errorf("invalid agent batch: %w", err)) } if batch.Events == nil { return streamFailure(StreamProtocol, errors.New("invalid agent batch: missing events")) } events := make([]AgentEvent, 0, len(batch.Events)) for _, raw := range batch.Events { event, err := decodeAgentEvent(raw) if err != nil { return streamFailure(StreamProtocol, err) } if event != nil { if event.Type == "overflow" && len(batch.Events) != 1 { return streamFailure(StreamProtocol, errors.New("overflow must be the sole agent event")) } events = append(events, *event) } } if len(events) == 0 && !first { continue } first = false if err := onBatch(events); err != nil { return streamFailure(StreamTerminal, err) } if len(events) == 1 && events[0].Type == "overflow" { return nil } } return nil }) return classifyStreamFailure(err)}
func CreateChild(ctx context.Context, conn *Connection, id string, body ChildRequest) (ChildResult, error) { var result ChildResult err := executeMutation(ctx, conn, operation{Name: "create child", Method: http.MethodPost, Path: sessionPath(id, "/children"), Body: body, Policy: authRecovery}, []int{201}, func(body []byte, _ int) error { var err error result, err = decodeChild(body) return err }) return result, err}
func GetAgents(ctx context.Context, conn *Connection, session string) (AgentsSnapshot, error) { var result AgentsSnapshot err := executeRead(ctx, conn, operation{Name: "get agents", Method: http.MethodGet, Path: "/agents?session=" + url.QueryEscape(session), Policy: readRecovery}, func(data []byte) error { fields, err := object(data) if err != nil { return err } if err = required(fields, "root", &result.Root); err != nil { return err } var rows []json.RawMessage if err = required(fields, "nodes", &rows); err != nil { return err } result.Nodes = make([]AgentNode, 0, len(rows)) for _, row := range rows { fields, err := object(row) if err != nil { return err } var node AgentNode var sessionData json.RawMessage if err = nullable(fields, "parent", &node.Parent); err != nil { return err } if err = nullable(fields, "address", &node.Address); err != nil { return err } if err = required(fields, "session", &sessionData); err != nil { return err } node.Session, err = decodeSession(sessionData) if err != nil { return err } if err = required(fields, "name", &node.Name); err != nil { return err } if err = required(fields, "depth", &node.Depth); err != nil { return err } if node.Depth < 0 { return fieldError("depth") } if err = required(fields, "running", &node.Running); err != nil { return err } if err = required(fields, "closed", &node.Closed); err != nil { return err } result.Nodes = append(result.Nodes, node) } return nil }) return result, err}
func decodeAgentEvent(raw json.RawMessage) (*AgentEvent, error) { var header struct { Type string `json:"type"` }48 unmodified lines } return &event, nil}
// StreamAgents delivers the first batch, including an empty readiness batch,// then nonempty batches until the stream ends or onBatch fails.// The caller controls cancellation, including any blocking work in onBatch.func StreamAgents(ctx context.Context, conn *Connection, onBatch func([]AgentEvent) error) error { operation := Operation{Name: "stream agents", Method: http.MethodGet, Path: "/agents/stream", Policy: ReadRecovery} first := true err := scanEventStream(ctx, conn, operation, streamLimits{requireSSE: true, lineBytes: 8 * 1024 * 1024}, func(scanner *bufio.Scanner) error { for scanner.Scan() { line := scanner.Text() if !strings.HasPrefix(line, "data:") { continue } var batch struct { Events []json.RawMessage `json:"events"` } if err := json.Unmarshal([]byte(strings.TrimSpace(line[5:])), &batch); err != nil { return streamFailure(StreamProtocol, fmt.Errorf("invalid agent batch: %w", err)) } if batch.Events == nil { return streamFailure(StreamProtocol, errors.New("invalid agent batch: missing events")) } events := make([]AgentEvent, 0, len(batch.Events)) for _, raw := range batch.Events { event, err := decodeAgentEvent(raw) if err != nil { return streamFailure(StreamProtocol, err) } if event != nil { if event.Type == "overflow" && len(batch.Events) != 1 { return streamFailure(StreamProtocol, errors.New("overflow must be the sole agent event")) } events = append(events, *event) } } if len(events) == 0 && !first { continue } first = false if err := onBatch(events); err != nil { return streamFailure(StreamTerminal, err) } if len(events) == 1 && events[0].Type == "overflow" { return nil }func decodeChild(data []byte) (ChildResult, error) { fields, err := object(data) if err != nil { return ChildResult{}, err } var result ChildResult raw, ok := fields["session"] if !ok { return result, fieldError("session") } result.Session, err = decodeSession(raw) if err != nil { return ChildResult{}, err } member, err := object(fields["member"]) if err != nil { return ChildResult{}, err } for _, field := range []struct { target any name string }{{&result.Member.Session, "session"}, {&result.Member.Parent, "parent"}, {&result.Member.Name, "name"}, {&result.Member.Depth, "depth"}, {&result.Member.Closed, "closed"}} { if err := required(member, field.name, field.target); err != nil { return ChildResult{}, err } return nil }) return classifyStreamFailure(err) } return result, nil}47 unmodified lines48495051525253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181195812012112212312416 unmodified lines1411421431448314514614714814915015115215315415515615715815916016116216316416516616716816917017117217389174175176177178941791801811821839918418518618718810418919019147 unmodified lines return "/auth/" + url.PathEscape(provider) + "/accounts/" + url.PathEscape(id)}
func SignInList(ctx context.Context, conn *Connection) (SignIns, error) { return RequestOperation[SignIns](ctx, conn, Operation{Name: "sign in list", Method: http.MethodGet, Path: "/auth", Body: nil, Policy: ReadRecovery}) var result SignIns err := executeRead(ctx, conn, operation{Name: "sign in list", Method: http.MethodGet, Path: "/auth", Policy: readRecovery}, func(data []byte) error { var wire struct { Logins []struct { Provider *string `json:"provider"` Label *string `json:"label"` Detail *string `json:"detail"` Protocol *string `json:"protocol"` } `json:"logins"` Accounts []struct { Provider *string `json:"provider"` ID *string `json:"id"` Label *string `json:"label"` Detail *string `json:"detail"` Selected *bool `json:"selected"` } `json:"accounts"` } if err := json.Unmarshal(data, &wire); err != nil { return err } if wire.Logins == nil { return fieldError("logins") } if wire.Accounts == nil { return fieldError("accounts") } result = SignIns{Logins: make([]SignIn, 0, len(wire.Logins)), Accounts: make([]Account, 0, len(wire.Accounts))} for _, login := range wire.Logins { if login.Provider == nil { return fieldError("provider") } if login.Label == nil { return fieldError("label") } if login.Detail == nil { return fieldError("detail") } if login.Protocol == nil { return fieldError("protocol") } result.Logins = append(result.Logins, SignIn{Provider: *login.Provider, Label: *login.Label, Detail: *login.Detail, Protocol: *login.Protocol}) } for _, account := range wire.Accounts { if account.Provider == nil { return fieldError("provider") } if account.ID == nil { return fieldError("id") } if account.Label == nil { return fieldError("label") } if account.Detail == nil { return fieldError("detail") } if account.Selected == nil { return fieldError("selected") } result.Accounts = append(result.Accounts, Account{Provider: *account.Provider, ID: *account.ID, Label: *account.Label, Detail: *account.Detail, Selected: *account.Selected}) } return nil }) return result, err}
// StartSignIn begins a sign-in and returns the url the browser should open.func StartSignIn(ctx context.Context, conn *Connection, provider string) (StartedSignIn, error) { var result StartedSignIn err := executeMutation(ctx, conn, Operation{Name: "start sign in", Method: http.MethodPost, Path: "/auth/" + url.PathEscape(provider), Policy: AuthRecovery}, []int{http.StatusCreated}, func(data []byte, _ int) error { err := executeMutation(ctx, conn, operation{Name: "start sign in", Method: http.MethodPost, Path: "/auth/" + url.PathEscape(provider), Policy: authRecovery}, []int{http.StatusCreated}, func(data []byte, _ int) error { var wire struct { ID *string `json:"id"` URL *string `json:"url"` }16 unmodified lines return result, err}
func PollSignIn(ctx context.Context, conn *Connection, id string) (SignInStatus, error) { return RequestOperation[SignInStatus](ctx, conn, Operation{Name: "poll sign in", Method: http.MethodGet, Path: "/auth/logins/" + url.PathEscape(id), Body: nil, Policy: ReadRecovery}) var result SignInStatus err := executeRead(ctx, conn, operation{Name: "poll sign in", Method: http.MethodGet, Path: "/auth/logins/" + url.PathEscape(id), Policy: readRecovery}, func(data []byte) error { var wire struct { State *string `json:"state"` Message *string `json:"message"` } if err := json.Unmarshal(data, &wire); err != nil { return err } if wire.State == nil { return fieldError("state") } if wire.Message == nil { return fieldError("message") } switch *wire.State { case "waiting", "exchanging", "done", "failed": default: return fieldError("state") } result = SignInStatus{State: *wire.State, Message: *wire.Message} return nil }) return result, err}
// SignInInput delivers a pasted callback url, code#state, query string, or// bare code; the daemon parses and races it against the browser callback.func SignInInput(ctx context.Context, conn *Connection, id, input string) error { err := acknowledge(ctx, conn, Operation{Name: "sign in input", Method: http.MethodPost, Path: "/auth/logins/" + url.PathEscape(id), Body: map[string]string{"input": input}, Policy: AuthRecovery}) err := acknowledge(ctx, conn, operation{Name: "sign in input", Method: http.MethodPost, Path: "/auth/logins/" + url.PathEscape(id), Body: map[string]string{"input": input}, Policy: authRecovery}) return err}
func CancelSignIn(ctx context.Context, conn *Connection, id string) error { err := acknowledge(ctx, conn, Operation{Name: "cancel sign in", Method: http.MethodDelete, Path: "/auth/logins/" + url.PathEscape(id), Body: nil, Policy: AuthRecovery}) err := acknowledge(ctx, conn, operation{Name: "cancel sign in", Method: http.MethodDelete, Path: "/auth/logins/" + url.PathEscape(id), Body: nil, Policy: authRecovery}) return err}
func SelectAccount(ctx context.Context, conn *Connection, provider, id string) error { err := acknowledge(ctx, conn, Operation{Name: "select account", Method: http.MethodPost, Path: accountPath(provider, id), Body: nil, Policy: AuthRecovery}) err := acknowledge(ctx, conn, operation{Name: "select account", Method: http.MethodPost, Path: accountPath(provider, id), Body: nil, Policy: authRecovery}) return err}
func RemoveAccount(ctx context.Context, conn *Connection, provider, id string) error { err := acknowledge(ctx, conn, Operation{Name: "remove account", Method: http.MethodDelete, Path: accountPath(provider, id), Body: nil, Policy: AuthRecovery}) err := acknowledge(ctx, conn, operation{Name: "remove account", Method: http.MethodDelete, Path: accountPath(provider, id), Body: nil, Policy: authRecovery}) return err}12345678921 unmodified lines31323334353637383940414243444546474849505152535455565758593560616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182383918318418518618718844189190package daemon
import ( "context" "encoding/json" "net/http")
// CapabilityCatalog describes fresh discovery, independently of prepared prompts.21 unmodified lines EffectiveEnabled bool `json:"effective_enabled"` Eligible bool `json:"eligible"`}
type CatalogCapabilityRequest struct { Revision string ID string Scope string Enabled *bool}
type catalogCandidateWire struct { ID *string `json:"id"` Kind *string `json:"kind"` Title *string `json:"title"` Source *string `json:"source"` Valid *bool `json:"valid"` EffectiveEnabled *bool `json:"effective_enabled"` Eligible *bool `json:"eligible"` Description json.RawMessage `json:"description"` ResolvedSource json.RawMessage `json:"resolved_source"` PreferenceKey json.RawMessage `json:"preference_key"` Diagnostic json.RawMessage `json:"diagnostic"` ShadowedBy json.RawMessage `json:"shadowed_by"` GlobalPreference json.RawMessage `json:"global_preference"` SessionOverride json.RawMessage `json:"session_override"`}
func GetCapabilityCatalog(ctx context.Context, conn *Connection, session string) (CapabilityCatalog, error) { return RequestOperation[CapabilityCatalog](ctx, conn, Operation{Name: "get capability catalog", Method: http.MethodGet, Path: sessionPath(session, "/catalog"), Policy: ReadRecovery}) var result CapabilityCatalog err := executeRead(ctx, conn, operation{Name: "get capability catalog", Method: http.MethodGet, Path: sessionPath(session, "/catalog"), Policy: readRecovery}, func(data []byte) error { var wire struct { Workspace *string `json:"workspace"` Revision *string `json:"revision"` Extensions map[string]*bool `json:"extensions"` Diagnostics *stringCollection `json:"diagnostics"` Candidates []catalogCandidateWire `json:"candidates"` } if err := json.Unmarshal(data, &wire); err != nil { return err } if wire.Workspace == nil { return fieldError("workspace") } if wire.Revision == nil { return fieldError("revision") } if wire.Extensions == nil { return fieldError("extensions") } if wire.Diagnostics == nil { return fieldError("diagnostics") } if wire.Candidates == nil { return fieldError("candidates") } result = CapabilityCatalog{ Workspace: *wire.Workspace, Revision: *wire.Revision, Extensions: make(map[string]bool, len(wire.Extensions)), Diagnostics: []string(*wire.Diagnostics), Candidates: make([]CatalogCandidate, 0, len(wire.Candidates)), } for name, enabled := range wire.Extensions { if enabled == nil { return fieldError("extensions." + name) } result.Extensions[name] = *enabled } for _, wire := range wire.Candidates { if wire.ID == nil { return fieldError("id") } if wire.Kind == nil { return fieldError("kind") } if wire.Title == nil { return fieldError("title") } if wire.Source == nil { return fieldError("source") } if wire.Valid == nil { return fieldError("valid") } if wire.EffectiveEnabled == nil { return fieldError("effective_enabled") } if wire.Eligible == nil { return fieldError("eligible") } candidate := CatalogCandidate{ ID: *wire.ID, Kind: *wire.Kind, Title: *wire.Title, Source: *wire.Source, Valid: *wire.Valid, EffectiveEnabled: *wire.EffectiveEnabled, Eligible: *wire.Eligible, } if candidate.ID == "" { return fieldError("id") } if wire.Description == nil { return fieldError("description") } if err := json.Unmarshal(wire.Description, &candidate.Description); err != nil { return err } if wire.ResolvedSource == nil { return fieldError("resolved_source") } if err := json.Unmarshal(wire.ResolvedSource, &candidate.ResolvedSource); err != nil { return err } if wire.PreferenceKey == nil { return fieldError("preference_key") } if err := json.Unmarshal(wire.PreferenceKey, &candidate.PreferenceKey); err != nil { return err } if wire.Diagnostic == nil { return fieldError("diagnostic") } if err := json.Unmarshal(wire.Diagnostic, &candidate.Diagnostic); err != nil { return err } if wire.ShadowedBy == nil { return fieldError("shadowed_by") } if err := json.Unmarshal(wire.ShadowedBy, &candidate.ShadowedBy); err != nil { return err } if wire.GlobalPreference == nil { return fieldError("global_preference") } if err := json.Unmarshal(wire.GlobalPreference, &candidate.GlobalPreference); err != nil { return err } if wire.SessionOverride == nil { return fieldError("session_override") } if err := json.Unmarshal(wire.SessionOverride, &candidate.SessionOverride); err != nil { return err } result.Candidates = append(result.Candidates, candidate) } return nil }) return result, err}
func SetCatalogCapability(ctx context.Context, conn *Connection, session, revision, id, scope string, enabled *bool) (ReloadResult, error) { return reloadSettings(ctx, conn, Operation{Name: "set catalog capability", Method: http.MethodPost, Path: sessionPath(session, "/catalog"), Body: struct {func SetCatalogCapability(ctx context.Context, conn *Connection, session string, request CatalogCapabilityRequest) (ReloadResult, error) { return reloadSettings(ctx, conn, operation{Name: "set catalog capability", Method: http.MethodPost, Path: sessionPath(session, "/catalog"), Body: struct { Enabled *bool `json:"enabled"` Revision string `json:"revision"` ID string `json:"id"` Scope string `json:"scope"` }{Enabled: enabled, Revision: revision, ID: id, Scope: scope}, Policy: AuthRecovery}) }{Enabled: request.Enabled, Revision: request.Revision, ID: request.ID, Scope: request.Scope}, Policy: authRecovery})}123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309// 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) }}