diff --git a/CHANGELOG.md b/CHANGELOG.md index 0db5c57..4924658 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/internal/gateway/anthropic.go b/internal/gateway/anthropic.go index 0b8f7b0..94e1672 100644 --- a/internal/gateway/anthropic.go +++ b/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 { diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index ca94c53..8d3c1b3 100644 --- a/internal/gateway/gateway.go +++ b/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) diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index d29a56a..13449b5 100644 --- a/internal/gateway/gateway_test.go +++ b/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) + } +} diff --git a/internal/gateway/openai.go b/internal/gateway/openai.go index ea4d132..b07638f 100644 --- a/internal/gateway/openai.go +++ b/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) diff --git a/internal/provider/openai.go b/internal/provider/openai.go index da79465..f5b6537 100644 --- a/internal/provider/openai.go +++ b/internal/provider/openai.go @@ -1,6 +1,7 @@ package provider import ( + "bufio" "bytes" "context" "encoding/json" @@ -47,10 +48,18 @@ func (o *OpenAI) Name() string { return "openai" } // --- wire types (OpenAI Chat Completions API) --- type openAIReq struct { - Model string `json:"model"` - Messages []openAIMessage `json:"messages"` - MaxTokens int `json:"max_tokens,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` + Model string `json:"model"` + Messages []openAIMessage `json:"messages"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + Stream bool `json:"stream,omitempty"` + StreamOptions *oaiStreamOptions `json:"stream_options,omitempty"` +} + +// oaiStreamOptions asks OpenAI to include a final usage chunk in the stream, so +// the gateway can meter a streamed call. +type oaiStreamOptions struct { + IncludeUsage bool `json:"include_usage"` } type openAIMessage struct { @@ -155,3 +164,112 @@ func (o *OpenAI) Chat(ctx context.Context, req Request) (*Response, error) { }, }, nil } + +// ChatStream implements the Streamer interface: a streaming chat completion that +// yields OpenAI's raw SSE chunks verbatim while accumulating usage. It asks for a +// final usage chunk (stream_options.include_usage) so the call can be metered. +func (o *OpenAI) ChatStream(ctx context.Context, req Request) (ChatStream, error) { + if o.apiKey == "" { + return nil, fmt.Errorf("openai: missing API key") + } + + msgs := make([]openAIMessage, 0, len(req.Messages)+1) + if req.System != "" { + msgs = append(msgs, openAIMessage{Role: string(RoleSystem), Content: req.System}) + } + for _, m := range req.Messages { + msgs = append(msgs, openAIMessage{Role: string(m.Role), Content: m.Content}) + } + + body, err := json.Marshal(openAIReq{ + Model: req.Model, Messages: msgs, MaxTokens: req.MaxTokens, Temperature: req.Temperature, + Stream: true, StreamOptions: &oaiStreamOptions{IncludeUsage: true}, + }) + if err != nil { + return nil, fmt.Errorf("openai: marshaling stream request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("openai: building stream request: %w", err) + } + httpReq.Header.Set("content-type", "application/json") + httpReq.Header.Set("authorization", "Bearer "+o.apiKey) + httpReq.Header.Set("accept", "text/event-stream") + + resp, err := o.http.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return nil, fmt.Errorf("openai: stream request failed: %w", err) + } + if resp.StatusCode != http.StatusOK { + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16)) + _ = resp.Body.Close() + var apiErr openAIError + if json.Unmarshal(raw, &apiErr) == nil && apiErr.Error.Message != "" { + return nil, fmt.Errorf("openai: %s (%s, http %d)", apiErr.Error.Message, apiErr.Error.Type, resp.StatusCode) + } + return nil, fmt.Errorf("openai: http %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) + } + return &oaiStream{body: resp.Body, r: bufio.NewReader(resp.Body)}, nil +} + +// oaiStream forwards OpenAI's SSE bytes line-by-line (verbatim, so the client sees +// authentic SSE) while sniffing each `data:` line for the model and the final +// usage chunk. +type oaiStream struct { + body io.ReadCloser + r *bufio.Reader + usage Usage + model string +} + +var sseData = []byte("data:") + +// Recv returns the next raw SSE line (including its trailing newline) to forward, +// or io.EOF at the end. Usage/model are updated from data lines as they pass. +func (s *oaiStream) Recv() ([]byte, error) { + line, err := s.r.ReadBytes('\n') + if len(line) > 0 { + s.sniff(line) + } + return line, err +} + +// sniff parses a `data: {json}` line for the model and a usage block, ignoring the +// terminal `data: [DONE]` and non-data lines (event:, blank, comments). +func (s *oaiStream) sniff(line []byte) { + t := bytes.TrimSpace(line) + if !bytes.HasPrefix(t, sseData) { + return + } + payload := bytes.TrimSpace(t[len(sseData):]) + if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) { + return + } + var chunk struct { + Model string `json:"model"` + Usage *oaiUsage `json:"usage"` + } + if json.Unmarshal(payload, &chunk) != nil { + return + } + if chunk.Model != "" { + s.model = chunk.Model + } + if chunk.Usage != nil { + s.usage = Usage{PromptTokens: chunk.Usage.PromptTokens, CompletionTokens: chunk.Usage.CompletionTokens} + } +} + +func (s *oaiStream) Usage() Usage { return s.usage } +func (s *oaiStream) Model() string { return s.model } +func (s *oaiStream) Close() error { return s.body.Close() } + +// oaiUsage is the OpenAI usage block (also used by the streaming sniffer). +type oaiUsage struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` +} diff --git a/internal/provider/openai_test.go b/internal/provider/openai_test.go index 020106f..690f994 100644 --- a/internal/provider/openai_test.go +++ b/internal/provider/openai_test.go @@ -3,6 +3,7 @@ package provider import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" "strings" @@ -124,3 +125,51 @@ func TestOpenAIChat_ContextCancel(t *testing.T) { t.Fatalf("expected context.Canceled, got %v", err) } } + +func TestOpenAIChatStream(t *testing.T) { + sse := "data: {\"model\":\"gpt-4o-2024\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\" there\"}}]}\n\n" + + "data: {\"model\":\"gpt-4o-2024\",\"choices\":[],\"usage\":{\"prompt_tokens\":12,\"completion_tokens\":3}}\n\n" + + "data: [DONE]\n\n" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var got openAIReq + _ = json.NewDecoder(r.Body).Decode(&got) + if !got.Stream || got.StreamOptions == nil || !got.StreamOptions.IncludeUsage { + t.Errorf("stream request must set stream + include_usage: %+v", got) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, sse) + })) + defer srv.Close() + + o := NewOpenAI("k") + o.baseURL = srv.URL + st, err := o.ChatStream(context.Background(), Request{Model: "gpt-4o", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != nil { + t.Fatalf("ChatStream: %v", err) + } + defer st.Close() + + var forwarded strings.Builder + for { + chunk, err := st.Recv() + forwarded.Write(chunk) + if err == io.EOF { + break + } + if err != nil { + t.Fatalf("Recv: %v", err) + } + } + // The client receives the provider's SSE verbatim. + if forwarded.String() != sse { + t.Errorf("forwarded stream != upstream:\n got %q\nwant %q", forwarded.String(), sse) + } + // Usage and model are extracted from the stream for metering. + if u := st.Usage(); u.PromptTokens != 12 || u.CompletionTokens != 3 { + t.Errorf("usage = %+v, want 12/3", u) + } + if st.Model() != "gpt-4o-2024" { + t.Errorf("model = %q, want gpt-4o-2024", st.Model()) + } +} diff --git a/internal/provider/stream.go b/internal/provider/stream.go new file mode 100644 index 0000000..262880a --- /dev/null +++ b/internal/provider/stream.go @@ -0,0 +1,25 @@ +package provider + +import "context" + +// Streamer is the optional interface a provider implements to support streaming +// chat. The gateway type-asserts for it; a provider that doesn't implement it +// (Anthropic, Ollama, the stubs, for now) makes a streaming request fall back to a +// clear "unsupported" error rather than silently buffering. +type Streamer interface { + // ChatStream starts a streaming completion. The returned ChatStream yields the + // provider's raw SSE chunks verbatim (so the client receives authentic, + // untranslated SSE) while it accumulates token usage. Honor ctx so the + // governor's kill switch / time budget interrupt an in-flight stream. + ChatStream(ctx context.Context, req Request) (ChatStream, error) +} + +// ChatStream is an open streaming response. Recv returns the next raw SSE chunk to +// forward to the client (io.EOF when the stream ends). After io.EOF, Usage and +// Model report the call's final accounting, parsed from the stream. +type ChatStream interface { + Recv() ([]byte, error) + Usage() Usage + Model() string + Close() error +}