Streaming, cancellation, drain

Overview

Make runs observable in real time without coupling execution to the HTTP request. Persist every run event to an append-only run_events table, fan out live events through an in-memory broadcaster, and expose GET /v1/runs/{id}/events as SSE with Last-Event-ID resume (replay from the table, then attach to the live stream). Add run cancellation, scheduler lanes with draining, panic-recovered goroutines, and an SSE-aware graceful shutdown. A client may disconnect and reattach without losing a single event.

Requirements

Functional

• run_events: append-only (run_id, seq) rows with kind, payload jsonb, created_at.
• The loop emits: run.started, turn.started, block.start, text.delta, thinking.delta, tool_input.delta, block.stop, tool.started, tool.finished, turn.finished, turn.aborted, usage, run.finished, error, stream_degraded.
• `EventSink.Emit` error policy, stated once and obeyed everywhere: log at warn with the run id and event kind, increment `hango.run_events.write_failures`, retry exactly once, and on a second failure emit a `stream_degraded` marker event (best effort) and continue. **A run is never failed because an event could not be written.** Callers therefore ignore `Emit`'s error except to propagate a context cancellation.
• A DB poller (`safego` ticker, `HANGO_QUEUED_POLL_INTERVAL`, default 5s) claims `queued` runs older than `HANGO_QUEUED_POLL_AGE` (default 10s) that are not in the in-memory queue, so a run whose enqueue was lost — process died between insert and enqueue, or the queue was full — still executes. It claims through the phase-5 conditional claim, so it cannot double-execute.
• `hango.run.oldest_queued_age` gauge (seconds since the oldest `queued` run was created) is emitted by the poller; phase 14 alerts on it. It is the single metric that proves runs are not silently stranded.
• GET /v1/runs/{id}/events streams text/event-stream with id: set to the event seq, event: set to the kind, data: a JSON payload.
• Last-Event-ID header (or ?last_event_id=) replays persisted events with a greater seq, then attaches live. Zero gaps, zero duplicates. Replay is a loop, not one query: ListAfter(runID, cursor, limit) is called repeatedly until both conditions hold — the page came back short (fewer rows than limit) and the highest replayed seq is greater than or equal to the first seq staged from the live subscription. Until then the handler keeps paginating, because a busy run can persist more events during replay than one page can carry.
• Events arriving from the live subscription *during* replay are staged in a separate byte-capped slice (`HANGO_SSE_REPLAY_STAGE_BYTES`, default 1MiB), not in the delivery buffer, so the replay phase can never overflow the subscriber buffer and drop the very events it is about to need. If the stage cap is hit, the handler abandons the live attach, closes with a `reconnect` error event and lets the client resume from its last `id:` — correctness by replay, always.
• `seq` is allocated **inside** the same transaction that inserts the `run_events` rows (`UPDATE runs SET last_event_seq = last_event_seq + $n ... RETURNING` then the multi-row insert), so a crash between allocation and insert cannot leave a permanent hole that the replay loop would wait for forever.
• `turn.aborted` event kind: a mid-stream provider failure emits `turn.aborted` carrying the partial message and resets the emitter's block-index state before any retry, so a retried turn restarts block indices from zero and a client never sees two `block.start` events for index 0 without an intervening abort.
• Keepalive : ping comment every HANGO_SSE_KEEPALIVE (default 15s).
• Terminal run states close the stream after emitting run.finished; already-terminal runs replay the full log and close immediately.
• POST /v1/runs/{id}/cancel sets cancel_requested_at, cancels the in-memory run context if local, and returns 202; the run finishes as cancelled with partial state persisted.
• Scheduler lanes: named lanes (main now, subagent in phase 9) with independent concurrency limits, MarkDraining, and Stats.
• Shutdown: mark draining, stop accepting runs, cancel SSE subscribers with a server.draining event, wait for active runs up to HANGO_DRAIN_TIMEOUT, then cancel them, persist the partial state and release the lease so another process (or this one on restart) requeues them at their turn boundary rather than failing them. HANGO_DRAIN_TIMEOUT must be at least one full turn (documented default 120s); config validation warns when it is shorter.
• Broadcaster topic lifetime: a topic lives until the run reaches a terminal state **or** its last subscriber leaves, whichever is later — it is reference-counted, not closed on terminal state. A suspended run (phase 15) therefore keeps no goroutine while SSE clients may stay subscribed; the bound on that is documented and observable through `hango.sse.subscribers` and `hango.broadcaster.topics`. The run-cancel registry entry is removed at exactly two points: when the run reaches a terminal state, and when it suspends (phase 15). Both are in the worker's `defer`.

Non-functional

