diff --git a/args.go b/args.go index c21489f4e..f5e0807e0 100644 --- a/args.go +++ b/args.go @@ -951,4 +951,16 @@ const ( // ArgAgentArchive treats the workspace payload as a tar archive. ArgAgentArchive = "archive" + + // ArgAgentProxyType selects which coding-agent protocol `start-proxy` impersonates. + ArgAgentProxyType = "type" + + // ArgAgentProxySession is the session ID or name the proxy bridges to. + ArgAgentProxySession = "session" + + // ArgAgentProxyPort is the local port `start-proxy` listens on. + ArgAgentProxyPort = "port" + + // ArgAgentProxyReplay replays the session's event history into the first thread on connect. + ArgAgentProxyReplay = "replay" ) diff --git a/commands/agents.go b/commands/agents.go index eed892ecd..735fefc06 100644 --- a/commands/agents.go +++ b/commands/agents.go @@ -29,6 +29,7 @@ import ( "io" "net/http" "os" + "os/signal" "path/filepath" "regexp" "strings" @@ -41,6 +42,8 @@ import ( "github.com/digitalocean/doctl/commands/charm" "github.com/digitalocean/doctl/commands/displayers" "github.com/digitalocean/doctl/do" + "github.com/digitalocean/doctl/internal/agentproxy" + "github.com/digitalocean/doctl/internal/agentproxy/codex" "github.com/digitalocean/godo" "github.com/muesli/termenv" "github.com/spf13/cobra" @@ -231,6 +234,22 @@ Use `+"`"+`--name`+"`"+` to name the session (this sets the manifest's `+"`"+`me AddStringFlag(cmdStart, doctl.ArgAgentName, "", "", "Name for the new session (sets the manifest's metadata.name). If omitted, the server auto-generates a name. Must be unique among your team's active sessions.") cmdStart.Example = `doctl agents start --spec agent-spec.yaml --name my-session` + cmdStartProxy := CmdBuilder(cmd, RunAgentsStartProxy, "start-proxy", + "Run a local facade that lets a coding-agent CLI drive a hosted session", + `Starts a local WebSocket server that impersonates a coding-agent's own app-server protocol, so the unmodified CLI can attach to a hosted session as if it were a local backend. + +`+"`"+`--type`+"`"+` selects which protocol to impersonate (v1: `+"`"+`codex`+"`"+` only; future agents get their own facade behind this same command, not a new command). Once `+"`"+`start-proxy`+"`"+` is listening, connect the real CLI, e.g. `+"`"+`codex --remote ws://127.0.0.1:1144`+"`"+`. + +Tested against `+"`"+`codex-cli `+codex.TestedVersion+"`"+`. Codex's WS/app-server transport is officially experimental and can change without notice — re-verify the protocol capture on every codex upgrade before trusting this against a newer CLI. + +The harness allows one connected device per session — a concurrent `+"`"+`doctl agents attach`+"`"+` on the same session conflicts with the proxy's own stream. Close one before opening the other.`, + Writer) + AddStringFlag(cmdStartProxy, doctl.ArgAgentProxyType, "", "codex", "Coding-agent protocol to impersonate (v1: codex)") + AddStringFlag(cmdStartProxy, doctl.ArgAgentProxySession, "", "", "Session ID or name to bridge to", requiredOpt()) + AddIntFlag(cmdStartProxy, doctl.ArgAgentProxyPort, "", 1144, "Local port to listen on") + AddBoolFlag(cmdStartProxy, doctl.ArgAgentProxyReplay, "", false, "Replay the session's event history into the first thread on connect (not yet implemented)") + cmdStartProxy.Example = `doctl agents start-proxy --type codex --session my-session --port 1144` + cmdAttach := CmdBuilder(cmd, RunAgentsAttach, "attach ", "Attach to an agent session", `Opens an interactive line-mode TUI on an existing session. Streams events from the server and accepts typed input. If the SSE connection drops, doctl shows Reconnecting... and retries automatically (5 attempts with backoff). If reconnection fails, it prints an error and stops the stream. @@ -345,6 +364,54 @@ func RunAgentsStart(c *CmdConfig) error { return c.Display(&displayers.HostedAgentSession{Sessions: []do.HostedAgentSession{*sess}}) } +// RunAgentsStartProxy runs a local WebSocket facade that impersonates a +// coding-agent's own app-server protocol, bridging an unmodified agent CLI to +// a hosted session. +func RunAgentsStartProxy(c *CmdConfig) error { + proxyType, err := c.Doit.GetString(c.NS, doctl.ArgAgentProxyType) + if err != nil { + return err + } + if proxyType != "codex" { + return fmt.Errorf("unsupported --type %q; v1 supports only \"codex\"", proxyType) + } + + sessionRef, err := c.Doit.GetString(c.NS, doctl.ArgAgentProxySession) + if err != nil { + return err + } + port, err := c.Doit.GetInt(c.NS, doctl.ArgAgentProxyPort) + if err != nil { + return err + } + // --replay is accepted now so the CLI surface doesn't change again in M1, + // but session-history replay isn't implemented until the harness bridge + // lands; read it only so an explicit --replay=true doesn't silently no-op + // without the flag help already having said so. + if _, err := c.Doit.GetBool(c.NS, doctl.ArgAgentProxyReplay); err != nil { + return err + } + + svc := c.HostedAgents() + sessionID, err := resolveSessionRef(svc, sessionRef) + if err != nil { + return err + } + if _, err := svc.GetSession(sessionID); err != nil { + return fmt.Errorf("session %q not found: %w", sessionRef, err) + } + + fmt.Fprintf(c.Out, "Proxying session %s as a codex app-server on ws://127.0.0.1:%d\n", sessionID, port) + fmt.Fprintf(c.Out, "Connect with: codex --remote ws://127.0.0.1:%d\n", port) + + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) + defer cancel() + + return agentproxy.Serve(ctx, port, func() agentproxy.Facade { + return &codex.Facade{SessionID: sessionID, Sessions: svc} + }) +} + // readManifest returns the spec file as raw bytes. path "-" reads from stdin. // The only client-side validation is "non-empty after trim" so a stray // `--spec /dev/null` fails fast instead of hitting the server. diff --git a/internal/agentproxy/agentproxytest/harness.go b/internal/agentproxy/agentproxytest/harness.go new file mode 100644 index 000000000..8222fef4c --- /dev/null +++ b/internal/agentproxy/agentproxytest/harness.go @@ -0,0 +1,172 @@ +// Package agentproxytest fakes the DigitalOcean Hosted Agents session HTTP +// API (see github.com/digitalocean/godo's hosted_agents.go) behind an +// httptest.Server, so agentproxy tests can drive a real +// do.HostedAgentsService end to end without a live harness backend. +// +// Nothing outside this package's own tests (or other packages' _test.go +// files) should ever import it: it's test scaffolding, not production code, +// even though it lives in an ordinary buildable package rather than a +// _test.go file — that's what lets other packages' tests (e.g. +// internal/agentproxy/codex's) import it across package boundaries. +package agentproxytest + +import ( + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" +) + +// Event is one canned SSE event the harness streams back from +// GET /v2/agents/sessions/{id}/stream, matching the event-specific part of +// godo.HostedAgentEvent's wire shape (see HostedAgentEventKind's doc comment +// for the canonical type strings). +type Event struct { + // Type is the wire "type" discriminator, e.g. "run.started", + // "run.token_delta", "run.completed", "run.failed". + Type string + // Data is the event-specific payload. A nil Data marshals as "{}". + Data json.RawMessage +} + +// eventWire is the on-the-wire SPI canonical event envelope this harness +// emits. It mirrors godo's unexported hostedAgentEventWire field-for-field +// (confirmed against vendor/github.com/digitalocean/godo/hosted_agents.go) +// rather than importing it, since that type isn't exported. +type eventWire struct { + EventID string `json:"event_id"` + RunID string `json:"run_id"` + TenantID string `json:"tenant_id"` + SessionID string `json:"session_id"` + Timestamp string `json:"timestamp"` + Seq int `json:"seq"` + Type string `json:"type"` + Data json.RawMessage `json:"data"` +} + +// Harness fakes the four Hosted Agents session routes an agentproxy test +// needs: session lookup, SSE stream, send-input, and HITL resolve. Point a +// *godo.Client at it with godo.SetBaseURL(h.Server.URL+"/") — do's +// NewHostedAgentsService wraps that client directly, so a Facade built +// against it talks to this harness exactly as it would a real backend. +type Harness struct { + Server *httptest.Server + + mu sync.Mutex + sessionID string + runID string // returned by the next POST .../input call + events []Event // streamed, in order, by the next GET .../stream call +} + +// New starts the fake harness and registers its shutdown via t.Cleanup. +// sessionID is the session id every route responds under — a test's +// Facade.SessionID should match it. +func New(t *testing.T, sessionID string) *Harness { + t.Helper() + h := &Harness{sessionID: sessionID} + + mux := http.NewServeMux() + mux.HandleFunc("GET /v2/agents/sessions/{id}", h.handleGetSession) + mux.HandleFunc("GET /v2/agents/sessions/{id}/stream", h.handleStream) + mux.HandleFunc("POST /v2/agents/sessions/{id}/input", h.handleInput) + mux.HandleFunc("POST /v2/agents/sessions/{id}/hitl/{requestID}", h.handleHITL) + + h.Server = httptest.NewServer(mux) + t.Cleanup(h.Server.Close) + return h +} + +// QueueRun arranges for the next POST .../input call to return runID, and +// for GET .../stream to then emit events (in order, tagged with runID), +// flushing after each one so a concurrent reader observes them incrementally +// rather than all at once at EOF. +// +// One (runID, events) pair is enough for a single-turn sequence test; queue +// again before the next turn/start if a future test drives more than one +// turn over the same harness. +func (h *Harness) QueueRun(runID string, events ...Event) { + h.mu.Lock() + defer h.mu.Unlock() + h.runID = runID + h.events = events +} + +func (h *Harness) handleGetSession(w http.ResponseWriter, _ *http.Request) { + h.mu.Lock() + sessionID := h.sessionID + h.mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"session":{"session_id":%q,"team_id":15726539,`+ + `"agent_kind":"AGENT_KIND_CODEX_CLI","status":"SESSION_STATUS_READY",`+ + `"created_at":"2026-01-01T00:00:00Z","last_event_at":"2026-01-01T00:00:00Z"}}`, + sessionID) +} + +func (h *Harness) handleInput(w http.ResponseWriter, r *http.Request) { + h.mu.Lock() + runID := h.runID + h.mu.Unlock() + + if runID == "" { + http.Error(w, "agentproxytest: no run id queued; call Harness.QueueRun first", http.StatusInternalServerError) + return + } + + // Best-effort decode only — the harness doesn't assert on the input + // text; callers that care can inspect the request in a custom handler. + var body struct { + Text string `json:"text"` + } + _ = json.NewDecoder(r.Body).Decode(&body) + + w.Header().Set("Content-Type", "application/json") + _ = json.NewEncoder(w).Encode(map[string]string{"run_id": runID}) +} + +func (h *Harness) handleStream(w http.ResponseWriter, _ *http.Request) { + h.mu.Lock() + runID := h.runID + events := h.events + sessionID := h.sessionID + h.mu.Unlock() + + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher, canFlush := w.(http.Flusher) + + for i, ev := range events { + data := ev.Data + if data == nil { + data = json.RawMessage("{}") + } + wire := eventWire{ + EventID: fmt.Sprintf("evt-%d", i), + RunID: runID, + TenantID: "15726539", + SessionID: sessionID, + Timestamp: "2026-01-01T00:00:00Z", + Seq: i, + Type: ev.Type, + Data: data, + } + encoded, err := json.Marshal(wire) + if err != nil { + // A test handed the harness a payload that can't marshal — a + // programmer error in the test, not a runtime condition. + panic(fmt.Sprintf("agentproxytest: event does not marshal to JSON: %v", err)) + } + + fmt.Fprintf(w, "event: %s\ndata: %s\n\n", ev.Type, encoded) + if canFlush { + flusher.Flush() + } + } +} + +func (h *Harness) handleHITL(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) +} diff --git a/internal/agentproxy/codex/facade.go b/internal/agentproxy/codex/facade.go new file mode 100644 index 000000000..3b50cf917 --- /dev/null +++ b/internal/agentproxy/codex/facade.go @@ -0,0 +1,787 @@ +// Package codex implements agentproxy.Facade for the codex CLI's app-server +// JSON-RPC protocol (github.com/openai/codex, codex-rs/app-server). +package codex + +import ( + "context" + "encoding/json" + "fmt" + "log" + "strings" + "sync" + "time" + + "github.com/digitalocean/doctl/do" + "github.com/digitalocean/doctl/internal/agentproxy" + "github.com/digitalocean/godo" +) + +// TestedVersion is the codex-cli version this facade's protocol handling was +// captured against (hosted-agents' docs/design/codex-app-server-protocol-capture.md: +// stdio send-message-v2 capture plus a live codex --remote / app-server pairing +// over WebSocket). Re-run that capture and update this pin on every codex +// upgrade — the WS/app-server transport is officially experimental and can +// change without notice. +const TestedVersion = "0.142.5" + +// defaultModel is the single model this facade advertises via model/list and +// echoes back from thread/start, kept as one constant so the two stubbed +// responses can't drift apart. Matches agent-spec.yaml's MODEL today; a real +// per-session model comes from the harness, not this facade, once it's wired. +const defaultModel = "gpt-5.5" + +// workspaceRoot is the fixed sandbox workspace root every hosted session +// uses (see doctl agents upload/download's --workspace-path help). +const workspaceRoot = "/workspace" + +// Facade implements agentproxy.Facade for codex --remote. +type Facade struct { + // SessionID is the hosted session this facade bridges to. thread/start's + // synthesized thread id is exactly this value — simplest possible way to + // satisfy "a synthesized thread whose id embeds the session id" while + // matching what the real capture showed too (thread.id == thread.sessionId + // there as well). + SessionID string + + // Sessions is the harness bridge: turn/start calls SendInput on it, and + // the background streaming goroutine reads StreamSession's SSE events to + // translate into codex notifications. The bootstrap methods (initialize, + // thread/start, etc.) don't touch it. + Sessions do.HostedAgentsService + + // notifier is set once per connection by the bridge (see SetNotifier) and + // used by the background turn-streaming goroutine to push notifications + // (turn/started, item/agentMessage/delta, turn/completed, ...) on the + // harness's timeline rather than in direct reply to a request. + notifier agentproxy.Notifier + + // mu guards turns and streamStarted. The harness allows exactly one + // connected SSE consumer per session (the same constraint that makes a + // concurrent `doctl agents attach` conflict with this proxy's own + // stream) — so this facade must never have more than one StreamSession + // call open at a time, no matter how many turns run, one after another or + // overlapping. Opening a second one mid-turn silently displaces the + // first, cutting off that turn's remaining events: typing a second + // message before the first turn finishes made only the second message + // ever get a response, because the first turn's stream got evicted the + // moment the second one opened its own. streamStarted tracks whether a + // StreamSession call is currently outstanding, not merely whether one was + // ever made: stopEventLoop clears it on exit so a later turn/start can + // open a fresh one if the previous one died. + mu sync.Mutex + turns map[string]*turnState + streamStarted bool +} + +// turnState is the in-flight accumulation state for one turn, keyed by run +// id in Facade.turns. All turns share one event loop (runEventLoop), so this +// replaces what used to be local variables inside a per-turn goroutine. +type turnState struct { + itemID string + text strings.Builder + startedAt int64 + itemStarted bool +} + +var _ agentproxy.Facade = (*Facade)(nil) +var _ agentproxy.NotifierAware = (*Facade)(nil) + +// SetNotifier implements agentproxy.NotifierAware. +func (f *Facade) SetNotifier(n agentproxy.Notifier) { + f.notifier = n +} + +// initializeResult is the shape a real codex app-server returns for +// `initialize`, captured both over stdio (send-message-v2) and from the live +// TUI over WebSocket: {"result":{"userAgent":...,"codexHome":..., +// "platformFamily":...,"platformOs":...}}. codexHome/platformFamily/platformOs +// describe codex's own local runtime environment in a real app-server; since +// this facade proxies to a hosted sandbox rather than running codex locally, +// these are synthetic placeholders, not a reflection of anything the sandbox +// actually reports. +// +// TODO(M3-M5): replace with the sandbox's real values once the harness +// exposes them. +type initializeResult struct { + UserAgent string `json:"userAgent"` + CodexHome string `json:"codexHome"` + PlatformFamily string `json:"platformFamily"` + PlatformOs string `json:"platformOs"` +} + +// accountReadResult is the shape a real codex app-server returns for +// `account/read`. Discovered to be load-bearing by testing against a real +// `codex --remote`: the TUI treats a failed account/read as fatal during +// bootstrap ("Error: account/read failed during TUI bootstrap") rather than +// tolerating it like other unhandled methods, so it needs a real (if +// synthetic) success response for the TUI to render at all. +// requiresOpenaiAuth is false because auth to the model happens server-side +// in the hosted sandbox, not via any local OpenAI login flow the TUI would +// otherwise try to trigger. +type accountReadResult struct { + Account struct { + Type string `json:"type"` + Email string `json:"email"` + PlanType string `json:"planType"` + } `json:"account"` + RequiresOpenaiAuth bool `json:"requiresOpenaiAuth"` +} + +// modelListResult is the shape a real codex app-server returns for +// `model/list`. Discovered fatal-if-missing the same way as account/read: the +// TUI bootstrap errors out on a failed model/list too. One entry matching +// agent-spec.yaml's default model (gpt-5.5) is enough to clear bootstrap. +// +// TODO(M3-M5): a real per-session model catalog, once the harness exposes +// one. +type modelListResult struct { + Data []modelInfo `json:"data"` + NextCursor *string `json:"nextCursor"` +} + +type modelInfo struct { + ID string `json:"id"` + Model string `json:"model"` + Upgrade *string `json:"upgrade"` + UpgradeInfo *string `json:"upgradeInfo"` + AvailabilityNux *string `json:"availabilityNux"` + DisplayName string `json:"displayName"` + Description string `json:"description"` + Hidden bool `json:"hidden"` + SupportedReasoningEfforts []reasoningEffortInfo `json:"supportedReasoningEfforts"` + DefaultReasoningEffort string `json:"defaultReasoningEffort"` + InputModalities []string `json:"inputModalities"` + SupportsPersonality bool `json:"supportsPersonality"` + AdditionalSpeedTiers []string `json:"additionalSpeedTiers"` + ServiceTiers []string `json:"serviceTiers"` + DefaultServiceTier *string `json:"defaultServiceTier"` + IsDefault bool `json:"isDefault"` +} + +type reasoningEffortInfo struct { + ReasoningEffort string `json:"reasoningEffort"` + Description string `json:"description"` +} + +// threadStartResult is the shape a real codex app-server returns for +// thread/start. Discovered fatal-if-missing the same way as account/read and +// model/list: the TUI's bootstrap sequence calls all three synchronously and +// errors out if any fails. This is a synthesized stub — SessionID as the +// thread id, static policy/sandbox defaults matching what the real capture +// showed. +// +// TODO(M3-M5): Turns always reports empty even once real turns run through +// turn/start; feed real turn history back into this shape. +type threadStartResult struct { + Thread thread `json:"thread"` + Model string `json:"model"` + ModelProvider string `json:"modelProvider"` + Cwd string `json:"cwd"` + RuntimeWorkspaceRoots []string `json:"runtimeWorkspaceRoots"` + ApprovalPolicy string `json:"approvalPolicy"` + ApprovalsReviewer string `json:"approvalsReviewer"` + Sandbox sandboxInfo `json:"sandbox"` + MultiAgentMode string `json:"multiAgentMode"` +} + +type sandboxInfo struct { + Type string `json:"type"` + NetworkAccess bool `json:"networkAccess"` +} + +type thread struct { + ID string `json:"id"` + SessionID string `json:"sessionId"` + ForkedFromID *string `json:"forkedFromId"` + ParentThreadID *string `json:"parentThreadId"` + Preview string `json:"preview"` + Ephemeral bool `json:"ephemeral"` + ModelProvider string `json:"modelProvider"` + CreatedAt int64 `json:"createdAt"` + UpdatedAt int64 `json:"updatedAt"` + RecencyAt int64 `json:"recencyAt"` + Status threadStatus `json:"status"` + Path string `json:"path"` + Cwd string `json:"cwd"` + CliVersion string `json:"cliVersion"` + // Source is a fixed placeholder, not a reflection of the real client: the + // live-TUI protocol capture showed a real app-server return "vscode" here + // even when the connecting client was the terminal TUI, not VS Code. Kept + // identical since that's the one value already proven to decode cleanly. + Source string `json:"source"` + ThreadSource string `json:"threadSource"` + AgentNickname *string `json:"agentNickname"` + AgentRole *string `json:"agentRole"` + GitInfo *string `json:"gitInfo"` + Name *string `json:"name"` + Turns []any `json:"turns"` +} + +type threadStatus struct { + Type string `json:"type"` +} + +// threadResumeParams is ThreadResumeParams from the codex source +// (codex-rs/app-server-protocol/src/protocol/v2/thread.rs): only ThreadID is +// read here. history/path are experimental, Codex-Cloud-only fields that +// don't apply to a hosted-session proxy. +type threadResumeParams struct { + ThreadID string `json:"threadId"` +} + +// threadResumeResult is a distinct name for documentation purposes only — +// ThreadResumeResponse is byte-for-byte identical to ThreadStartResponse in +// the codex source, so the underlying fields (and JSON tags) are shared via +// threadStartResult. +type threadResumeResult threadStartResult + +// hooksListResult is HooksListResponse (plugin.rs). data is required but an +// empty list is valid — no hosted-session hooks catalog exists yet. +type hooksListResult struct { + Data []any `json:"data"` +} + +// skillsListResult is SkillsListResponse (plugin.rs). Same story as hooks: +// empty is valid, and clears the "failed to load skills" toast the TUI +// otherwise shows when this method 404s. +type skillsListResult struct { + Data []any `json:"data"` +} + +// pluginListResult is PluginListResponse (plugin.rs). marketplaces is +// required; the other two fields have #[serde(default)] on the codex side +// but are included explicitly since Go's zero value for a nil slice +// marshals to `null`, not `[]`, and the real server always sends arrays. +type pluginListResult struct { + Marketplaces []any `json:"marketplaces"` + MarketplaceLoadErrors []any `json:"marketplaceLoadErrors"` + FeaturedPluginIDs []string `json:"featuredPluginIds"` +} + +// appListResult is AppsListResponse (apps.rs). Both fields are required keys +// on the wire (nextCursor is a required key with a nullable value, not an +// omittable one). +type appListResult struct { + Data []any `json:"data"` + NextCursor *string `json:"nextCursor"` +} + +// threadUnsubscribeResult is ThreadUnsubscribeResponse (thread.rs). status is +// one of "notLoaded" | "notSubscribed" | "unsubscribed" — "unsubscribed" is +// always a truthful, safe answer here since this facade holds no real +// subscription state to be wrong about. +type threadUnsubscribeResult struct { + Status string `json:"status"` +} + +// --- turn/start, turn/interrupt, and the notifications a turn produces --- +// +// Shapes confirmed against the codex source +// (codex-rs/app-server-protocol/src/protocol/v2/turn.rs and thread_data.rs), +// not guessed. Worth calling out: there is no `turn/failed` method anywhere +// in the codex protocol. Failure is represented purely via `turn/completed` +// with `turn.status == "failed"` and `turn.error` populated — finishTurn +// below reflects that; there is no separate failed-turn notification to +// send. + +// turnStartParams is TurnStartParams (v2/turn.rs) — only the two fields this +// v1 "one-way text" facade reads; the many optional per-turn overrides +// (model, sandboxPolicy, approvalPolicy, ...) aren't applied here. +type turnStartParams struct { + ThreadID string `json:"threadId"` + Input []userInputItem `json:"input"` +} + +// userInputItem is one variant of the UserInput enum (v2/turn.rs), internally +// tagged by "type". Only the "text" variant is read; image/localImage/skill/ +// mention inputs are silently ignored in v1. +type userInputItem struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` +} + +// turnStartResult is TurnStartResponse (v2/turn.rs) — just the Turn itself. +type turnStartResult struct { + Turn turnObj `json:"turn"` +} + +// turnObj is Turn (thread_data.rs). itemsView is omitted: it has +// #[serde(default)] on the codex side, so leaving it out entirely (rather +// than sending a zero value that might not match the default variant) is the +// safer decode-safe choice, same reasoning as the omitted optional fields on +// thread/start's Thread. +type turnObj struct { + ID string `json:"id"` + Items []any `json:"items"` + Status string `json:"status"` + Error *turnError `json:"error"` + StartedAt *int64 `json:"startedAt"` + CompletedAt *int64 `json:"completedAt"` + DurationMs *int64 `json:"durationMs"` +} + +type turnError struct { + Message string `json:"message"` +} + +// turnStartedNotification is TurnStartedNotification (v2/turn.rs), method +// "turn/started" — maps from the canonical run.started event. +type turnStartedNotification struct { + ThreadID string `json:"threadId"` + Turn turnObj `json:"turn"` +} + +// turnCompletedNotification is TurnCompletedNotification, method +// "turn/completed" — maps from both canonical run.completed AND run.failed +// (see the no-turn/failed note above). +type turnCompletedNotification struct { + ThreadID string `json:"threadId"` + Turn turnObj `json:"turn"` +} + +// itemStartedNotification is ItemStartedNotification, method "item/started". +type itemStartedNotification struct { + Item agentMessageItem `json:"item"` + ThreadID string `json:"threadId"` + TurnID string `json:"turnId"` + StartedAtMs int64 `json:"startedAtMs"` +} + +// itemCompletedNotification is ItemCompletedNotification, method +// "item/completed". +type itemCompletedNotification struct { + Item agentMessageItem `json:"item"` + ThreadID string `json:"threadId"` + TurnID string `json:"turnId"` + CompletedAtMs int64 `json:"completedAtMs"` +} + +// agentMessageItem is the ThreadItem::AgentMessage variant, internally +// tagged "type": "agentMessage". phase/memoryCitation are omitted — both +// optional on the codex side. +type agentMessageItem struct { + Type string `json:"type"` + ID string `json:"id"` + Text string `json:"text"` +} + +// agentMessageDeltaNotification is AgentMessageDeltaNotification, method +// "item/agentMessage/delta" — maps from canonical run.token_delta. +type agentMessageDeltaNotification struct { + ThreadID string `json:"threadId"` + TurnID string `json:"turnId"` + ItemID string `json:"itemId"` + Delta string `json:"delta"` +} + +// turnInterruptResult is TurnInterruptResponse (v2/turn.rs) — genuinely +// empty on the wire. Best-effort no-op for now: the harness has no +// cancel-input surface yet. +// +// TODO(M3-M5): wire this to a real cancel once the harness exposes one. +type turnInterruptResult struct{} + +// synthesizedThread builds the one thread this facade ever knows about: a +// stand-in for the hosted session, id equal to SessionID, with static +// policy/sandbox defaults matching what the protocol capture showed for a +// real thread/start. Shared by thread/start and thread/resume so the two +// can't drift apart. Turns is always an empty placeholder — see the +// threadStartResult TODO above. +func (f *Facade) synthesizedThread() threadStartResult { + now := time.Now().Unix() + return threadStartResult{ + Thread: thread{ + ID: f.SessionID, + SessionID: f.SessionID, + ModelProvider: "openai", + CreatedAt: now, + UpdatedAt: now, + RecencyAt: now, + Status: threadStatus{Type: "idle"}, + Path: workspaceRoot + "/.codex-session-" + f.SessionID + ".jsonl", + Cwd: workspaceRoot, + CliVersion: TestedVersion, + Source: "vscode", + ThreadSource: "user", + Turns: []any{}, + }, + Model: defaultModel, + ModelProvider: "openai", + Cwd: workspaceRoot, + RuntimeWorkspaceRoots: []string{workspaceRoot}, + ApprovalPolicy: "on-request", + ApprovalsReviewer: "user", + Sandbox: sandboxInfo{Type: "readOnly", NetworkAccess: false}, + MultiAgentMode: "explicitRequestOnly", + } +} + +// Dispatch implements agentproxy.Facade. +func (f *Facade) Dispatch(ctx context.Context, method string, params json.RawMessage) (any, error) { + switch method { + case "initialize": + return initializeResult{ + UserAgent: "doctl-agents-proxy/" + TestedVersion, + CodexHome: "/hosted", + PlatformFamily: "unix", + PlatformOs: "linux", + }, nil + + case "initialized": + // Client notification (no id, no reply expected); this facade has + // nothing to do in response. + return nil, nil + + case "account/read": + result := accountReadResult{RequiresOpenaiAuth: false} + result.Account.Type = "chatgpt" + result.Account.Email = "hosted-session@digitalocean.com" + result.Account.PlanType = "free" + return result, nil + + case "model/list": + return modelListResult{ + Data: []modelInfo{{ + ID: defaultModel, + Model: defaultModel, + DisplayName: defaultModel, + Description: "Model configured for this hosted session.", + SupportedReasoningEfforts: []reasoningEffortInfo{ + {ReasoningEffort: "medium", Description: "Balances speed and reasoning depth for everyday tasks"}, + }, + DefaultReasoningEffort: "medium", + InputModalities: []string{"text"}, + AdditionalSpeedTiers: []string{}, + ServiceTiers: []string{}, + IsDefault: true, + }}, + }, nil + + case "thread/start": + return f.synthesizedThread(), nil + + case "thread/resume": + // ThreadResumeParams.thread_id (camelCase threadId on the wire) is the + // only field this stub reads — history/path are experimental, + // codex-cloud-only fields per the codex source and don't apply here. + var p threadResumeParams + if err := json.Unmarshal(params, &p); err != nil { + return nil, &agentproxy.RPCError{Code: -32602, Message: "invalid params: " + err.Error()} + } + // This proxy is bound 1:1 to a single hosted session for its whole + // lifetime (see Facade.SessionID), so the only thread id it can ever + // resume is its own — anything else is a real "not found," not a gap + // to stub around. + if p.ThreadID != f.SessionID { + return nil, &agentproxy.RPCError{ + Code: -32001, + Message: fmt.Sprintf("thread %q not found; this proxy only serves thread %q", p.ThreadID, f.SessionID), + } + } + // ThreadResumeResponse is identical in shape to ThreadStartResponse + // (verified against the codex source, codex-rs/app-server-protocol/ + // src/protocol/v2/thread.rs) — same synthesized thread either way. + return threadResumeResult(f.synthesizedThread()), nil + + // The remaining cases are real methods the TUI calls during/after + // bootstrap that don't crash it (unlike account/read, model/list, + // thread/start above) but, left unhandled, get logged and — for + // skills/list — surface a visible (non-fatal) error toast. Shapes + // confirmed against the codex source + // (codex-rs/app-server-protocol/src/protocol/v2/{plugin,apps,thread}.rs), + // not guessed — empty catalogs are valid, decode-safe responses. + // + // TODO(M3-M5): a real per-session hooks/skills/plugin/app catalog, once + // the harness exposes one. + case "hooks/list": + return hooksListResult{Data: []any{}}, nil + + case "skills/list": + return skillsListResult{Data: []any{}}, nil + + case "plugin/list": + return pluginListResult{ + Marketplaces: []any{}, + MarketplaceLoadErrors: []any{}, + FeaturedPluginIDs: []string{}, + }, nil + + case "app/list": + return appListResult{Data: []any{}}, nil + + case "thread/unsubscribe": + // thread_id isn't checked: unsubscribing from the one thread this + // facade ever serves (or from anything else) is always a safe no-op. + return threadUnsubscribeResult{Status: "unsubscribed"}, nil + + case "turn/start": + var p turnStartParams + if err := json.Unmarshal(params, &p); err != nil { + return nil, &agentproxy.RPCError{Code: -32602, Message: "invalid params: " + err.Error()} + } + if p.ThreadID != f.SessionID { + return nil, &agentproxy.RPCError{ + Code: -32001, + Message: fmt.Sprintf("thread %q not found; this proxy only serves thread %q", p.ThreadID, f.SessionID), + } + } + // Open (or confirm) the event-loop stream before SendInput, not after: + // SendInput is what makes the harness start the run, so a reader + // attached only afterward can race a fast run to completion and miss + // every event it emits. See ensureEventLoop. + if err := f.ensureEventLoop(ctx); err != nil { + return nil, &agentproxy.RPCError{Code: -32000, Message: "opening event stream failed: " + err.Error()} + } + + var text strings.Builder + for _, item := range p.Input { + if item.Type != "text" { + continue // image/localImage/skill/mention: out of scope for v1 "one-way text" + } + if text.Len() > 0 { + text.WriteString("\n") + } + text.WriteString(item.Text) + } + resp, err := f.Sessions.SendInput(f.SessionID, &godo.HostedAgentSendInputRequest{Text: text.String()}) + if err != nil { + return nil, &agentproxy.RPCError{Code: -32000, Message: "SendInput failed: " + err.Error()} + } + // The harness's run id doubles as this facade's turn id — already a + // unique per-turn identifier, no separate id scheme needed. + f.trackTurn(resp.RunID) + return turnStartResult{ + Turn: turnObj{ID: resp.RunID, Items: []any{}, Status: "inProgress"}, + }, nil + + case "turn/interrupt": + return turnInterruptResult{}, nil + + default: + // Every other method is logged by the bridge as "unhandled: " + // and, if it was a request, answered with a JSON-RPC error so codex + // never hangs. + return nil, agentproxy.ErrMethodNotFound + } +} + +// notify pushes one notification and logs (rather than propagating) a +// failure — from a background goroutine there's no request to answer with an +// error, and if the connection is dead every subsequent Notify in this turn +// will fail the same way, so the caller uses the bool to stop early instead +// of error-spamming a dead connection for the rest of the stream. +func (f *Facade) notify(method string, params any) bool { + if err := f.notifier.Notify(method, params); err != nil { + log.Printf("codex facade: notify %s failed: %v", method, err) + return false + } + return true +} + +// ensureEventLoop makes sure exactly one StreamSession read loop is running +// for this facade's connection before returning, opening one synchronously +// if none is active yet — either because this is the first turn for this +// facade's connection, or because the previous loop died and stopEventLoop +// cleared streamStarted. Callers (turn/start) must call this, and see it +// succeed, BEFORE calling SendInput: SendInput is what makes the harness +// start actually running and emitting events, so opening the stream only +// afterward (e.g. lazily from inside a newly spawned goroutine) leaves a +// window where a fast run can start, emit every event, and complete before +// any reader is attached to see them — codex would then hang on a turn +// that silently never got its notifications. See the Facade.mu doc comment +// for why there must never be more than one StreamSession call open at a +// time. +func (f *Facade) ensureEventLoop(ctx context.Context) error { + f.mu.Lock() + if f.streamStarted { + f.mu.Unlock() + return nil + } + f.mu.Unlock() + + stream, err := f.Sessions.StreamSession(ctx, f.SessionID, nil) + if err != nil { + return err + } + + f.mu.Lock() + f.streamStarted = true + f.mu.Unlock() + + go f.runEventLoop(stream) + return nil +} + +// trackTurn registers a new in-flight turn, keyed by run id, once +// ensureEventLoop has confirmed the read loop is running to see its events. +func (f *Facade) trackTurn(runID string) { + f.mu.Lock() + if f.turns == nil { + f.turns = make(map[string]*turnState) + } + f.turns[runID] = &turnState{itemID: runID + "-msg"} + f.mu.Unlock() +} + +// eventClaimPoll/eventClaimWait bound how long lookupTurn retries an +// unrecognized run id before giving up on it. +const ( + eventClaimPoll = 2 * time.Millisecond + eventClaimWait = 200 * time.Millisecond +) + +// lookupTurn returns the tracked turnState for runID, retrying briefly if +// it's not there yet. turn/start calls SendInput, which is what makes the +// harness create the run and start emitting its events, and only registers +// runID in f.turns once SendInput returns — so this goroutine can observe +// that run's first event fractionally before trackTurn's map write becomes +// visible to it. A caller that still gets ok == false after this returns +// should treat the event as genuinely untracked (e.g. one that predates this +// connection), not as this race. +func (f *Facade) lookupTurn(runID string) (*turnState, bool) { + deadline := time.Now().Add(eventClaimWait) + for { + f.mu.Lock() + ts, ok := f.turns[runID] + f.mu.Unlock() + if ok || time.Now().After(deadline) { + return ts, ok + } + time.Sleep(eventClaimPoll) + } +} + +// runEventLoop is the single, long-lived reader of stream — opened by +// ensureEventLoop before this goroutine was even spawned — for this facade's +// connection lifetime. It dispatches each event to whichever tracked turn it +// belongs to (by run id) and translates into the "one-way text" notification +// sequence: turn/started -> item/started -> item/agentMessage/delta* -> +// item/completed -> turn/completed. Everything else falls through untouched. +// Runs until the ctx passed to ensureEventLoop's StreamSession call is +// canceled (the connection closes) or the stream ends/errors — NOT until one +// turn finishes, since later turns depend on this same loop still running. +// The deferred stopEventLoop is what lets a later turn/start recover: +// without it, streamStarted would stay true forever after any exit here and +// every subsequent turn/start would return "inProgress" with no loop left +// running to ever notify it. +// +// TODO(M3): translate tool-call events instead of dropping them. +// TODO(M4): translate HITL events instead of dropping them. +func (f *Facade) runEventLoop(stream *godo.HostedAgentSessionStream) { + defer f.stopEventLoop() + defer stream.Close() + + for stream.Next() { + ev := stream.Current() + + ts, ok := f.lookupTurn(ev.RunID) + if !ok { + continue // event for a run this facade isn't tracking (e.g. predates this connection) + } + + switch ev.Kind { + case godo.HostedAgentEventKindRunStarted: + ts.startedAt = time.Now().Unix() + if !f.notify("turn/started", turnStartedNotification{ + ThreadID: f.SessionID, + Turn: turnObj{ID: ev.RunID, Items: []any{}, Status: "inProgress", StartedAt: &ts.startedAt}, + }) { + return + } + ts.itemStarted = f.notify("item/started", itemStartedNotification{ + Item: agentMessageItem{Type: "agentMessage", ID: ts.itemID, Text: ""}, + ThreadID: f.SessionID, + TurnID: ev.RunID, + StartedAtMs: ts.startedAt * 1000, + }) + + case godo.HostedAgentEventKindTokenChunk: + var payload struct { + Text string `json:"text"` + } + if err := json.Unmarshal(ev.Payload, &payload); err != nil { + continue + } + ts.text.WriteString(payload.Text) + if !f.notify("item/agentMessage/delta", agentMessageDeltaNotification{ + ThreadID: f.SessionID, + TurnID: ev.RunID, + ItemID: ts.itemID, + Delta: payload.Text, + }) { + return + } + + case godo.HostedAgentEventKindRunCompleted: + f.finishTurn(ev.RunID, ts, "completed", nil) + + case godo.HostedAgentEventKindRunFailed: + var payload struct { + Message string `json:"message"` + } + _ = json.Unmarshal(ev.Payload, &payload) + f.finishTurn(ev.RunID, ts, "failed", &turnError{Message: payload.Message}) + } + } + if err := stream.Err(); err != nil { + log.Printf("codex facade: stream error: %v", err) + } +} + +// stopEventLoop runs when runEventLoop returns for any reason — the stream +// ended, errored, or a notify failed (ensureEventLoop only spawns +// runEventLoop after StreamSession has already succeeded, so a failure to +// connect never reaches here; it fails turn/start directly instead). It +// clears streamStarted so the next turn/start reopens the stream instead of +// silently returning "inProgress" for a turn nothing will ever finish, and +// fails out any turns still tracked at that point (mid-flight when the loop +// died) so codex isn't left waiting forever on those either. +func (f *Facade) stopEventLoop() { + f.mu.Lock() + f.streamStarted = false + stranded := f.turns + f.turns = nil + f.mu.Unlock() + + for runID, ts := range stranded { + f.finishTurn(runID, ts, "failed", &turnError{Message: "lost connection to hosted session"}) + } +} + +// finishTurn sends the closing item/completed (if item/started ever landed) +// and turn/completed pair, then stops tracking runID. status is "completed" +// or "failed" — there is no turn/failed method in the codex protocol +// (confirmed against the source); failure is turn/completed with status +// "failed" and turnErr populated. +func (f *Facade) finishTurn(runID string, ts *turnState, status string, turnErr *turnError) { + completedAt := time.Now().Unix() + if ts.itemStarted { + f.notify("item/completed", itemCompletedNotification{ + Item: agentMessageItem{Type: "agentMessage", ID: ts.itemID, Text: ts.text.String()}, + ThreadID: f.SessionID, + TurnID: runID, + CompletedAtMs: completedAt * 1000, + }) + } + var durationMs *int64 + if ts.startedAt > 0 { + d := (completedAt - ts.startedAt) * 1000 + durationMs = &d + } + f.notify("turn/completed", turnCompletedNotification{ + ThreadID: f.SessionID, + Turn: turnObj{ + ID: runID, + Items: []any{}, + Status: status, + Error: turnErr, + StartedAt: &ts.startedAt, + CompletedAt: &completedAt, + DurationMs: durationMs, + }, + }) + + f.mu.Lock() + delete(f.turns, runID) + f.mu.Unlock() +} diff --git a/internal/agentproxy/codex/facade_test.go b/internal/agentproxy/codex/facade_test.go new file mode 100644 index 000000000..841957f04 --- /dev/null +++ b/internal/agentproxy/codex/facade_test.go @@ -0,0 +1,375 @@ +package codex + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/digitalocean/doctl/do" + "github.com/digitalocean/doctl/internal/agentproxy" + "github.com/digitalocean/doctl/internal/agentproxy/agentproxytest" + "github.com/digitalocean/godo" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const testSessionID = "sess_test123" + +// recordedNotification is one call a Facade made through Notifier.Notify, +// captured verbatim (not round-tripped through JSON) so assertions can type +// assert the concrete notification struct directly. +type recordedNotification struct { + method string + params any +} + +// notifierRecorder is a fake agentproxy.Notifier that queues every call so a +// test can assert on the exact notification sequence a facade produced, +// in order, without racing its background event-loop goroutine. +type notifierRecorder struct { + ch chan recordedNotification +} + +func newNotifierRecorder() *notifierRecorder { + return ¬ifierRecorder{ch: make(chan recordedNotification, 32)} +} + +func (r *notifierRecorder) Notify(method string, params any) error { + r.ch <- recordedNotification{method: method, params: params} + return nil +} + +// next blocks for the next notification, failing the test rather than +// hanging the suite if the facade's event loop never produces one. +func (r *notifierRecorder) next(t *testing.T) recordedNotification { + t.Helper() + select { + case n := <-r.ch: + return n + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for a notification") + return recordedNotification{} + } +} + +func (r *notifierRecorder) expectNone(t *testing.T) { + t.Helper() + select { + case n := <-r.ch: + t.Fatalf("expected no further notifications, got %q", n.method) + case <-time.After(100 * time.Millisecond): + } +} + +var _ agentproxy.Notifier = (*notifierRecorder)(nil) + +// newTestFacade wires a Facade against a fresh agentproxytest.Harness (a real +// do.HostedAgentsService talking to an httptest.Server) and installs a +// notifierRecorder, mirroring how agentproxy.Serve wires a real *wsNotifier +// via NotifierAware before handing the facade any messages. +func newTestFacade(t *testing.T) (*Facade, *agentproxytest.Harness, *notifierRecorder) { + t.Helper() + h := agentproxytest.New(t, testSessionID) + client, err := godo.New(nil, godo.SetBaseURL(h.Server.URL+"/")) + require.NoError(t, err) + svc := do.NewHostedAgentsService(client) + + f := &Facade{SessionID: testSessionID, Sessions: svc} + rec := newNotifierRecorder() + f.SetNotifier(rec) + return f, h, rec +} + +// dispatch marshals params (if any) to json.RawMessage and calls Dispatch, +// matching what proxy.go's handleConn does with an incoming JSON-RPC frame's +// params field. +func dispatch(t *testing.T, f *Facade, method string, params any) (any, error) { + t.Helper() + var raw json.RawMessage + if params != nil { + b, err := json.Marshal(params) + require.NoError(t, err) + raw = b + } + return f.Dispatch(context.Background(), method, raw) +} + +func TestFacade_Initialize(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "initialize", nil) + require.NoError(t, err) + + init, ok := result.(initializeResult) + require.True(t, ok, "result should be initializeResult, got %T", result) + assert.Equal(t, "unix", init.PlatformFamily) + assert.Equal(t, "linux", init.PlatformOs) + assert.Contains(t, init.UserAgent, TestedVersion) +} + +func TestFacade_Initialized_NoOp(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "initialized", nil) + assert.NoError(t, err) + assert.Nil(t, result) +} + +func TestFacade_AccountRead(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "account/read", nil) + require.NoError(t, err) + + acct, ok := result.(accountReadResult) + require.True(t, ok, "result should be accountReadResult, got %T", result) + assert.False(t, acct.RequiresOpenaiAuth) + assert.Equal(t, "chatgpt", acct.Account.Type) +} + +func TestFacade_ModelList(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "model/list", nil) + require.NoError(t, err) + + list, ok := result.(modelListResult) + require.True(t, ok, "result should be modelListResult, got %T", result) + require.Len(t, list.Data, 1) + assert.Equal(t, defaultModel, list.Data[0].ID) + assert.True(t, list.Data[0].IsDefault) +} + +func TestFacade_ThreadStart(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "thread/start", nil) + require.NoError(t, err) + + start, ok := result.(threadStartResult) + require.True(t, ok, "result should be threadStartResult, got %T", result) + assert.Equal(t, testSessionID, start.Thread.ID) + assert.Equal(t, testSessionID, start.Thread.SessionID) + assert.Equal(t, "idle", start.Thread.Status.Type) +} + +func TestFacade_ThreadResume(t *testing.T) { + f, _, _ := newTestFacade(t) + + t.Run("matching thread id", func(t *testing.T) { + result, err := dispatch(t, f, "thread/resume", threadResumeParams{ThreadID: testSessionID}) + require.NoError(t, err) + + resume, ok := result.(threadResumeResult) + require.True(t, ok, "result should be threadResumeResult, got %T", result) + assert.Equal(t, testSessionID, resume.Thread.ID) + }) + + t.Run("unknown thread id is a real not-found, not a stub gap", func(t *testing.T) { + _, err := dispatch(t, f, "thread/resume", threadResumeParams{ThreadID: "not-this-session"}) + require.Error(t, err) + + var rpcErr *agentproxy.RPCError + require.ErrorAs(t, err, &rpcErr) + assert.Equal(t, -32001, rpcErr.Code) + }) +} + +func TestFacade_UnhandledMethod(t *testing.T) { + f, _, _ := newTestFacade(t) + + _, err := dispatch(t, f, "totally/unknown/method", nil) + assert.ErrorIs(t, err, agentproxy.ErrMethodNotFound) +} + +func TestFacade_TurnInterrupt_NoOp(t *testing.T) { + f, _, _ := newTestFacade(t) + + result, err := dispatch(t, f, "turn/interrupt", nil) + require.NoError(t, err) + assert.Equal(t, turnInterruptResult{}, result) +} + +func TestFacade_TurnStart_WrongThreadID(t *testing.T) { + f, _, _ := newTestFacade(t) + + _, err := dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: "not-this-session", + Input: []userInputItem{{Type: "text", Text: "hi"}}, + }) + require.Error(t, err) + + var rpcErr *agentproxy.RPCError + require.ErrorAs(t, err, &rpcErr) + assert.Equal(t, -32001, rpcErr.Code) +} + +// TestFacade_TurnStart_SendInputError exercises the harness's "no run queued" +// failure path (no QueueRun call before turn/start) to confirm a SendInput +// failure surfaces as a JSON-RPC error rather than a hang or panic. +func TestFacade_TurnStart_SendInputError(t *testing.T) { + f, _, _ := newTestFacade(t) + + _, err := dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: testSessionID, + Input: []userInputItem{{Type: "text", Text: "hi"}}, + }) + require.Error(t, err) + + var rpcErr *agentproxy.RPCError + require.ErrorAs(t, err, &rpcErr) + assert.Equal(t, -32000, rpcErr.Code) +} + +// TestFacade_TurnStart_StreamsToCompletion drives one full "one-way text" +// turn end to end: turn/start -> SendInput against the fake harness -> the +// harness's queued run.started/run.token_delta/run.completed events -> +// translated into the turn/started, item/started, item/agentMessage/delta*, +// item/completed, turn/completed notification sequence the design doc's +// event-mapping table calls for. +func TestFacade_TurnStart_StreamsToCompletion(t *testing.T) { + f, h, rec := newTestFacade(t) + + h.QueueRun("run-1", + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunStarted)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindTokenChunk), Data: json.RawMessage(`{"text":"Hel"}`)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindTokenChunk), Data: json.RawMessage(`{"text":"lo"}`)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunCompleted)}, + ) + + result, err := dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: testSessionID, + Input: []userInputItem{{Type: "text", Text: "hi"}}, + }) + require.NoError(t, err) + + start, ok := result.(turnStartResult) + require.True(t, ok, "result should be turnStartResult, got %T", result) + assert.Equal(t, "run-1", start.Turn.ID) + assert.Equal(t, "inProgress", start.Turn.Status) + + started := rec.next(t) + assert.Equal(t, "turn/started", started.method) + assert.Equal(t, "run-1", started.params.(turnStartedNotification).Turn.ID) + + itemStarted := rec.next(t) + assert.Equal(t, "item/started", itemStarted.method) + + delta1 := rec.next(t) + require.Equal(t, "item/agentMessage/delta", delta1.method) + assert.Equal(t, "Hel", delta1.params.(agentMessageDeltaNotification).Delta) + + delta2 := rec.next(t) + require.Equal(t, "item/agentMessage/delta", delta2.method) + assert.Equal(t, "lo", delta2.params.(agentMessageDeltaNotification).Delta) + + itemCompleted := rec.next(t) + require.Equal(t, "item/completed", itemCompleted.method) + assert.Equal(t, "Hello", itemCompleted.params.(itemCompletedNotification).Item.Text) + + turnCompleted := rec.next(t) + require.Equal(t, "turn/completed", turnCompleted.method) + tc := turnCompleted.params.(turnCompletedNotification) + assert.Equal(t, "completed", tc.Turn.Status) + assert.Equal(t, "run-1", tc.Turn.ID) + assert.Nil(t, tc.Turn.Error) + + rec.expectNone(t) +} + +// waitEventLoopStopped polls until the facade's event-loop goroutine has +// exited and stopEventLoop has reset streamStarted. Observing a +// turn/completed notification only means the loop finished that turn, not +// that it has since hit EOF on the underlying stream and stopped — a test +// that starts a second turn right after turn/completed would otherwise race +// stopEventLoop rather than deterministically exercising life after it. +func waitEventLoopStopped(t *testing.T, f *Facade) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + f.mu.Lock() + stopped := !f.streamStarted + f.mu.Unlock() + if stopped { + return + } + time.Sleep(time.Millisecond) + } + t.Fatal("timed out waiting for the event loop to stop") +} + +// TestFacade_TurnStart_ReopensStreamAfterPreviousLoopEnded confirms a second +// turn/start can still get notified after the harness's SSE stream for the +// first turn ends (agentproxytest.Harness.handleStream closes the response +// once its queued events run out) — regression coverage for streamStarted +// getting stuck true forever, which left every later turn/start returning +// "inProgress" with no event-loop goroutine left running to ever notify it. +func TestFacade_TurnStart_ReopensStreamAfterPreviousLoopEnded(t *testing.T) { + f, h, rec := newTestFacade(t) + + h.QueueRun("run-1", + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunStarted)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunCompleted)}, + ) + _, err := dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: testSessionID, + Input: []userInputItem{{Type: "text", Text: "hi"}}, + }) + require.NoError(t, err) + + _ = rec.next(t) // turn/started + _ = rec.next(t) // item/started + _ = rec.next(t) // item/completed + turnCompleted := rec.next(t) + require.Equal(t, "turn/completed", turnCompleted.method) + waitEventLoopStopped(t, f) + + h.QueueRun("run-2", + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunStarted)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunCompleted)}, + ) + _, err = dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: testSessionID, + Input: []userInputItem{{Type: "text", Text: "again"}}, + }) + require.NoError(t, err) + + started := rec.next(t) + require.Equal(t, "turn/started", started.method) + assert.Equal(t, "run-2", started.params.(turnStartedNotification).Turn.ID) +} + +// TestFacade_TurnStart_Failure confirms a canonical run.failed event maps to +// turn/completed with status "failed" and the error populated — there is no +// turn/failed method in the codex protocol (see the facade.go comment above +// turnStartParams), so this is the only failure signal the client ever sees. +func TestFacade_TurnStart_Failure(t *testing.T) { + f, h, rec := newTestFacade(t) + + h.QueueRun("run-fail", + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunStarted)}, + agentproxytest.Event{Type: string(godo.HostedAgentEventKindRunFailed), Data: json.RawMessage(`{"message":"boom"}`)}, + ) + + _, err := dispatch(t, f, "turn/start", turnStartParams{ + ThreadID: testSessionID, + Input: []userInputItem{{Type: "text", Text: "hi"}}, + }) + require.NoError(t, err) + + _ = rec.next(t) // turn/started + _ = rec.next(t) // item/started + + // finishTurn always closes a started item before completing the turn, + // success or failure — an item that started should never dangle "open". + itemCompleted := rec.next(t) + require.Equal(t, "item/completed", itemCompleted.method) + + turnCompleted := rec.next(t) + require.Equal(t, "turn/completed", turnCompleted.method) + tc := turnCompleted.params.(turnCompletedNotification) + assert.Equal(t, "failed", tc.Turn.Status) + require.NotNil(t, tc.Turn.Error) + assert.Equal(t, "boom", tc.Turn.Error.Message) +} diff --git a/internal/agentproxy/codex/wire_test.go b/internal/agentproxy/codex/wire_test.go new file mode 100644 index 000000000..6166aee29 --- /dev/null +++ b/internal/agentproxy/codex/wire_test.go @@ -0,0 +1,281 @@ +package codex + +import ( + "context" + "encoding/json" + "net" + "net/http" + "testing" + "time" + + "github.com/digitalocean/doctl/do" + "github.com/digitalocean/doctl/internal/agentproxy" + "github.com/digitalocean/doctl/internal/agentproxy/agentproxytest" + "github.com/digitalocean/godo" + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" +) + +// This file is the wire-level complement to facade_test.go: those tests call +// Facade.Dispatch directly with a fake Notifier, which is fast and thorough +// for facade logic but never exercises proxy.go's actual JSON-RPC/WebSocket +// framing (ServeListener, handleConn's request/notification split, the +// SetNotifier/NotifierAware wiring, or whether a struct's json tags really +// serialize the way the Go-level tests assume). wsTestClient below drives a +// real agentproxy.ServeListener over a real *websocket.Conn, so a struct-tag +// typo or a framing regression in proxy.go would fail here even if every +// Dispatch-level test still passes. + +// wsTestClient is a minimal scripted WebSocket JSON-RPC client: send one +// message, read the next one back, with a short deadline so a regression +// that breaks the reply/notification sequence fails fast and loud instead of +// hanging the test suite. +type wsTestClient struct { + t *testing.T + conn *websocket.Conn +} + +func dialTestClient(t *testing.T, url string) *wsTestClient { + t.Helper() + conn, _, err := websocket.DefaultDialer.Dial(url, nil) + require.NoError(t, err, "dial %s", url) + t.Cleanup(func() { conn.Close() }) + return &wsTestClient{t: t, conn: conn} +} + +// dialTestClientRetry dials until it succeeds or 2s pass, for reconnect tests +// where a just-closed previous connection's single "client slot" (see +// ServeListener) may not be released by the time this dial races it. +func dialTestClientRetry(t *testing.T, url string) *wsTestClient { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + conn, _, err := websocket.DefaultDialer.Dial(url, nil) + if err == nil { + t.Cleanup(func() { conn.Close() }) + return &wsTestClient{t: t, conn: conn} + } + if time.Now().After(deadline) { + t.Fatalf("dial %s: %v", url, err) + } + time.Sleep(10 * time.Millisecond) + } +} + +func (c *wsTestClient) send(id, method string, params any) { + c.t.Helper() + msg := map[string]any{"method": method} + if id != "" { + msg["id"] = id + } + if params != nil { + msg["params"] = params + } + require.NoError(c.t, c.conn.WriteJSON(msg), "send %s", method) +} + +// recv reads the next frame within 5s, or fails the test — never polls or +// sleeps to wait for a notification. +func (c *wsTestClient) recv() map[string]any { + c.t.Helper() + require.NoError(c.t, c.conn.SetReadDeadline(time.Now().Add(5*time.Second))) + var msg map[string]any + require.NoError(c.t, c.conn.ReadJSON(&msg), "recv") + return msg +} + +// result returns msg["result"] as a map, failing the test if msg carries a +// JSON-RPC error instead or result isn't an object. +func (c *wsTestClient) result(msg map[string]any) map[string]any { + c.t.Helper() + if errObj, ok := msg["error"]; ok { + c.t.Fatalf("got error reply instead of result: %v", errObj) + } + result, ok := msg["result"].(map[string]any) + require.True(c.t, ok, "result is not an object: %v", msg) + return result +} + +// TestWire_InitializeThreadStartTurnStartDeltasTurnCompleted is the M2 "done +// when" sequence test at the wire level: initialize -> thread/start -> +// turn/start -> deltas -> turn/completed, over a real WebSocket connection to +// a real agentproxy.ServeListener, against a fake harness instead of a live +// one. facade_test.go's TestFacade_TurnStart_StreamsToCompletion covers the +// same sequence at the Dispatch level; this one additionally proves the JSON +// wire shapes and proxy.go's framing are correct. Extend this file (or +// facade_test.go) for M3 (tool-call items) and M4 (approval round-trip) +// rather than retrofitting tests later. +func TestWire_InitializeThreadStartTurnStartDeltasTurnCompleted(t *testing.T) { + const sessionID = "sess-wire-test" + const runID = "run-wire-1" + + harness := agentproxytest.New(t, sessionID) + harness.QueueRun(runID, + agentproxytest.Event{Type: "run.started", Data: json.RawMessage(`{"agent":"codex"}`)}, + agentproxytest.Event{Type: "run.token_delta", Data: json.RawMessage(`{"text":"Four"}`)}, + agentproxytest.Event{Type: "run.token_delta", Data: json.RawMessage(`{"text":"!"}`)}, + agentproxytest.Event{Type: "run.completed", Data: json.RawMessage(`{}`)}, + ) + + godoClient, err := godo.New(http.DefaultClient, godo.SetBaseURL(harness.Server.URL+"/")) + require.NoError(t, err) + svc := do.NewHostedAgentsService(godoClient) + newFacade := func() agentproxy.Facade { + return &Facade{SessionID: sessionID, Sessions: svc} + } + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := ln.Addr().String() + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + serveErr := make(chan error, 1) + go func() { serveErr <- agentproxy.ServeListener(ctx, ln, newFacade) }() + + client := dialTestClient(t, "ws://"+addr+"/") + + // initialize + client.send("1", "initialize", map[string]any{ + "clientInfo": map[string]any{"name": "test", "version": "0"}, + "capabilities": map[string]any{}, + }) + initResult := client.result(client.recv()) + require.NotEmpty(t, initResult["userAgent"], "initialize result missing userAgent: %v", initResult) + + // thread/start + client.send("2", "thread/start", map[string]any{}) + threadResult := client.result(client.recv()) + thread, ok := threadResult["thread"].(map[string]any) + require.True(t, ok, "thread/start result missing thread object: %v", threadResult) + require.Equal(t, sessionID, thread["id"]) + + // turn/start + client.send("3", "turn/start", map[string]any{ + "threadId": sessionID, + "input": []any{map[string]any{"type": "text", "text": "test prompt"}}, + }) + turnStartResult := client.result(client.recv()) + turn, ok := turnStartResult["turn"].(map[string]any) + require.True(t, ok, "turn/start result missing turn object: %v", turnStartResult) + require.Equal(t, runID, turn["id"]) + require.Equal(t, "inProgress", turn["status"]) + + // Async notification sequence, in order, with the fields that matter. + notif := client.recv() + require.Equal(t, "turn/started", notif["method"], "full: %v", notif) + + notif = client.recv() + require.Equal(t, "item/started", notif["method"], "full: %v", notif) + params, _ := notif["params"].(map[string]any) + item, _ := params["item"].(map[string]any) + itemID, _ := item["id"].(string) + require.NotEmpty(t, itemID, "item/started missing item.id: %v", notif) + + for _, want := range []string{"Four", "!"} { + notif = client.recv() + require.Equal(t, "item/agentMessage/delta", notif["method"], "full: %v", notif) + params, _ = notif["params"].(map[string]any) + require.Equal(t, want, params["delta"]) + require.Equal(t, itemID, params["itemId"]) + } + + notif = client.recv() + require.Equal(t, "item/completed", notif["method"], "full: %v", notif) + params, _ = notif["params"].(map[string]any) + item, _ = params["item"].(map[string]any) + require.Equal(t, "Four!", item["text"]) + + notif = client.recv() + require.Equal(t, "turn/completed", notif["method"], "full: %v", notif) + params, _ = notif["params"].(map[string]any) + finalTurn, _ := params["turn"].(map[string]any) + require.Equal(t, "completed", finalTurn["status"]) + + cancel() + select { + case err := <-serveErr: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("ServeListener did not shut down within 5s of ctx cancellation") + } +} + +// TestWire_ReconnectGetsFreshFacadeState proves a second connection over the +// same ServeListener, after the first disconnects, gets a brand-new Facade +// rather than reusing the first connection's — regression coverage for a +// review comment noting that a single shared Facade instance would carry its +// first connection's turns map and streamStarted over into the reconnect: +// codex --remote reconnecting would then hang on turn/start forever, since +// the old connection's event loop is dead but streamStarted would still read +// true, so ensureEventLoop would never open a new one. +func TestWire_ReconnectGetsFreshFacadeState(t *testing.T) { + const sessionID = "sess-wire-reconnect" + + harness := agentproxytest.New(t, sessionID) + godoClient, err := godo.New(http.DefaultClient, godo.SetBaseURL(harness.Server.URL+"/")) + require.NoError(t, err) + svc := do.NewHostedAgentsService(godoClient) + newFacade := func() agentproxy.Facade { + return &Facade{SessionID: sessionID, Sessions: svc} + } + + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := ln.Addr().String() + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + serveErr := make(chan error, 1) + go func() { serveErr <- agentproxy.ServeListener(ctx, ln, newFacade) }() + + // First connection: run one full turn to completion, then disconnect + // without a clean handshake, mirroring a codex --remote that just drops. + harness.QueueRun("run-1", + agentproxytest.Event{Type: "run.started"}, + agentproxytest.Event{Type: "run.completed"}, + ) + client1 := dialTestClient(t, "ws://"+addr+"/") + client1.send("1", "turn/start", map[string]any{ + "threadId": sessionID, + "input": []any{map[string]any{"type": "text", "text": "hi"}}, + }) + _ = client1.result(client1.recv()) // turn/start reply + for { + notif := client1.recv() + if notif["method"] == "turn/completed" { + break + } + } + client1.conn.Close() + + // Second connection, over the same listener: a fresh turn/start must + // still work end to end. If the old Facade were reused, streamStarted + // would still read true from the dead first connection's event loop, and + // this turn/start would return "inProgress" with no turn/started ever + // following it. + harness.QueueRun("run-2", + agentproxytest.Event{Type: "run.started"}, + agentproxytest.Event{Type: "run.completed"}, + ) + client2 := dialTestClientRetry(t, "ws://"+addr+"/") + client2.send("1", "turn/start", map[string]any{ + "threadId": sessionID, + "input": []any{map[string]any{"type": "text", "text": "hi again"}}, + }) + turnStartResult := client2.result(client2.recv()) + turn, ok := turnStartResult["turn"].(map[string]any) + require.True(t, ok, "turn/start result missing turn object: %v", turnStartResult) + require.Equal(t, "run-2", turn["id"]) + + notif := client2.recv() + require.Equal(t, "turn/started", notif["method"], "full: %v", notif) + + cancel() + select { + case err := <-serveErr: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("ServeListener did not shut down within 5s of ctx cancellation") + } +} diff --git a/internal/agentproxy/facade.go b/internal/agentproxy/facade.go new file mode 100644 index 000000000..05247e413 --- /dev/null +++ b/internal/agentproxy/facade.go @@ -0,0 +1,53 @@ +// Package agentproxy implements the local WebSocket facade that lets an +// unmodified coding-agent CLI (codex today, others later) drive a hosted MARS +// session as if the session were its own local backend. +package agentproxy + +import ( + "context" + "encoding/json" +) + +// RPCError is a JSON-RPC 2.0 error object. +type RPCError struct { + Code int `json:"code"` + Message string `json:"message"` +} + +func (e *RPCError) Error() string { return e.Message } + +// ErrMethodNotFound is returned by a Facade for any method it doesn't (yet) +// implement. The bridge logs it as "unhandled: " and, for requests +// (a message with an id), turns it into a JSON-RPC error response — so the +// client never hangs waiting on a reply that was logged-and-dropped instead +// of sent. +var ErrMethodNotFound = &RPCError{Code: -32601, Message: "method not found"} + +// Facade is the per-agent seam start-proxy dispatches JSON-RPC messages to. +// One implementation per --type (v1: codex; future: claude-code, opencode). +// +// Dispatch handles a single JSON-RPC message, request or notification alike; +// it doesn't need to know which one it's looking at. The bridge is the one +// that knows: a message with a non-nil id is a request and its (result, err) +// becomes a reply; a message with a nil id is a notification and (result, +// err) only affects logging, never the wire. +type Facade interface { + Dispatch(ctx context.Context, method string, params json.RawMessage) (result any, err error) +} + +// Notifier lets a Facade push a server-initiated JSON-RPC notification to the +// connected client at any time — not just synchronously in reply to a +// request. Needed for events that arrive asynchronously from the harness +// (streamed tokens, turn completion) on their own timeline, well after the +// request that kicked off the turn has already been answered. +type Notifier interface { + Notify(method string, params any) error +} + +// NotifierAware is implemented by facades that need to push asynchronous +// notifications rather than only answer synchronous requests. The bridge +// calls SetNotifier once per connection, before handing it any messages, so +// Dispatch can stash it and use it later from a background goroutine. +type NotifierAware interface { + SetNotifier(Notifier) +} diff --git a/internal/agentproxy/proxy.go b/internal/agentproxy/proxy.go new file mode 100644 index 000000000..420056d53 --- /dev/null +++ b/internal/agentproxy/proxy.go @@ -0,0 +1,220 @@ +package agentproxy + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log" + "net" + "net/http" + "sync" + "time" + + "github.com/gorilla/websocket" +) + +// rpcMessage is a JSON-RPC 2.0 envelope, minus the "jsonrpc" field itself — +// the real codex app-server protocol omits "jsonrpc":"2.0" on the wire +// (confirmed by the protocol capture), and codex --remote doesn't require it +// either, so the proxy matches what's actually sent instead of adding a +// field nothing checks. +type rpcMessage struct { + ID json.RawMessage `json:"id,omitempty"` + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` + Result any `json:"result,omitempty"` + Error *RPCError `json:"error,omitempty"` +} + +var upgrader = websocket.Upgrader{ + // Loopback-only listener (see Serve); no browser reaches this, so origin + // checking buys nothing and would only get in websocat's way. + CheckOrigin: func(r *http.Request) bool { return true }, +} + +// wsNotifier serializes every write to a WebSocket connection — replies from +// the main read/dispatch loop and asynchronous notifications from a facade's +// background goroutines alike — behind one mutex. gorilla/websocket conns +// aren't safe for concurrent writers; once a facade can push notifications on +// its own timeline (Notifier), there are always at least two potential +// writers, so this is required, not defensive. +type wsNotifier struct { + mu sync.Mutex + conn *websocket.Conn +} + +func (w *wsNotifier) writeMessage(msg rpcMessage) error { + out, err := json.Marshal(msg) + if err != nil { + return fmt.Errorf("marshal %s: %w", msg.Method, err) + } + w.mu.Lock() + defer w.mu.Unlock() + return w.conn.WriteMessage(websocket.TextMessage, out) +} + +// Notify implements agentproxy.Notifier. +func (w *wsNotifier) Notify(method string, params any) error { + return w.writeMessage(rpcMessage{Method: method, Params: mustRawMessage(params)}) +} + +func mustRawMessage(v any) json.RawMessage { + if v == nil { + return nil + } + data, err := json.Marshal(v) + if err != nil { + // A facade handed us a value that can't marshal — a programmer error + // in that facade, not a runtime condition to recover from gracefully. + panic(fmt.Sprintf("agentproxy: params do not marshal to JSON: %v", err)) + } + return data +} + +// Serve binds a WebSocket listener on 127.0.0.1:port — never "localhost" +// (IPv6 ::1 resolution can fail to connect) or 0.0.0.0 (codex requires auth +// for non-loopback listeners, and there's no reason to expose this beyond the +// machine anyway) — and speaks newFacade's JSON-RPC protocol to exactly one +// connected client at a time, until ctx is canceled. +func Serve(ctx context.Context, port int, newFacade func() Facade) error { + ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) + if err != nil { + return err + } + return ServeListener(ctx, ln, newFacade) +} + +// ServeListener is Serve, except it speaks newFacade's protocol over an +// already-bound listener instead of binding one from a port number itself. +// +// This exists for tests: picking a free ephemeral port ahead of time (e.g. +// via net.Listen("tcp", "127.0.0.1:0")) and having a *separate* call to Serve +// re-bind that same port number is inherently racy — anything else on the +// machine could grab it in between. Handing the already-bound *net.Listener +// straight to ServeListener closes that race window entirely: the listener +// is accepting connections (into the kernel backlog, at least) from the +// moment net.Listen returns, before this function is even called. +// +// newFacade is called once per accepted connection, not once for the whole +// listener: a facade like codex's carries per-connection state (in-flight +// turns, whether its event loop is running), and only one client connects at +// a time anyway (see the slot below), so reusing one Facade instance across +// a disconnect/reconnect would leak that state into the new connection — +// notifications from a still-unwinding previous connection's background +// goroutine could even land on the new socket via a shared notifier. A fresh +// Facade per connection starts clean and makes the previous one's goroutines +// (if still winding down) entirely self-contained. +func ServeListener(ctx context.Context, ln net.Listener, newFacade func() Facade) error { + // One slot: the first upgrade takes it, a second concurrent attempt is + // refused until the first disconnects and returns it. + slot := make(chan struct{}, 1) + slot <- struct{}{} + + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + select { + case <-slot: + default: + http.Error(w, "a client is already connected", http.StatusConflict) + return + } + defer func() { slot <- struct{}{} }() + + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Printf("agentproxy: upgrade error: %v", err) + return + } + defer conn.Close() + + // Own context for this connection's lifetime, not r.Context(): after + // Upgrade hijacks the connection, the request context's cancellation + // semantics on disconnect are no longer guaranteed. This is what a + // facade's background goroutines (e.g. streaming a turn's tokens) + // watch to know when to stop. + connCtx, cancel := context.WithCancel(context.Background()) + defer cancel() + + facade := newFacade() + notifier := &wsNotifier{conn: conn} + if na, ok := facade.(NotifierAware); ok { + na.SetNotifier(notifier) + } + + log.Printf("agentproxy: client connected from %s", r.RemoteAddr) + handleConn(connCtx, conn, facade, notifier) + log.Printf("agentproxy: client disconnected") + }) + + server := &http.Server{Handler: mux} + + errCh := make(chan error, 1) + go func() { errCh <- server.Serve(ln) }() + + select { + case err := <-errCh: + if err != nil && !errors.Is(err, http.ErrServerClosed) { + return err + } + return nil + case <-ctx.Done(): + shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + return server.Shutdown(shutdownCtx) + } +} + +// handleConn pumps one WebSocket connection: read a text frame, dispatch it +// to facade, write a response if (and only if) the message was a request. +// +// Every request gets *some* reply, even an error — that invariant lives here, +// structurally, rather than in each facade's Dispatch, so a facade can never +// recreate the "TUI hangs forever on a logged-and-dropped request" bug by +// forgetting to answer. +func handleConn(ctx context.Context, conn *websocket.Conn, facade Facade, notifier *wsNotifier) { + for { + _, data, err := conn.ReadMessage() + if err != nil { + if !websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + log.Printf("agentproxy: read error: %v", err) + } + return + } + + var msg rpcMessage + if err := json.Unmarshal(data, &msg); err != nil { + log.Printf("agentproxy: malformed JSON-RPC frame, dropping: %v", err) + continue + } + + result, dispatchErr := facade.Dispatch(ctx, msg.Method, msg.Params) + + if errors.Is(dispatchErr, ErrMethodNotFound) { + log.Printf("unhandled: %s", msg.Method) + } else if dispatchErr != nil { + log.Printf("agentproxy: %s: %v", msg.Method, dispatchErr) + } + + if len(msg.ID) == 0 { + // Notification: no id, no reply — even on error. + continue + } + + reply := rpcMessage{ID: msg.ID} + if dispatchErr != nil { + var rpcErr *RPCError + if !errors.As(dispatchErr, &rpcErr) { + rpcErr = &RPCError{Code: -32603, Message: dispatchErr.Error()} + } + reply.Error = rpcErr + } else { + reply.Result = result + } + + if err := notifier.writeMessage(reply); err != nil { + log.Printf("agentproxy: write error: %v", err) + return + } + } +}