feat(gateway): streaming proxy with budget enforcement (OpenAI)

CHANGELOG.md

@@ -10,6 +10,14 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md).
 ## [Unreleased]
 
 ### Added
+- **Streaming proxy.** `POST /v1/chat/completions` now supports `stream:true`: the
+  budget is enforced before the stream opens, the OpenAI provider's SSE is forwarded
+  to the client verbatim (authentic chunks, no translation) while token usage is
+  metered from the final usage chunk, and the run's context — time budget, kill
+  switch, or client disconnect — cuts a live stream. Dollar/token budgets are
+  checked pre-stream and recorded after (so the next call is refused if it went
+  over). A provider without streaming, and the Anthropic `/v1/messages` endpoint,
+  return a clear 501 rather than silently buffering.
 - **Prometheus `/metrics` endpoint.** Scrape the daemon's own state: governed runs
   by status, halted runs by halt reason, total spend in dollars and tokens, priced
   model calls, and the pending-approval queue depth. Plain Prometheus text

internal/gateway/anthropic.go

@@ -59,8 +59,10 @@ func (g *Gateway) handleMessages(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 	if req.Stream {
+		// The OpenAI-compatible /v1/chat/completions path streams; native Anthropic
+		// /v1/messages streaming is not wired yet (its SSE event format differs).
 		httpx.WriteError(w, http.StatusNotImplemented, "streaming_unsupported",
-			"streaming is not supported in v0.1; set stream:false")
+			"streaming is not yet supported on /v1/messages; set stream:false (or use /v1/chat/completions)")
 		return
 	}
 	if req.Model == "" || len(req.Messages) == 0 {

internal/gateway/gateway.go

@@ -6,8 +6,12 @@
 // priced into the cost ledger, and forwarded to the real provider with the
 // user's key.
 //
-// Streaming is not supported in v0.1 (mid-stream budget enforcement is deferred);
-// stream requests are rejected with a clear error rather than silently degraded.
+// Streaming (`stream:true`) is supported on the OpenAI-compatible endpoint: the
+// budget is enforced before the stream opens, the provider's SSE is forwarded to
+// the client verbatim while token usage is metered from it, and the run's context
+// (time budget / kill switch / client disconnect) cuts a live stream. Providers
+// that don't implement streaming, and the Anthropic /v1/messages endpoint, reject
+// a stream request with a clear error rather than silently degrading.
 package gateway
 
 import (
@@ -195,6 +199,100 @@ func maxF(a, b float64) float64 {
 	return b
 }
 
+// streamCall proxies a streaming completion: enforce the budget before opening the
+// stream, forward the provider's SSE chunks verbatim to the client (flushing each),
+// and meter the call from the stream's final usage. The run's context (time budget
+// / kill switch) and client disconnect cut a live stream. Only providers that
+// implement provider.Streamer support this; others get a clear 501. Dollar/token
+// budgets are checked before the stream and recorded after (so the next call is
+// refused if over); the time budget and kill switch cut mid-stream via the context.
+func (g *Gateway) streamCall(w http.ResponseWriter, r *http.Request, run *runs.Run, preq provider.Request) {
+	step, err := run.BeginStep()
+	if err != nil {
+		budgetError(err).write(w)
+		return
+	}
+	if err := run.CanProceed(); err != nil {
+		budgetError(err).write(w)
+		return
+	}
+
+	prov, err := g.providers.Get(routeModel(preq.Model))
+	if err != nil {
+		(&gwError{http.StatusBadRequest, "unknown_provider", err.Error()}).write(w)
+		return
+	}
+	streamer, ok := prov.(provider.Streamer)
+	if !ok {
+		(&gwError{http.StatusNotImplemented, "streaming_unsupported",
+			"streaming is not supported for provider " + prov.Name() + "; set stream:false"}).write(w)
+		return
+	}
+	flusher, ok := w.(http.Flusher)
+	if !ok {
+		(&gwError{http.StatusInternalServerError, "internal_error", "response writer does not support streaming"}).write(w)
+		return
+	}
+
+	// The stream dies if the run is governed-cancelled/expired (parent) or the
+	// client goes away.
+	callCtx, cancel := context.WithCancel(run.Context())
+	defer cancel()
+	stop := context.AfterFunc(r.Context(), cancel)
+	defer stop()
+
+	start := time.Now()
+	stream, serr := streamer.ChatStream(callCtx, preq)
+	if serr != nil {
+		if run.Halted() {
+			haltGWError(run.HaltReason()).write(w)
+			return
+		}
+		(&gwError{http.StatusBadGateway, "provider_error", serr.Error()}).write(w)
+		return
+	}
+	defer stream.Close()
+
+	// Past here the response is committed (200 + SSE); all budget pre-checks passed.
+	h := w.Header()
+	h.Set("Content-Type", "text/event-stream")
+	h.Set("Cache-Control", "no-cache")
+	h.Set("Connection", "keep-alive")
+	h.Set(HeaderRunID, run.ID)
+	h.Set(headerStep, strconv.Itoa(int(step)))
+	w.WriteHeader(http.StatusOK)
+	flusher.Flush()
+
+	for {
+		chunk, rerr := stream.Recv()
+		if len(chunk) > 0 {
+			if _, werr := w.Write(chunk); werr != nil {
+				break // client disconnected
+			}
+			flusher.Flush()
+		}
+		if rerr != nil {
+			break // io.EOF (clean end) or an upstream/context error (truncates the stream)
+		}
+	}
+	end := time.Now()
+
+	// Meter the (possibly truncated) call from the usage the stream reported, so the
+	// ledger and budget reflect it and the next call is refused if it went over.
+	model := stream.Model()
+	if model == "" {
+		model = preq.Model
+	}
+	usage := stream.Usage()
+	cost, priced := g.prices.Cost(model, usage.PromptTokens, usage.CompletionTokens)
+	_ = run.RecordCall(runs.Call{
+		StepIndex: step, Provider: prov.Name(), Model: model,
+		PromptTokens: usage.PromptTokens, CompletionTokens: usage.CompletionTokens,
+		Dollars: cost, Priced: priced,
+	})
+	g.emitCallSpan(run, step, prov.Name(), preq, &provider.Response{Model: model, Usage: usage}, cost, priced, run.HaltReason(), start, end)
+}
+
 // stampHeaders writes the governance headers onto a successful proxied response.
 func stampHeaders(w http.ResponseWriter, run *runs.Run, resp *provider.Response, meta callMeta) {
 	w.Header().Set(HeaderRunID, run.ID)

internal/gateway/gateway_test.go

@@ -265,3 +265,116 @@ func TestDecodeContent(t *testing.T) {
 		t.Errorf("nil content = %q", got)
 	}
 }
+
+// --- streaming (#22) ---
+
+type fakeStream struct {
+	chunks [][]byte
+	i      int
+	usage  provider.Usage
+	model  string
+}
+
+func (s *fakeStream) Recv() ([]byte, error) {
+	if s.i >= len(s.chunks) {
+		return nil, io.EOF
+	}
+	c := s.chunks[s.i]
+	s.i++
+	return c, nil
+}
+func (s *fakeStream) Usage() provider.Usage { return s.usage }
+func (s *fakeStream) Model() string         { return s.model }
+func (s *fakeStream) Close() error          { return nil }
+
+// fakeStreamer is a fakeProvider that also supports streaming.
+type fakeStreamer struct {
+	fakeProvider
+	stream    *fakeStream
+	streamErr error
+}
+
+func (f *fakeStreamer) ChatStream(ctx context.Context, _ provider.Request) (provider.ChatStream, error) {
+	atomic.AddInt32(&f.calls, 1)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if f.streamErr != nil {
+		return nil, f.streamErr
+	}
+	return f.stream, nil
+}
+
+func newStreamGateway(t *testing.T, budget governor.Budget, fp provider.Provider) *Gateway {
+	t.Helper()
+	reg, err := provider.NewRegistry(fp.Name(), fp)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return New(reg, runs.NewManager(budget), pricing.NewTable(nil), otel.Disabled(), slog.New(slog.NewTextHandler(io.Discard, nil)))
+}
+
+func newSSEStreamer() *fakeStreamer {
+	return &fakeStreamer{
+		fakeProvider: fakeProvider{name: "openai"},
+		stream: &fakeStream{
+			chunks: [][]byte{
+				[]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n"),
+				[]byte("data: [DONE]\n\n"),
+			},
+			usage: provider.Usage{PromptTokens: 10, CompletionTokens: 5},
+			model: "gpt-4o-2024",
+		},
+	}
+}
+
+func TestStreamingProxy_ForwardsAndMeters(t *testing.T) {
+	g := newStreamGateway(t, governor.Budget{Tokens: 1000}, newSSEStreamer())
+	w := postChat(g, "stream-run", `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"stream":true}`)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("status = %d, body=%s", w.Code, w.Body.String())
+	}
+	if ct := w.Header().Get("Content-Type"); ct != "text/event-stream" {
+		t.Errorf("content-type = %q", ct)
+	}
+	body := w.Body.String()
+	if !strings.Contains(body, `"delta"`) || !strings.Contains(body, "[DONE]") {
+		t.Errorf("client did not receive the SSE verbatim: %q", body)
+	}
+	// The streamed call is metered against the run from the stream's usage.
+	run, ok := g.runs.Get("stream-run")
+	if !ok {
+		t.Fatal("run not found")
+	}
+	if v := run.View(); v.Usage.Tokens() != 15 || v.Usage.Loops != 1 {
+		t.Fatalf("streamed usage not recorded: %+v", v.Usage)
+	}
+}
+
+func TestStreamingProxy_BudgetRefusedBeforeStream(t *testing.T) {
+	fs := newSSEStreamer()
+	g := newStreamGateway(t, governor.Budget{Loops: 1}, fs)
+	body := `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"stream":true}`
+
+	if w := postChat(g, "r", body); w.Code != http.StatusOK { // 1st: loops 0→1
+		t.Fatalf("first stream status = %d", w.Code)
+	}
+	// 2nd: loop budget is spent → refused at 402 BEFORE the stream opens.
+	w := postChat(g, "r", body)
+	if w.Code != http.StatusPaymentRequired {
+		t.Fatalf("second stream status = %d, want 402", w.Code)
+	}
+	if got := atomic.LoadInt32(&fs.calls); got != 1 {
+		t.Fatalf("ChatStream called %d times; the refused call must not reach the provider", got)
+	}
+}
+
+func TestStreamingProxy_UnsupportedProvider(t *testing.T) {
+	// A plain provider (no Streamer) → a clear 501, not a silent buffer.
+	g := newStreamGateway(t, governor.Budget{Tokens: 1000}, &fakeProvider{name: "openai", resp: &provider.Response{}})
+	w := postChat(g, "r", `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"stream":true}`)
+	if w.Code != http.StatusNotImplemented {
+		t.Fatalf("status = %d, want 501", w.Code)
+	}
+}

internal/gateway/openai.go

@@ -62,11 +62,6 @@ func (g *Gateway) handleChatCompletions(w http.ResponseWriter, r *http.Request)
 		httpx.WriteError(w, http.StatusBadRequest, "bad_request", "invalid JSON: "+err.Error())
 		return
 	}
-	if req.Stream {
-		httpx.WriteError(w, http.StatusNotImplemented, "streaming_unsupported",
-			"streaming is not supported in v0.1 (mid-stream budget enforcement is deferred); set stream:false")
-		return
-	}
 	if req.Model == "" || len(req.Messages) == 0 {
 		httpx.WriteError(w, http.StatusBadRequest, "bad_request", "model and messages are required")
 		return
@@ -86,6 +81,15 @@ func (g *Gateway) handleChatCompletions(w http.ResponseWriter, r *http.Request)
 	}
 
 	run := g.resolveRun(r)
+
+	// Streaming: forward the provider's SSE verbatim while metering it. The budget
+	// is enforced before the stream opens; the run's context (time budget / kill
+	// switch / client disconnect) cuts a live stream.
+	if req.Stream {
+		g.streamCall(w, r, run, preq)
+		return
+	}
+
 	resp, meta, gwErr := g.governedCall(r, run, preq)
 	if gwErr != nil {
 		gwErr.write(w)