• High-volume deltas must not write one row per token: deltas are coalesced by the emitter into time/size-bounded batches (HANGO_EVENT_FLUSH_INTERVAL default 100ms, max 8KiB) before persistence, and each flushed batch is published to live subscribers only after it is persisted (step 4 and its persist-before-publish test).
• http.NewResponseController(w).SetWriteDeadline is reset per write so no server timeout kills a long stream.
• Every goroutine in hango is started through safego.Go, which recovers panics, logs the stack and increments a counter.
• Leak detection is behavioural, not profile-based: tests assert `runtime.NumGoroutine()` returns to its pre-test value (with a bounded settle wait) after run completion, subscriber disconnect and cancellation, **and** that `Broadcaster.TopicCount()` returns to zero. The Go 1.27 `goroutineleak` profile is used only in the phase-16 soak test, where a long-running process makes it meaningful; it cannot see the leaks these phases actually risk (a blocked-forever goroutine that is still reachable).

Architecture

// internal/safego/safego.go
func Go(ctx context.Context, log *slog.Logger, component, op string, fn func(context.Context))
func Recover(log *slog.Logger, component, op string) // defer safego.Recover(...)()

// internal/run/broadcaster.go
type Broadcaster struct{} // per-run topics, bounded per-subscriber buffers

type Subscription struct {
	Events <-chan store.RunEvent
	Cancel func()
}

func (b *Broadcaster) Publish(runID uuid.UUID, ev store.RunEvent)
func (b *Broadcaster) Subscribe(runID uuid.UUID, buffer int) *Subscription // slow consumer is dropped, not blocking
func (b *Broadcaster) CloseRun(runID uuid.UUID) // marks terminal; topic is freed when refcount hits 0
func (b *Broadcaster) TopicCount() int          // test + gauge. Red team: F11

// internal/run/emitter.go
type Emitter struct{} // implements agent.EventSink

func (e *Emitter) Emit(ctx context.Context, ev agent.Event) error // seq alloc + coalesce + persist + publish
func (e *Emitter) Flush(ctx context.Context) error                // called at turn boundaries and on finish

// internal/run/scheduler.go (extended)
type Lane struct {
	Name        string
	Concurrency int
}

func (s *Scheduler) MarkDraining()
func (s *Scheduler) Stats() map[string]LaneStats
func (s *Scheduler) Cancel(runID uuid.UUID) bool // true when the run was local and cancelled

// internal/server/sse.go
type SSEWriter struct{}

func (w *SSEWriter) Event(id int64, kind string, payload []byte) error // resets write deadline, flushes
func (w *SSEWriter) Ping() error

-- migrations/00004_run_events.sql
alter table runs add column last_event_seq bigint not null default 0; -- Red team: F8 — seq source

create table run_events (
  id uuid primary key,
  tenant_id uuid not null references tenants(id) on delete cascade,
  run_id uuid not null references runs(id) on delete cascade,
  seq bigint not null,
  kind text not null,
  payload jsonb not null,
  created_at timestamptz not null default now(),
  unique (run_id, seq)
);
create index run_events_tenant_run_seq_idx on run_events (tenant_id, run_id, seq);

Instrumentation

• SSE handler span: sse.run_events SERVER span with hango.run.id, sse.resume_from_seq, sse.replayed_events, sse.live_events, sse.close_reason (client_disconnect|run_finished|draining|lagged).
• Response headers: the SSE response carries X-Hango-Run-Id, X-Trace-Id (the SSE handler’s trace) and traceparent carrying the run’s root span context, so a client can correlate its stream with the run trace it is watching.
• Span events on the run’s root span: run.cancel_requested (with the requesting user), run.drained.
• Metrics (names in internal/telemetry/metrics_names.go): hango.sse.subscribers gauge, hango.broadcaster.topics gauge, hango.sse.events_published counter (attribute kind), hango.sse.lagged_subscribers counter, hango.run_events.rows histogram per run (write-amplification watchdog), hango.run_events.write_failures counter, hango.run.oldest_queued_age gauge, hango.sse.replay_pages histogram, hango.scheduler.draining gauge.
• Log fields: the emitter logs at debug per flush with hango.run.id, batch size and byte count; drops and drains log at warn.

Related Code Files

Create

• internal/safego/safego.go, internal/safego/safego_test.go.
• internal/run/broadcaster.go, internal/run/broadcaster_test.go.
• internal/run/emitter.go, internal/run/emitter_test.go.
• internal/server/sse.go, internal/server/handlers_run_events.go, internal/server/sse_test.go.
• internal/store/pg/run_event.go, internal/store/pg/run_event_test.go.
• internal/run/leak_test.go — before/after runtime.NumGoroutine() assertions plus Broadcaster.TopicCount() == 0, for run completion, client disconnect and cancellation.
• internal/run/poller.go, internal/run/poller_test.go — the queued-run DB poller and its age gauge.
• internal/server/sse_headers_test.go — asserts X-Hango-Run-Id, X-Trace-Id and a well-formed traceparent on the SSE response.
• migrations/00004_run_events.sql.
• docs/learning/06-streaming-and-cancellation.md.

