diff --git a/pkg/api/types.go b/pkg/api/types.go index 2beb910b19..a4b2995474 100644 --- a/pkg/api/types.go +++ b/pkg/api/types.go @@ -21,6 +21,22 @@ type Agent struct { Multi bool `json:"multi"` } +// ErrCodeUnknownSession is the [ErrorResponse.Code] for a request that names +// a session the server does not know. Distinguishing it matters to clients +// like the board: a 404 with this code means "the server is current but the +// session does not exist (yet)" — keep waiting or recreate the session — +// while a 404 without it means the route itself is missing (an older binary +// that predates the endpoint). +const ErrCodeUnknownSession = "unknown_session" + +// ErrorResponse is the body of an API error that clients may need to react +// to programmatically. Message is human-readable; Code, when set, is a +// stable machine-readable classifier (e.g. [ErrCodeUnknownSession]). +type ErrorResponse struct { + Code string `json:"code,omitempty"` + Message string `json:"message"` +} + // CreateAgentRequest represents a request to create an agent type CreateAgentRequest struct { Prompt string `json:"prompt"` @@ -327,6 +343,10 @@ type SessionSnapshotResponse struct { // Zero when the session has no event stream (no events yet, or not // attached to a control plane). LastEventSeq uint64 `json:"last_event_seq"` + // Cost is the session's cumulative cost in US dollars, including + // sub-sessions and item-level costs (e.g. compaction). Clients should + // prefer it over summing per-message costs, which misses those. + Cost float64 `json:"cost"` } // RunAgentRequest is the body of POST /api/sessions/:id/agent/:agent[/:agent_name]. diff --git a/pkg/app/app.go b/pkg/app/app.go index c473dcb997..07f4f97a46 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -805,6 +805,12 @@ func (a *App) removeSubscriber(ch chan tea.Msg) { // scatters every message to all currently-registered subscribers. Sends are // non-blocking; if a subscriber's buffer is full the event is dropped for // that subscriber so one slow consumer cannot stall the others. +// +// Turn-boundary events are the exception: dropping a stream_started or +// stream_stopped skews a consumer's turn accounting for good (the SSE replay +// buffer never sees the event, so reconnecting cannot recover it). For those, +// the oldest pending message — almost always a content delta, which the next +// delta supersedes — is evicted to make room instead. func (a *App) startFanOut() { throttled := a.throttleEvents(a.ctx(), a.events) go func() { @@ -816,13 +822,47 @@ func (a *App) startFanOut() { select { case ch <- msg: default: - slog.Warn("app: subscriber buffer full, dropping event") + if !isTurnBoundaryEvent(msg) { + slog.Warn("app: subscriber buffer full, dropping event") + continue + } + // Evict the oldest pending message (racing the subscriber's + // own receive is fine: either way a slot frees up), then + // retry once. Still full means the subscriber is wedged; + // drop as before rather than block the fan-out. + select { + case <-ch: + default: + } + select { + case ch <- msg: + default: + slog.Warn("app: subscriber buffer full, dropping turn-boundary event") + } } } } }() } +// isTurnBoundaryEvent reports whether msg is one of the events consumers use +// to track turn state (running/waiting/failed/paused) and identity (title). +// These are low-frequency and irrecoverable when lost, unlike the content +// deltas that dominate the stream, so the fan-out prefers them on overflow. +func isTurnBoundaryEvent(msg tea.Msg) bool { + switch msg.(type) { + case *runtime.StreamStartedEvent, + *runtime.StreamStoppedEvent, + *runtime.UserMessageEvent, + *runtime.ErrorEvent, + *runtime.PausedEvent, + *runtime.SessionTitleEvent: + return true + default: + return false + } +} + // Resume resumes the runtime with the given confirmation request func (a *App) Resume(req runtime.ResumeRequest) { a.runtime.Resume(a.ctx(), req) diff --git a/pkg/app/fanout_test.go b/pkg/app/fanout_test.go new file mode 100644 index 0000000000..0100674bf4 --- /dev/null +++ b/pkg/app/fanout_test.go @@ -0,0 +1,86 @@ +package app + +import ( + "context" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/runtime" +) + +// TestFanOut_TurnBoundaryEventEvictsPendingDelta verifies that when a +// subscriber's buffer is full, a turn-boundary event (which SSE consumers +// cannot recover if lost) evicts the oldest pending message instead of being +// dropped. +func TestFanOut_TurnBoundaryEventEvictsPendingDelta(t *testing.T) { + t.Parallel() + + events := make(chan tea.Msg, 16) + app := &App{ + ctx: func() context.Context { return t.Context() }, + events: events, + throttleDuration: time.Millisecond, + } + + // A one-slot subscriber makes the overflow deterministic. The subscriber + // never reads, standing in for a consumer that fell behind. + ch := make(chan tea.Msg, 1) + app.addSubscriber(ch) + app.fanoutOnce.Do(app.startFanOut) + + // Fill the subscriber's buffer with a droppable event. + filler := runtime.NewTokenUsageEvent("sess", "root", &runtime.Usage{}) + events <- filler + require.Eventually(t, func() bool { return len(ch) == 1 }, 2*time.Second, time.Millisecond) + + // The turn-boundary event must displace the pending filler. + stopped := runtime.StreamStopped("sess", "root", "normal") + events <- stopped + + require.Eventually(t, func() bool { + select { + case msg := <-ch: + return assert.ObjectsAreEqual(stopped, msg) + default: + return false + } + }, 2*time.Second, time.Millisecond, "the stream_stopped event must survive the overflow") +} + +// TestFanOut_DroppableEventIsDroppedOnOverflow verifies the pre-existing +// behavior for non-boundary events: on overflow they are dropped, never +// evicting anything. +func TestFanOut_DroppableEventIsDroppedOnOverflow(t *testing.T) { + t.Parallel() + + events := make(chan tea.Msg, 16) + app := &App{ + ctx: func() context.Context { return t.Context() }, + events: events, + throttleDuration: time.Millisecond, + } + + ch := make(chan tea.Msg, 1) + app.addSubscriber(ch) + // The witness is registered after ch, so once a message reaches it the + // fan-out has already made its keep-or-drop decision for ch. + witness := make(chan tea.Msg, 16) + app.addSubscriber(witness) + app.fanoutOnce.Do(app.startFanOut) + + first := runtime.NewTokenUsageEvent("sess", "root", &runtime.Usage{}) + events <- first + require.Eventually(t, func() bool { return len(ch) == 1 }, 2*time.Second, time.Millisecond) + + // A second droppable event overflows and is dropped; the first stays. + second := runtime.NewTokenUsageEvent("sess", "other", &runtime.Usage{}) + events <- second + require.Eventually(t, func() bool { return len(witness) == 2 }, 2*time.Second, time.Millisecond) + + msg := <-ch + assert.Equal(t, first, msg) +} diff --git a/pkg/board/client.go b/pkg/board/client.go index af3ddf2c7f..f5aba4dda4 100644 --- a/pkg/board/client.go +++ b/pkg/board/client.go @@ -13,6 +13,8 @@ import ( "net/url" "strconv" "strings" + "sync/atomic" + "time" ) // Event types the board reacts to on the session event stream. Every other @@ -42,6 +44,17 @@ const ( // cleanly, as opposed to "error", "canceled", "hook_blocked"... const reasonNormal = "normal" +// streamIdleTimeout is how long StreamEvents tolerates a silent connection +// once the server has proven it sends heartbeats (": ping" SSE comments, +// emitted every 15s). Three missed heartbeats means the transport is hung — +// e.g. the agent's VM was paused — not that the session is quiet, so the +// stream is aborted and the watcher reconnects. Servers that predate +// heartbeats never arm the watchdog, keeping long-lived idle streams working. +var streamIdleTimeout = 45 * time.Second + +// errStreamIdle reports a stream aborted by the idle watchdog. +var errStreamIdle = errors.New("event stream idle: heartbeats stopped") + // event is the subset of a runtime event the board cares about. type event struct { Type string `json:"type"` @@ -147,6 +160,8 @@ func (c *client) Followup(ctx context.Context, idempotencyKey, message string) e // false stops the stream cleanly. It returns nil on a clean stop and an // error when the connection fails. func (c *client) StreamEvents(ctx context.Context, since uint64, onEvent func(event) bool) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() u := c.endpoint("events") if since > 0 { u += "?since=" + strconv.FormatUint(since, 10) @@ -167,9 +182,37 @@ func (c *client) StreamEvents(ctx context.Context, since uint64, onEvent func(ev scanner := bufio.NewScanner(resp.Body) scanner.Buffer(make([]byte, 64*1024), 4*1024*1024) + + // The idle watchdog is armed by the first heartbeat and re-armed by any + // subsequent line; when it fires it cancels the request, failing the read. + var idle atomic.Bool + var watchdog *time.Timer + defer func() { + if watchdog != nil { + watchdog.Stop() + } + }() + resetWatchdog := func() { + if watchdog != nil { + watchdog.Reset(streamIdleTimeout) + } + } + var seq uint64 for scanner.Scan() { line := scanner.Text() + if strings.HasPrefix(line, ":") { + // Heartbeat comment: the server sends them, so silence now means + // a hung transport. Arm the watchdog on the first one. + if watchdog == nil { + watchdog = time.AfterFunc(streamIdleTimeout, func() { + idle.Store(true) + cancel() + }) + } + continue + } + resetWatchdog() if id, ok := strings.CutPrefix(line, "id:"); ok { seq, _ = strconv.ParseUint(strings.TrimSpace(id), 10, 64) continue @@ -188,6 +231,9 @@ func (c *client) StreamEvents(ctx context.Context, since uint64, onEvent func(ev return nil } } + if idle.Load() { + return errStreamIdle + } if err := scanner.Err(); err != nil { return err } diff --git a/pkg/board/client_test.go b/pkg/board/client_test.go new file mode 100644 index 0000000000..6c4dfe34b6 --- /dev/null +++ b/pkg/board/client_test.go @@ -0,0 +1,86 @@ +package board + +import ( + "fmt" + "net" + "net/http" + "os" + "path/filepath" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// serveUnix runs an HTTP handler on a unix socket and returns the socket path. +func serveUnix(t *testing.T, handler http.Handler) string { + t.Helper() + // Not t.TempDir(): its per-test path is long enough to overflow the + // ~104-byte unix sun_path limit under long test names. + dir, err := os.MkdirTemp("", "board-client") //nolint:forbidigo,usetesting // see above + require.NoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(dir) }) + + socket := filepath.Join(dir, "cp.sock") + var lc net.ListenConfig + ln, err := lc.Listen(t.Context(), "unix", socket) + require.NoError(t, err) + srv := &http.Server{Handler: handler, ReadHeaderTimeout: time.Second} + go func() { _ = srv.Serve(ln) }() + t.Cleanup(func() { _ = srv.Close() }) + return socket +} + +// Heartbeat comments are invisible to the event callback, and once one has +// been seen, a stream that goes silent is aborted with an error (instead of +// blocking forever on a hung transport) so the watcher reconnects. +func TestStreamEventsIdleWatchdogAbortsSilentStream(t *testing.T) { + old := streamIdleTimeout + streamIdleTimeout = 50 * time.Millisecond + t.Cleanup(func() { streamIdleTimeout = old }) + + socket := serveUnix(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + f := w.(http.Flusher) + fmt.Fprint(w, ": ping\n\n") + fmt.Fprint(w, "data: {\"type\":\"stream_started\"}\n\n") + f.Flush() + // Then hang without closing, like a wedged transport. + <-r.Context().Done() + })) + + c := newClient(socket, "sess-1") + var got []event + err := c.StreamEvents(t.Context(), 0, func(ev event) bool { + got = append(got, ev) + return true + }) + require.ErrorIs(t, err, errStreamIdle) + require.Len(t, got, 1, "heartbeat comments must not reach the callback") + assert.Equal(t, eventStreamStarted, got[0].Type) +} + +// Without any heartbeat from the server (an older docker-agent), the watchdog +// stays unarmed: a quiet stream is left alone and events keep flowing. +func TestStreamEventsNoHeartbeatNoWatchdog(t *testing.T) { + old := streamIdleTimeout + streamIdleTimeout = 50 * time.Millisecond + t.Cleanup(func() { streamIdleTimeout = old }) + + socket := serveUnix(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + f := w.(http.Flusher) + fmt.Fprint(w, "data: {\"type\":\"stream_started\"}\n\n") + f.Flush() + time.Sleep(120 * time.Millisecond) //nolint:forbidigo // real quiet time is the thing under test + fmt.Fprint(w, "data: {\"type\":\"stream_stopped\"}\n\n") + })) + + c := newClient(socket, "sess-1") + var got []event + err := c.StreamEvents(t.Context(), 0, func(ev event) bool { + got = append(got, ev) + return len(got) < 2 + }) + require.NoError(t, err) + require.Len(t, got, 2) +} diff --git a/pkg/server/server.go b/pkg/server/server.go index bf8ee621f5..aea0d5a9e6 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -13,6 +13,7 @@ import ( "slices" "strconv" "strings" + "sync" "time" "github.com/labstack/echo/v4" @@ -31,6 +32,11 @@ type Server struct { e *echo.Echo sm *SessionManager authToken string + // heartbeatInterval is how often an idle /events stream emits an SSE + // comment (": ping") so clients can tell a quiet session from a dead + // transport. SSE comments are invisible to EventSource clients and carry + // no id, so they never interfere with the sequenced stream or its replay. + heartbeatInterval time.Duration } func New(ctx context.Context, sessionStore session.Store, runConfig *config.RuntimeConfig, refreshInterval time.Duration, agentSources config.Sources, authToken string) (*Server, error) { @@ -50,7 +56,7 @@ func NewWithManager(sm *SessionManager, authToken string) *Server { e.Use(BearerTokenMiddleware(authToken)) } - s := &Server{e: e, sm: sm, authToken: authToken} + s := &Server{e: e, sm: sm, authToken: authToken, heartbeatInterval: defaultEventsHeartbeatInterval} s.registerRoutes() return s } @@ -319,10 +325,18 @@ func (s *Server) getSessionStatus(c echo.Context) error { // (stored fields + live runtime state + last event sequence number) so a // client can rebuild its view and then tail /events?since= // without missing any transition. +// +// An unknown session yields a 404 whose body carries +// [api.ErrCodeUnknownSession], so clients can tell "this server does not +// have that session" (wait or recreate) from a route-less 404 produced by a +// binary that predates this endpoint (upgrade/relaunch). func (s *Server) getSessionSnapshot(c echo.Context) error { snapshot, err := s.sm.GetSessionSnapshot(c.Request().Context(), c.Param("id")) if err != nil { - return echo.NewHTTPError(http.StatusNotFound, fmt.Sprintf("session not found: %v", err)) + return c.JSON(http.StatusNotFound, api.ErrorResponse{ + Code: api.ErrCodeUnknownSession, + Message: fmt.Sprintf("session not found: %v", err), + }) } return c.JSON(http.StatusOK, snapshot) } @@ -566,6 +580,9 @@ func (s *Server) steerSession(c echo.Context) error { return c.JSON(http.StatusAccepted, map[string]string{"status": "queued"}) } +// defaultEventsHeartbeatInterval is the default for [Server.heartbeatInterval]. +const defaultEventsHeartbeatInterval = 15 * time.Second + // sessionEvents streams events for a session as Server-Sent Events. The // stream lasts until the client disconnects or the session ends. // @@ -583,6 +600,12 @@ func (s *Server) steerSession(c echo.Context) error { // and the stream closes. A client that receives session_exited should stop. A // stream that closes WITHOUT a session_exited event indicates a transport // drop; the client should reconnect with its last id to replay and continue. +// +// Liveness contract: while connected, the stream emits an SSE comment +// (": ping") every [Server.heartbeatInterval] even when the session is quiet. +// A client that has seen at least one heartbeat can treat a prolonged silence +// as a hung transport and reconnect, instead of waiting forever on a +// connection that will never fail a read (e.g. across a paused VM). func (s *Server) sessionEvents(c echo.Context) error { if !s.sm.HasEventSource(c.Param("id")) { return echo.NewHTTPError(http.StatusNotFound, "no event source for session") @@ -596,11 +619,35 @@ func (s *Server) sessionEvents(c echo.Context) error { c.Response().WriteHeader(http.StatusOK) c.Response().Flush() + // Event writes and heartbeat writes come from different goroutines; + // interleaved partial frames would corrupt the stream, so serialize them. + var writeMu sync.Mutex + + heartbeatCtx, stopHeartbeat := context.WithCancel(c.Request().Context()) + defer stopHeartbeat() + go func() { + ticker := time.NewTicker(s.heartbeatInterval) + defer ticker.Stop() + for { + select { + case <-heartbeatCtx.Done(): + return + case <-ticker.C: + writeMu.Lock() + fmt.Fprint(c.Response(), ": ping\n\n") + c.Response().Flush() + writeMu.Unlock() + } + } + }() + s.sm.StreamEvents(c.Request().Context(), c.Param("id"), since, func(seq uint64, event any) { data, err := json.Marshal(event) if err != nil { return } + writeMu.Lock() + defer writeMu.Unlock() // seq 0 marks a per-connection control event (e.g. gap) that is not // part of the sequenced stream, so it carries no id. if seq > 0 { diff --git a/pkg/server/server_heartbeat_test.go b/pkg/server/server_heartbeat_test.go new file mode 100644 index 0000000000..beb28a102a --- /dev/null +++ b/pkg/server/server_heartbeat_test.go @@ -0,0 +1,127 @@ +package server + +import ( + "bufio" + "context" + "encoding/json" + "net/http" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/api" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/session" +) + +// TestAttachedServer_EventStreamSendsHeartbeats verifies that an idle /events +// stream emits SSE comments so clients can distinguish a quiet session from a +// hung transport. +func TestAttachedServer_EventStreamSendsHeartbeats(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + store := session.NewInMemorySessionStore() + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + sm.AttachRuntime(t.Context(), sess.ID, &fakeRuntime{}, sess) + sm.RegisterEventSource(sess.ID, func(ctx context.Context, _ func(any)) { + <-ctx.Done() // a session that stays quiet + }) + + srv := NewWithManager(sm, "") + srv.heartbeatInterval = 10 * time.Millisecond + + ln, err := Listen(ctx, "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = srv.Serve(ctx, ln) }() + addr := "http://" + ln.Addr().String() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, addr+"/api/sessions/"+sess.ID+"/events", http.NoBody) + require.NoError(t, err) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + // With no events at all, the only traffic is heartbeats. + reader := bufio.NewReader(resp.Body) + for range 2 { + line, err := reader.ReadString('\n') + require.NoError(t, err) + if strings.TrimSpace(line) == "" { + continue + } + assert.Equal(t, ": ping", strings.TrimSpace(line)) + } +} + +// TestAttachedServer_SnapshotReportsAggregateCost verifies the snapshot +// carries the session's total cost so clients need not sum per-message costs. +func TestAttachedServer_SnapshotReportsAggregateCost(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + store := session.NewInMemorySessionStore() + sess := session.New() + sess.AddMessage(session.UserMessage("hello")) + m1 := session.NewAgentMessage("root", &chat.Message{Role: chat.MessageRoleAssistant, Content: "a"}) + m1.Message.Cost = 0.25 + sess.AddMessage(m1) + m2 := session.NewAgentMessage("root", &chat.Message{Role: chat.MessageRoleAssistant, Content: "b"}) + m2.Message.Cost = 0.50 + sess.AddMessage(m2) + require.NoError(t, store.AddSession(ctx, sess)) + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + sm.AttachRuntime(t.Context(), sess.ID, &fakeRuntime{}, sess) + + srv := NewWithManager(sm, "") + ln, err := Listen(ctx, "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = srv.Serve(ctx, ln) }() + addr := "http://" + ln.Addr().String() + + resp := httpDoTCP(t, ctx, http.MethodGet, addr+"/api/sessions/"+sess.ID+"/snapshot", nil) + + var snap api.SessionSnapshotResponse + require.NoError(t, json.Unmarshal(resp, &snap)) + assert.InDelta(t, 0.75, snap.Cost, 1e-9) +} + +// TestAttachedServer_SnapshotUnknownSessionCarriesErrorCode verifies that a +// snapshot of an unknown session 404s with a machine-readable code, so +// clients can tell it apart from a route-less 404 (older binary). +func TestAttachedServer_SnapshotUnknownSessionCarriesErrorCode(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + + sm := NewSessionManager(ctx, config.Sources{}, session.NewInMemorySessionStore(), 0, &config.RuntimeConfig{}) + srv := NewWithManager(sm, "") + ln, err := Listen(ctx, "127.0.0.1:0") + require.NoError(t, err) + go func() { _ = srv.Serve(ctx, ln) }() + addr := "http://" + ln.Addr().String() + + req, err := http.NewRequestWithContext(ctx, http.MethodGet, addr+"/api/sessions/nope/snapshot", http.NoBody) + require.NoError(t, err) + resp, err := http.DefaultClient.Do(req) + require.NoError(t, err) + defer resp.Body.Close() + + assert.Equal(t, http.StatusNotFound, resp.StatusCode) + var body api.ErrorResponse + require.NoError(t, json.NewDecoder(resp.Body).Decode(&body)) + assert.Equal(t, api.ErrCodeUnknownSession, body.Code) +} diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 8d9f05a6ad..e8227c077f 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -345,6 +345,7 @@ func (sm *SessionManager) GetSessionSnapshot(ctx context.Context, id string) (*a Streaming: streaming, Agent: agent, LastEventSeq: lastSeq, + Cost: sess.TotalCost(), }, nil }