diff --git a/CHANGELOG.md b/CHANGELOG.md index 904aaab..cd8e37d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,15 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] ### Added +- **Streaming proxy.** Both `POST /v1/chat/completions` and `POST /v1/messages` now + support `stream:true`: the budget is enforced before the stream opens, the + provider's SSE is forwarded to the client verbatim (authentic OpenAI or Anthropic + chunks, no translation) while token usage is metered from the stream's own + accounting, 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 whose + backend doesn't implement streaming returns a clear 501 rather than silently + buffering. - **Python SDK: LlamaIndex adapter.** `RiskKernelCallbackHandler` (from `riskkernel.adapters.llama_index`) is a LlamaIndex `BaseCallbackHandler` that ticks one governed step per LLM call (`CBEventType.LLM`), so a run's loop/time budget is diff --git a/internal/gateway/anthropic.go b/internal/gateway/anthropic.go index 94e1672..6e68014 100644 --- a/internal/gateway/anthropic.go +++ b/internal/gateway/anthropic.go @@ -58,13 +58,6 @@ func (g *Gateway) handleMessages(w http.ResponseWriter, r *http.Request) { httpx.WriteError(w, http.StatusBadRequest, "bad_request", "invalid JSON: "+err.Error()) 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 yet supported on /v1/messages; set stream:false (or use /v1/chat/completions)") - return - } if req.Model == "" || len(req.Messages) == 0 { httpx.WriteError(w, http.StatusBadRequest, "bad_request", "model and messages are required") return @@ -85,6 +78,15 @@ func (g *Gateway) handleMessages(w http.ResponseWriter, r *http.Request) { } run := g.resolveRun(r) + + // Streaming: forward Anthropic's SSE events verbatim while metering them. 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/gateway/gateway.go b/internal/gateway/gateway.go index 8d3c1b3..73e9f49 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -6,12 +6,13 @@ // priced into the cost ledger, and forwarded to the real provider with the // user's key. // -// 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. +// Streaming (`stream:true`) is supported on both endpoints (OpenAI +// /v1/chat/completions and Anthropic /v1/messages): 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. A provider whose backend doesn't +// implement streaming rejects a stream request with a clear error rather than +// silently degrading. package gateway import ( diff --git a/internal/gateway/gateway_test.go b/internal/gateway/gateway_test.go index 13449b5..63eadae 100644 --- a/internal/gateway/gateway_test.go +++ b/internal/gateway/gateway_test.go @@ -378,3 +378,48 @@ func TestStreamingProxy_UnsupportedProvider(t *testing.T) { t.Fatalf("status = %d, want 501", w.Code) } } + +// newAnthropicSSEStreamer is an Anthropic streamer emitting authentic Anthropic +// SSE events (event: + data: lines), used to exercise the /v1/messages path. +func newAnthropicSSEStreamer() *fakeStreamer { + return &fakeStreamer{ + fakeProvider: fakeProvider{name: "anthropic"}, + stream: &fakeStream{ + chunks: [][]byte{ + []byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":11,\"output_tokens\":1}}}\n\n"), + []byte("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n"), + []byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"), + }, + usage: provider.Usage{PromptTokens: 11, CompletionTokens: 7}, + model: "claude-sonnet-4-5", + }, + } +} + +func TestMessagesStreamingProxy_ForwardsAndMeters(t *testing.T) { + g := newStreamGateway(t, governor.Budget{Tokens: 1000}, newAnthropicSSEStreamer()) + r := httptest.NewRequest(http.MethodPost, "/v1/messages", + strings.NewReader(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}],"stream":true}`)) + r.Header.Set(HeaderRunID, "msg-stream-run") + w := httptest.NewRecorder() + g.handleMessages(w, r) + + 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, "message_start") || !strings.Contains(body, "message_stop") { + t.Errorf("client did not receive the Anthropic SSE verbatim: %q", body) + } + // The streamed call is metered against the run from the stream's usage. + run, ok := g.runs.Get("msg-stream-run") + if !ok { + t.Fatal("run not found") + } + if v := run.View(); v.Usage.Tokens() != 18 || v.Usage.Loops != 1 { + t.Fatalf("streamed usage not recorded: %+v", v.Usage) + } +} diff --git a/internal/provider/anthropic.go b/internal/provider/anthropic.go index 3606d9d..2bee7df 100644 --- a/internal/provider/anthropic.go +++ b/internal/provider/anthropic.go @@ -1,6 +1,7 @@ package provider import ( + "bufio" "bytes" "context" "encoding/json" @@ -60,6 +61,7 @@ type anthropicReq struct { System string `json:"system,omitempty"` Messages []anthropicMessage `json:"messages"` Temperature *float64 `json:"temperature,omitempty"` + Stream bool `json:"stream,omitempty"` } type anthropicMessage struct { @@ -94,30 +96,10 @@ func (a *Anthropic) Chat(ctx context.Context, req Request) (*Response, error) { if a.apiKey == "" { return nil, fmt.Errorf("anthropic: missing API key") } - maxTokens := req.MaxTokens - if maxTokens <= 0 { - maxTokens = defaultMaxTokens - } - - // Anthropic takes the system prompt as a top-level field. If the caller put a - // system message in Messages, lift it out; otherwise use req.System. - system := req.System - msgs := make([]anthropicMessage, 0, len(req.Messages)) - for _, m := range req.Messages { - if m.Role == RoleSystem { - if system == "" { - system = m.Content - } else { - system = system + "\n\n" + m.Content - } - continue - } - msgs = append(msgs, anthropicMessage{Role: string(m.Role), Content: m.Content}) - } - + system, msgs := splitSystem(req) body, err := json.Marshal(anthropicReq{ Model: req.Model, - MaxTokens: maxTokens, + MaxTokens: anthropicMaxTokens(req), System: system, Messages: msgs, Temperature: req.Temperature, @@ -181,3 +163,162 @@ func (a *Anthropic) Chat(ctx context.Context, req Request) (*Response, error) { }, }, nil } + +// anthropicMaxTokens returns the request's MaxTokens, falling back to the default +// (Anthropic requires the field to be present and positive). +func anthropicMaxTokens(req Request) int { + if req.MaxTokens > 0 { + return req.MaxTokens + } + return defaultMaxTokens +} + +// splitSystem lifts any system message out of req.Messages and merges it with +// req.System — Anthropic takes the system prompt as a top-level field — returning +// the system prompt and the remaining conversation messages. +func splitSystem(req Request) (string, []anthropicMessage) { + system := req.System + msgs := make([]anthropicMessage, 0, len(req.Messages)) + for _, m := range req.Messages { + if m.Role == RoleSystem { + if system == "" { + system = m.Content + } else { + system = system + "\n\n" + m.Content + } + continue + } + msgs = append(msgs, anthropicMessage{Role: string(m.Role), Content: m.Content}) + } + return system, msgs +} + +// ChatStream implements the Streamer interface: a streaming completion that yields +// Anthropic's raw SSE events verbatim (so the client receives authentic Anthropic +// SSE) while accumulating token usage for metering. Usage is assembled from the +// stream's own accounting: message_start carries input_tokens (and the model), +// message_delta carries the final cumulative output_tokens. +func (a *Anthropic) ChatStream(ctx context.Context, req Request) (ChatStream, error) { + if a.apiKey == "" { + return nil, fmt.Errorf("anthropic: missing API key") + } + + system, msgs := splitSystem(req) + body, err := json.Marshal(anthropicReq{ + Model: req.Model, + MaxTokens: anthropicMaxTokens(req), + System: system, + Messages: msgs, + Temperature: req.Temperature, + Stream: true, + }) + if err != nil { + return nil, fmt.Errorf("anthropic: marshaling stream request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.baseURL+"/v1/messages", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("anthropic: building stream request: %w", err) + } + httpReq.Header.Set("content-type", "application/json") + httpReq.Header.Set("x-api-key", a.apiKey) + httpReq.Header.Set("anthropic-version", anthropicAPIVersion) + httpReq.Header.Set("accept", "text/event-stream") + + resp, err := a.http.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return nil, fmt.Errorf("anthropic: stream request failed: %w", err) + } + if resp.StatusCode != http.StatusOK { + raw, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<16)) + _ = resp.Body.Close() + var apiErr anthropicError + if json.Unmarshal(raw, &apiErr) == nil && apiErr.Error.Message != "" { + return nil, fmt.Errorf("anthropic: %s (%s, http %d)", apiErr.Error.Message, apiErr.Error.Type, resp.StatusCode) + } + return nil, fmt.Errorf("anthropic: http %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) + } + return &antStream{body: resp.Body, r: bufio.NewReader(resp.Body)}, nil +} + +// antStream forwards Anthropic's SSE bytes line-by-line (verbatim, so the client +// sees authentic Anthropic events) while sniffing the data lines for the model and +// token usage. +type antStream struct { + body io.ReadCloser + r *bufio.Reader + usage Usage + model string +} + +// 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 *antStream) 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 (message_start) and token usage +// (input_tokens on message_start, final output_tokens on message_delta), ignoring +// `event:` lines, blanks, and content deltas. +func (s *antStream) sniff(line []byte) { + t := bytes.TrimSpace(line) + if !bytes.HasPrefix(t, sseData) { + return + } + payload := bytes.TrimSpace(t[len(sseData):]) + if len(payload) == 0 { + return + } + var ev struct { + Type string `json:"type"` + Message *struct { + Model string `json:"model"` + Usage *antStreamUsage `json:"usage"` + } `json:"message"` + Usage *antStreamUsage `json:"usage"` + } + if json.Unmarshal(payload, &ev) != nil { + return + } + switch ev.Type { + case "message_start": + if ev.Message == nil { + return + } + if ev.Message.Model != "" { + s.model = ev.Message.Model + } + if u := ev.Message.Usage; u != nil { + s.usage.PromptTokens = u.InputTokens + s.usage.CompletionTokens = u.OutputTokens + } + case "message_delta": + // message_delta carries the running (final, at stream end) output token + // count, and on cache paths an updated input count. + if u := ev.Usage; u != nil { + if u.OutputTokens > 0 { + s.usage.CompletionTokens = u.OutputTokens + } + if u.InputTokens > 0 { + s.usage.PromptTokens = u.InputTokens + } + } + } +} + +func (s *antStream) Usage() Usage { return s.usage } +func (s *antStream) Model() string { return s.model } +func (s *antStream) Close() error { return s.body.Close() } + +// antStreamUsage is the usage block carried on Anthropic stream events. +type antStreamUsage struct { + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` +} diff --git a/internal/provider/anthropic_test.go b/internal/provider/anthropic_test.go index d1a42e1..fc98514 100644 --- a/internal/provider/anthropic_test.go +++ b/internal/provider/anthropic_test.go @@ -3,6 +3,7 @@ package provider import ( "context" "encoding/json" + "io" "net/http" "net/http/httptest" "strings" @@ -101,3 +102,89 @@ func TestAnthropicChat_ContextCancel(t *testing.T) { t.Fatalf("expected context.Canceled, got %v", err) } } + +func TestAnthropicChatStream(t *testing.T) { + // Authentic Anthropic SSE: input_tokens arrive on message_start, the final + // (cumulative) output_tokens on message_delta. + sse := "event: message_start\n" + + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"model\":\"claude-sonnet-4-5-20250101\",\"usage\":{\"input_tokens\":11,\"output_tokens\":1}}}\n\n" + + "event: content_block_delta\n" + + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hi\"}}\n\n" + + "event: content_block_delta\n" + + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\" there\"}}\n\n" + + "event: message_delta\n" + + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":7}}\n\n" + + "event: message_stop\n" + + "data: {\"type\":\"message_stop\"}\n\n" + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var got anthropicReq + _ = json.NewDecoder(r.Body).Decode(&got) + if !got.Stream { + t.Errorf("stream request must set stream:true: %+v", got) + } + if r.Header.Get("accept") != "text/event-stream" { + t.Errorf("accept = %q, want text/event-stream", r.Header.Get("accept")) + } + if r.Header.Get("x-api-key") != "k" { + t.Errorf("missing x-api-key: %q", r.Header.Get("x-api-key")) + } + w.Header().Set("Content-Type", "text/event-stream") + _, _ = io.WriteString(w, sse) + })) + defer srv.Close() + + a := NewAnthropic("k") + a.baseURL = srv.URL + st, err := a.ChatStream(context.Background(), Request{Model: "claude-sonnet-4-5", 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 Anthropic's SSE verbatim. + if forwarded.String() != sse { + t.Errorf("forwarded stream != upstream:\n got %q\nwant %q", forwarded.String(), sse) + } + // Usage is assembled from the stream: input from message_start, the final + // output from message_delta (overriding message_start's initial 1). + if u := st.Usage(); u.PromptTokens != 11 || u.CompletionTokens != 7 { + t.Errorf("usage = %+v, want 11/7", u) + } + if st.Model() != "claude-sonnet-4-5-20250101" { + t.Errorf("model = %q, want claude-sonnet-4-5-20250101", st.Model()) + } +} + +func TestAnthropicChatStream_APIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"type":"error","error":{"type":"invalid_request_error","message":"bad model"}}`)) + })) + defer srv.Close() + + a := NewAnthropic("k") + a.baseURL = srv.URL + _, err := a.ChatStream(context.Background(), Request{Model: "x", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "bad model") { + t.Fatalf("expected API error surfaced, got %v", err) + } +} + +func TestAnthropicChatStream_MissingKey(t *testing.T) { + a := NewAnthropic("") + _, err := a.ChatStream(context.Background(), Request{Model: "x", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "missing API key") { + t.Fatalf("expected missing-key error, got %v", err) + } +}