Modify

• internal/agent/loop.go — emit events at every boundary; honour ctx.Err() promptly; finalize on context.WithoutCancel.
• internal/run/scheduler.go, internal/run/worker.go — lanes, draining, cancel registry, emitter wiring.
• internal/server/handlers_runs.go — cancel endpoint; GET /v1/runs/{id} includes last_event_seq.
• internal/server/server.go — SSE route, no write timeout, draining event on shutdown.
• internal/app/app.go — broadcaster, emitter factory, drain order: HTTP draining → scheduler draining → wait → cancel → close pool.
• internal/config/config.go — HANGO_SSE_KEEPALIVE, HANGO_SSE_WRITE_TIMEOUT, HANGO_SSE_REPLAY_STAGE_BYTES, HANGO_EVENT_FLUSH_INTERVAL, HANGO_EVENT_MAX_BATCH_BYTES, HANGO_DRAIN_TIMEOUT (≥ one turn), HANGO_LANE_MAIN_CONCURRENCY, HANGO_QUEUED_POLL_INTERVAL, HANGO_QUEUED_POLL_AGE.
• internal/store/pg/tenancy_test.go — add the run_events case.
• docs/system-architecture.md, docs/learning/README.md, README.md (curl SSE example).

Implementation Steps

1. Write internal/safego and mechanically convert every existing go func() in the codebase to safego.Go.
2. Write migrations/00004_run_events.sql (adds runs.last_event_seq and the run_events table) and RunEventStore: Append(ctx, runID, events...) allocating seq inside the same transaction as the insert (UPDATE runs SET last_event_seq = last_event_seq + $n RETURNING, then the multi-row insert, then commit), and ListAfter(ctx, runID, afterSeq, limit).
3. Write internal/run/broadcaster.go: map[uuid.UUID]*topic, each topic reference-counted by its subscribers and freed only when the run is terminal and the refcount is zero. Each topic holds a set of subscriber channels with a bounded buffer; on a full buffer, drop the subscriber and mark it lagged so the SSE handler can close the connection with an error event telling the client to reconnect with Last-Event-ID (correctness via replay, never blocking the loop). Expose TopicCount().
4. Write internal/run/emitter.go: implements agent.EventSink (defined in internal/agent; the dependency runs one way). Deltas of the same kind and index accumulate in a buffer flushed on interval, size, or any non-delta event; non-delta events flush immediately. Each flushed event is persisted then published, in that order, so a subscriber never sees an event that is not yet resumable. Apply the Emit error policy: log, count, retry once, then emit stream_degraded and continue — never return an error that could fail the run. On turn.aborted, reset the block-index state.
5. Extend the loop to emit the full event vocabulary. Provider deltas map one-to-one from provider.Event; add tool.started/tool.finished carrying tool name, ForUser output and IsError.
6. Write internal/server/sse.go: set Content-Type: text/event-stream, Cache-Control: no-cache, Connection: keep-alive, X-Accel-Buffering: no; on each write call rc.SetWriteDeadline(time.Now().Add(HANGO_SSE_WRITE_TIMEOUT)) then rc.Flush().
7. Write handlers_run_events.go: authorize the run against the request scope (creator-scoped for member, 404 on denial); parse Last-Event-ID; subscribe first and stage live events into the byte-capped replay stage; then loop ListAfter until the page is short and the highest replayed seq ≥ the first staged live seq; then flush the stage (skipping replayed seqs) and drain the subscription live; select over subscription, keepalive ticker, r.Context().Done() and the server draining channel. Record hango.sse.replay_pages. 7a. Write internal/run/poller.go: a safego ticker that selects queued runs older than HANGO_QUEUED_POLL_AGE that are absent from the in-memory queue, enqueues them (dedup makes this safe), and publishes hango.run.oldest_queued_age.
8. Add POST /v1/runs/{id}/cancel: persist cancel_requested_at, call scheduler.Cancel(runID), return 202 with the run. A queued run is cancelled by the worker at claim time. The handler calls SetAuditAction(c, "run.cancel", "run", runID), so a cancelled run always names the human who cancelled it.
9. Extend the loop to check ctx.Err() between turns and after each tool call, and to persist the partial assistant message before returning StatusCancelled.
10. Extend the scheduler with lanes and MarkDraining; Enqueue returns an error once draining. Wire the shutdown order in internal/app.
11. Write tests: broadcaster fan-out, slow-subscriber drop, and refcounted topic lifetime (topic survives terminal state while a subscriber remains, and TopicCount() returns to zero when it leaves); emitter coalescing and ordering (persist-before-publish asserted) and its error policy (two consecutive write failures produce stream_degraded and a succeeded run); resume under load — a run producing more events than one ListAfter page while the client reconnects, asserting the concatenated stream equals the log exactly once (this is the test that catches an unpaginated replay); replay-stage overflow closing with reconnect rather than dropping; turn.aborted resetting block indices across a provider retry; cancel test; drain test asserting a drained run is requeued at its turn boundary, not failed; poller test (a run inserted without enqueue still executes); before/after goroutine-count and topic-count assertions.
12. Update docs; add a curl -N SSE example to the README.

Learning Doc Outline

docs/learning/06-streaming-and-cancellation.md must contain:

• Concept: why the run is the unit of durability and SSE is a subscriber, not the executor.
• The SSE wire format: id:/event:/data:/comment lines, Last-Event-ID, reconnect semantics, proxy buffering, and why X-Accel-Buffering: no matters.
• Design choices: persist-then-publish, subscribe-then-replay with de-dup, delta coalescing for write amplification, dropping slow subscribers instead of blocking the loop, ResponseController write-deadline reset, keepalive pings, cancel via a persisted flag plus an in-memory context, finalize on context.WithoutCancel.
• Trade-offs: SSE vs WebSocket vs long polling; in-memory broadcaster vs listen/notify vs a queue (and exactly what breaks at two replicas).
• goclaw comparison: WebSocket RPC driven by chat channels, onChunk callbacks bubbling to AgentEvent, FinalizeStage on a background context, activeRuns atomic for drain, MarkDraining in the lane manager, safego.Recover.
• How to debug this layer: “the client saw nothing” is answered by the sse.run_events span (was it subscribed, what was resume_from_seq, what was close_reason) plus hango.sse.lagged_subscribers; “events were duplicated” is answered by comparing the persisted run_events seq range with the handler’s replayed_events/live_events counts; a hung shutdown is answered by hango.scheduler.draining and the run.drained span events.
• deepagents comparison: streaming is LangGraph’s astream; suspension and resumption go through a checkpointer and thread_id, which is the same insight hango implements as durable runs plus an event log.
• Why hango diverges: SSE only, a persisted event log as the source of truth for replay (goclaw cannot resume a dropped stream).
• Why "replay then attach" is a catch-up loop and not one query, and why live events must be staged outside the delivery buffer during replay.
• Why goroutine leaks here are best detected by counting, not by the `goroutineleak` profile.
• Exercises: kill curl mid-stream and reattach; make a subscriber slow with a debugger breakpoint and observe the lag policy; insert a queued run straight into Postgres and watch the poller pick it up.
• Time budget: 14h implementation + 2h learning doc.

Success Criteria

• go build ./... && go vet ./... && go test ./... pass.
• SSE resume test: disconnect after N events, reconnect with Last-Event-ID: N, and the concatenated stream equals the full event log exactly once — including the variant where the run emits more events than one replay page while the client is reconnecting.
• A terminal run’s SSE endpoint replays the whole log and closes without hanging.
• Cancel test: run reaches cancelled, run.finished is emitted, the partial assistant message is persisted, and an audit_events row with action run.cancel names the deciding user.
• Drain test: SIGTERM during an active run produces succeeded, or a requeued run whose lease was released and whose resume_turn is its last completed turn — never a lost row and never a failed run caused only by shutdown; the process exits within HANGO_DRAIN_TIMEOUT + 5s.
• Goroutine count returns to its pre-test value and Broadcaster.TopicCount() returns to zero after run completion, client disconnect and cancellation.
• Two consecutive run_events write failures produce stream_degraded and a run that still reaches succeeded; hango.run_events.write_failures increments.
• A run inserted directly into the database without an enqueue is picked up by the poller within HANGO_QUEUED_POLL_INTERVAL + HANGO_QUEUED_POLL_AGE, and hango.run.oldest_queued_age returns to near zero.
• Delta coalescing: a 2000-token fake response produces fewer than 50 run_events rows (asserted).
• go test ./internal/store/pg/ -run TestCrossTenant includes run_events.
• The SSE response carries X-Hango-Run-Id, X-Trace-Id and a traceparent pointing at the run’s root span (test).
• In Grafana, an sse.run_events span with resume and close-reason attributes is visible, and hango.sse.subscribers/hango.sse.events_published series exist.

Risk Assessment

Rollback

goose down one step drops run_events and the runs.last_event_seq column. Removing the SSE route returns clients to phase 5 polling, which keeps working because the run row remains the authoritative state.