diff --git a/docs/advanced/anthropic-messages-api.mdx b/docs/advanced/anthropic-messages-api.mdx index 5695b0d12..b834584e6 100644 --- a/docs/advanced/anthropic-messages-api.mdx +++ b/docs/advanced/anthropic-messages-api.mdx @@ -31,7 +31,8 @@ different name), then relays the provider-native response or SSE stream unchanged. This preserves everything the canonical translation cannot — `cache_control` breakpoints, thinking-block signatures, `anthropic-beta` headers — which coding agents like Claude Code depend on. Rate limits, -budgets, audit logging, and streaming usage tracking still apply. +budgets, audit logging, and usage tracking still apply, for both streaming +and non-streaming responses. Native forwarding is automatic. Requests fall back to the translated pipeline when a feature that operates on the canonical request is in play: guardrails diff --git a/docs/features/passthrough-api.mdx b/docs/features/passthrough-api.mdx index 9386eb454..5a432887f 100644 --- a/docs/features/passthrough-api.mdx +++ b/docs/features/passthrough-api.mdx @@ -69,6 +69,11 @@ Because passthrough is provider-native, the response is also provider-native. For Anthropic messages, the response uses Anthropic's message schema, not an OpenAI chat completion schema. +When usage tracking is enabled, successful passthrough inference responses +are recorded: token counts are read from SSE usage events on streaming +responses and from the `usage` member of JSON responses, so costs and budgets +account for passthrough traffic like any other route. + ## Anthropic SDK example Set the Anthropic SDK base URL to GoModel's Anthropic passthrough route. Use the diff --git a/internal/server/messages_native_test.go b/internal/server/messages_native_test.go index 1838be0a8..b02eb40ce 100644 --- a/internal/server/messages_native_test.go +++ b/internal/server/messages_native_test.go @@ -108,6 +108,104 @@ func TestBuildPassthroughHeadersDropsAcceptEncoding(t *testing.T) { } } +const anthropicNonStreamingJSON = `{"id":"msg_1","type":"message","role":"assistant","model":"claude-fable-5","content":[{"type":"text","text":"ok"}],"stop_reason":"end_turn","usage":{"input_tokens":100,"cache_creation_input_tokens":7,"cache_read_input_tokens":9,"output_tokens":25}}` + +// A non-streaming /v1/messages request through the native forwarding path +// must record a usage entry from the response object's usage member, so cost +// accounting and budgets see the spend just like on the translated pipeline. +func TestMessages_NativeNonStreamingLogsUsage(t *testing.T) { + provider := &mockProvider{ + supportedModels: []string{"claude-fable-5"}, + providerTypes: map[string]string{"claude-fable-5": "anthropic"}, + passthroughResponse: &core.PassthroughResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(anthropicNonStreamingJSON)), + }, + } + usageLogger := &collectingUsageLogger{config: usage.Config{Enabled: true}} + + e := echo.New() + handler := NewHandler(provider, nil, usageLogger, nil) + + reqBody := `{"model":"claude-fable-5","max_tokens":64,"messages":[{"role":"user","content":"Hi"}]}` + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(reqBody)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + if err := handler.Messages(e.NewContext(req, rec)); err != nil { + t.Fatalf("Messages: %v", err) + } + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String()) + } + if provider.lastPassthroughReq == nil { + t.Fatal("native path was not taken") + } + if rec.Body.String() != anthropicNonStreamingJSON { + t.Fatalf("body not relayed verbatim: %s", rec.Body.String()) + } + + if len(usageLogger.entries) != 1 { + t.Fatalf("usage entries = %d, want 1", len(usageLogger.entries)) + } + entry := usageLogger.entries[0] + if entry.InputTokens != 100 { + t.Errorf("InputTokens = %d, want 100", entry.InputTokens) + } + if entry.OutputTokens != 25 { + t.Errorf("OutputTokens = %d, want 25", entry.OutputTokens) + } + if entry.RawData["cache_creation_input_tokens"] != 7 { + t.Errorf("cache_creation_input_tokens = %v, want 7", entry.RawData["cache_creation_input_tokens"]) + } + if entry.RawData["cache_read_input_tokens"] != 9 { + t.Errorf("cache_read_input_tokens = %v, want 9", entry.RawData["cache_read_input_tokens"]) + } + if entry.ProviderID != "msg_1" { + t.Errorf("ProviderID = %q, want msg_1", entry.ProviderID) + } +} + +// A provider body that fails mid-relay must not produce a usage entry: the +// client received an incomplete response and there is no trustworthy usage. +func TestMessages_NativeNonStreamingBodyErrorSkipsUsage(t *testing.T) { + provider := &mockProvider{ + supportedModels: []string{"claude-fable-5"}, + providerTypes: map[string]string{"claude-fable-5": "anthropic"}, + passthroughResponse: &core.PassthroughResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{"Content-Type": {"application/json"}}, + Body: io.NopCloser(io.MultiReader( + strings.NewReader(anthropicNonStreamingJSON[:40]), + &failingReader{}, + )), + }, + } + usageLogger := &collectingUsageLogger{config: usage.Config{Enabled: true}} + + e := echo.New() + handler := NewHandler(provider, nil, usageLogger, nil) + + reqBody := `{"model":"claude-fable-5","max_tokens":64,"messages":[{"role":"user","content":"Hi"}]}` + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(reqBody)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + + if err := handler.Messages(e.NewContext(req, rec)); err == nil { + t.Fatal("Messages: expected relay error, got nil") + } + if len(usageLogger.entries) != 0 { + t.Fatalf("usage entries = %d, want 0", len(usageLogger.entries)) + } +} + +type failingReader struct{} + +func (*failingReader) Read([]byte) (int, error) { + return 0, io.ErrUnexpectedEOF +} + type recordingFeedbackObserver struct { input, read, write int observed bool @@ -183,3 +281,75 @@ func TestMessages_NativeStreamingNotifiesFeedbackObservers(t *testing.T) { t.Errorf("cacheWriteInputTokens = %d, want 100", observer.write) } } + +// Extensions that requested response feedback must also hear about +// non-streaming native responses, with usage read from the response object. +func TestMessages_NativeNonStreamingNotifiesFeedbackObservers(t *testing.T) { + provider := &mockProvider{ + supportedModels: []string{"claude-fable-5"}, + providerTypes: map[string]string{"claude-fable-5": "anthropic"}, + passthroughResponse: &core.PassthroughResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(anthropicNonStreamingJSON)), + }, + } + + e := echo.New() + handler := NewHandler(provider, nil, nil, nil) + + reqBody := `{"model":"claude-fable-5","max_tokens":64,"messages":[{"role":"user","content":"Hi"}]}` + req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(reqBody)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + observer := &recordingFeedbackObserver{} + setResponseFeedbackObservers(c, []ext.ResponseFeedbackObserver{observer}) + + if err := handler.Messages(c); err != nil { + t.Fatalf("Messages: %v", err) + } + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, body = %s", rec.Code, rec.Body.String()) + } + if observer.calls != 1 { + t.Fatalf("ObserveResponse calls = %d, want 1", observer.calls) + } + if !observer.observed { + t.Error("usageObserved = false, want true") + } + if observer.input != 100 { + t.Errorf("inputTokens = %d, want 100", observer.input) + } + if observer.read != 9 { + t.Errorf("cachedInputTokens = %d, want 9", observer.read) + } + if observer.write != 7 { + t.Errorf("cacheWriteInputTokens = %d, want 7", observer.write) + } +} + +// The capture buffer must abandon oversized bodies without disturbing the +// relay: writes keep succeeding, and Captured reports nothing usable. +func TestCappedCaptureBufferOverflow(t *testing.T) { + capture := newCappedCaptureBuffer(8) + for range 3 { + n, err := capture.Write([]byte("abcde")) + if n != 5 || err != nil { + t.Fatalf("Write = (%d, %v), want (5, nil)", n, err) + } + } + if body, ok := capture.Captured(); ok { + t.Fatalf("Captured = (%q, true), want abandoned", body) + } + + capture = newCappedCaptureBuffer(8) + if _, err := capture.Write([]byte("abcde")); err != nil { + t.Fatalf("Write: %v", err) + } + body, ok := capture.Captured() + if !ok || string(body) != "abcde" { + t.Fatalf("Captured = (%q, %v), want (abcde, true)", body, ok) + } +} diff --git a/internal/server/passthrough_support.go b/internal/server/passthrough_support.go index d7258b2f3..5d741e7dc 100644 --- a/internal/server/passthrough_support.go +++ b/internal/server/passthrough_support.go @@ -1,13 +1,17 @@ package server import ( + "bytes" "context" "fmt" "io" + "mime" "net/http" "sort" "strings" + "github.com/goccy/go-json" + "github.com/labstack/echo/v5" "github.com/enterpilot/gomodel/internal/auditlog" @@ -317,14 +321,8 @@ func proxyPassthroughResponse(c *echo.Context, logger auditlog.LoggerInterface, observers = append(observers, observer) } } - if usageLogger != nil && usageLogger.Config().Enabled && (workflow == nil || workflow.UsageEnabled()) { - if observer := usage.NewStreamUsageObserver(usageLogger, model, providerType, requestID, usagePath, pricingResolver, core.UserPathFromContext(c.Request().Context())); observer != nil { - observer.SetProviderName(providerName) - observer.SetSessionID(core.SessionIDFromContext(c.Request().Context())) - observer.SetLabels(core.RequestLabelsFromContext(c.Request().Context())) - observer.SetRewriteTokensSaved(core.RewriteTokensSavedFromContext(c.Request().Context())) - observers = append(observers, observer) - } + if observer := passthroughUsageObserver(c, usageLogger, pricingResolver, workflow, model, providerType, providerName, requestID, usagePath); observer != nil { + observers = append(observers, observer) } observers = append(observers, extraObservers...) wrappedStream := streaming.NewObservedSSEStream(resp.Body, observers...) @@ -342,16 +340,168 @@ func proxyPassthroughResponse(c *echo.Context, logger auditlog.LoggerInterface, return nil } + // Non-streaming JSON responses carry usage inside the response object + // itself (top-level "usage" for Anthropic messages and OpenAI chat + // completions). Tee the relay into a bounded buffer and feed the complete + // body to the same observers as a single synthetic event, so usage + // accounting and response feedback match the SSE behavior. The audit + // stream observer is deliberately absent: non-streaming responses are + // audited by the regular audit middleware. + var observers []streaming.Observer + if isObservablePassthroughStatus(resp.StatusCode) { + observers = passthroughJSONResponseObservers(c, usageLogger, pricingResolver, providerType, providerName, endpoint, info, extraObservers) + } + if len(observers) == 0 || !isJSONContentType(resp.Headers) { + c.Response().WriteHeader(resp.StatusCode) + if _, err := io.Copy(c.Response(), resp.Body); err != nil { + return err + } + if f, ok := c.Response().(http.Flusher); ok { + f.Flush() + } + return nil + } + + capture := newCappedCaptureBuffer(maxObservedJSONResponseBytes) c.Response().WriteHeader(resp.StatusCode) - if _, err := io.Copy(c.Response(), resp.Body); err != nil { + if _, err := io.Copy(c.Response(), io.TeeReader(resp.Body, capture)); err != nil { + // The client received an incomplete body; do not account for it. return err } if f, ok := c.Response().(http.Flusher); ok { f.Flush() } + if body, ok := capture.Captured(); ok { + notifyObserversWithJSONBody(body, observers) + } return nil } +// maxObservedJSONResponseBytes caps how much of a non-streaming JSON response +// is buffered for usage extraction. Inference responses are bounded by +// max_tokens and stay far below this; anything larger (e.g. a passthrough +// file download with a JSON content type) skips observation rather than +// holding the body in memory. +const maxObservedJSONResponseBytes = 8 << 20 + +// cappedCaptureBuffer records writes up to a fixed cap. Once the cap is +// exceeded the capture is abandoned (Captured reports false) while writes +// keep succeeding, so the client relay is never affected. +type cappedCaptureBuffer struct { + buf bytes.Buffer + max int + overflow bool +} + +func newCappedCaptureBuffer(maxBytes int) *cappedCaptureBuffer { + return &cappedCaptureBuffer{max: maxBytes} +} + +func (b *cappedCaptureBuffer) Write(p []byte) (int, error) { + if !b.overflow { + if b.buf.Len()+len(p) > b.max { + b.overflow = true + b.buf.Reset() + } else { + b.buf.Write(p) + } + } + return len(p), nil +} + +func (b *cappedCaptureBuffer) Captured() ([]byte, bool) { + if b.overflow || b.buf.Len() == 0 { + return nil, false + } + return b.buf.Bytes(), true +} + +// passthroughUsageObserver builds the stream usage observer for a passthrough +// response when usage logging is enabled for the workflow, or nil. +func passthroughUsageObserver(c *echo.Context, usageLogger usage.LoggerInterface, pricingResolver usage.PricingResolver, workflow *core.Workflow, model, providerType, providerName, requestID, usagePath string) *usage.StreamUsageObserver { + if usageLogger == nil || !usageLogger.Config().Enabled || (workflow != nil && !workflow.UsageEnabled()) { + return nil + } + observer := usage.NewStreamUsageObserver(usageLogger, model, providerType, requestID, usagePath, pricingResolver, core.UserPathFromContext(c.Request().Context())) + if observer == nil { + return nil + } + observer.SetProviderName(providerName) + observer.SetSessionID(core.SessionIDFromContext(c.Request().Context())) + observer.SetLabels(core.RequestLabelsFromContext(c.Request().Context())) + observer.SetRewriteTokensSaved(core.RewriteTokensSavedFromContext(c.Request().Context())) + return observer +} + +// passthroughJSONResponseObservers assembles the observers interested in a +// completed non-streaming JSON passthrough response: the usage accounting +// observer plus any caller-supplied extras (response feedback). +func passthroughJSONResponseObservers(c *echo.Context, usageLogger usage.LoggerInterface, pricingResolver usage.PricingResolver, providerType, providerName, endpoint string, info *core.PassthroughRouteInfo, extraObservers []streaming.Observer) []streaming.Observer { + workflow := core.GetWorkflow(c.Request().Context()) + requestID := requestIDFromContextOrHeader(c.Request()) + usagePath := passthroughAuditPath(c, providerType, endpoint, info) + if requestPath := strings.TrimSpace(c.Request().URL.Path); requestPath != "" { + usagePath = requestPath + } + model := "" + if info != nil { + model = strings.TrimSpace(info.Model) + } + model = resolvedModelFromWorkflow(workflow, model) + + observers := make([]streaming.Observer, 0, 1+len(extraObservers)) + if observer := passthroughUsageObserver(c, usageLogger, pricingResolver, workflow, model, providerType, providerName, requestID, usagePath); observer != nil { + observers = append(observers, observer) + } + return append(observers, extraObservers...) +} + +// notifyObserversWithJSONBody replays a complete JSON response body to stream +// observers as one synthetic event followed by close, mirroring how the same +// payload would reach them as the final event of an SSE stream. Bodies that +// do not decode to a JSON object still close the observers out, so response +// feedback fires exactly once per response. +func notifyObserversWithJSONBody(body []byte, observers []streaming.Observer) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err == nil && payload != nil { + for _, observer := range observers { + if filter, ok := observer.(streaming.EventFilter); ok && !filter.WantsJSONEvent(body) { + continue + } + observer.OnJSONEvent(payload) + } + } + for _, observer := range observers { + observer.OnStreamClose() + } +} + +// isObservablePassthroughStatus reports whether a passthrough response status +// can carry a complete, accountable response body: any success status except +// 206 Partial Content, whose body is by definition incomplete. +func isObservablePassthroughStatus(status int) bool { + return status >= http.StatusOK && status < http.StatusMultipleChoices && status != http.StatusPartialContent +} + +func isJSONContentType(headers map[string][]string) bool { + for key, values := range headers { + if !strings.EqualFold(key, "Content-Type") { + continue + } + for _, value := range values { + mediaType, _, err := mime.ParseMediaType(value) + if err != nil { + continue + } + mediaType = strings.ToLower(mediaType) + if mediaType == "application/json" || strings.HasSuffix(mediaType, "+json") { + return true + } + } + } + return false +} + func passthroughErrorResponseHeaders(providerType string, statusCode int, src http.Header) http.Header { if providerType != "llmd" || statusCode != http.StatusTooManyRequests { return nil diff --git a/internal/server/passthrough_support_test.go b/internal/server/passthrough_support_test.go index cdc233a9e..53dd3aff0 100644 --- a/internal/server/passthrough_support_test.go +++ b/internal/server/passthrough_support_test.go @@ -2,11 +2,17 @@ package server import ( "context" + "io" "net/http" + "net/http/httptest" "slices" + "strings" "testing" + "github.com/labstack/echo/v5" + "github.com/enterpilot/gomodel/internal/core" + "github.com/enterpilot/gomodel/internal/usage" ) func TestBuildPassthroughHeadersSkipsConfiguredUserPathHeader(t *testing.T) { @@ -38,3 +44,115 @@ func TestDefaultEnabledPassthroughProvidersIncludesHetzner(t *testing.T) { t.Fatalf("defaultEnabledPassthroughProviders = %v, want hetzner included", defaultEnabledPassthroughProviders) } } + +// A successful non-streaming JSON passthrough response must produce a usage +// entry from its usage member — the same accounting SSE streams get from the +// stream usage observer. Covers the /p/{provider} surface directly. +func TestProxyPassthroughNonStreamingLogsUsage(t *testing.T) { + body := `{"id":"msg_p","type":"message","model":"claude-fable-5","usage":{"input_tokens":42,"output_tokens":6}}` + resp := &core.PassthroughResponse{ + StatusCode: http.StatusOK, + Headers: map[string][]string{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } + usageLogger := &collectingUsageLogger{config: usage.Config{Enabled: true}} + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/p/anthropic/messages", strings.NewReader(`{}`)) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + info := &core.PassthroughRouteInfo{Provider: "anthropic", RawEndpoint: "messages", Model: "claude-fable-5"} + if err := proxyPassthroughResponse(c, nil, usageLogger, nil, "anthropic", "anthropic", "messages", info, resp); err != nil { + t.Fatalf("proxyPassthroughResponse: %v", err) + } + if rec.Body.String() != body { + t.Fatalf("body not relayed verbatim: %s", rec.Body.String()) + } + if len(usageLogger.entries) != 1 { + t.Fatalf("usage entries = %d, want 1", len(usageLogger.entries)) + } + entry := usageLogger.entries[0] + if entry.InputTokens != 42 || entry.OutputTokens != 6 { + t.Errorf("tokens = (%d, %d), want (42, 6)", entry.InputTokens, entry.OutputTokens) + } + if entry.ProviderID != "msg_p" { + t.Errorf("ProviderID = %q, want msg_p", entry.ProviderID) + } +} + +// Any complete-body success status must be accounted, not only 200: /p/ is +// provider-generic and a 201/202 JSON response can carry usage too. +func TestProxyPassthroughNonStreamingLogsUsageForNon200Success(t *testing.T) { + for _, status := range []int{http.StatusCreated, http.StatusAccepted} { + body := `{"id":"msg_p","model":"claude-fable-5","usage":{"input_tokens":42,"output_tokens":6}}` + resp := &core.PassthroughResponse{ + StatusCode: status, + Headers: map[string][]string{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + } + usageLogger := &collectingUsageLogger{config: usage.Config{Enabled: true}} + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/p/anthropic/messages", strings.NewReader(`{}`)) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + info := &core.PassthroughRouteInfo{Provider: "anthropic", RawEndpoint: "messages"} + if err := proxyPassthroughResponse(c, nil, usageLogger, nil, "anthropic", "anthropic", "messages", info, resp); err != nil { + t.Fatalf("status %d: proxyPassthroughResponse: %v", status, err) + } + if len(usageLogger.entries) != 1 { + t.Fatalf("status %d: usage entries = %d, want 1", status, len(usageLogger.entries)) + } + if rec.Code != status { + t.Fatalf("status = %d, want %d", rec.Code, status) + } + } +} + +// Non-JSON media types and incomplete-body statuses must relay untouched with +// no usage entry: there is nothing trustworthy to account for. The media type +// is parsed, so JSON-adjacent types and parameters must not slip through. +func TestProxyPassthroughNonStreamingSkipsNonAccountableResponses(t *testing.T) { + cases := []struct { + name string + status int + contentType string + }{ + {name: "non-JSON content type", status: http.StatusOK, contentType: "application/x-jsonl"}, + {name: "JSON-adjacent media type", status: http.StatusOK, contentType: "application/json-seq"}, + {name: "JSON only in a parameter", status: http.StatusOK, contentType: `text/plain; profile="application/json"`}, + {name: "partial content", status: http.StatusPartialContent, contentType: "application/json"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + body := `{"usage":{"input_tokens":42,"output_tokens":6}}` + resp := &core.PassthroughResponse{ + StatusCode: tc.status, + Headers: map[string][]string{"Content-Type": {tc.contentType}}, + Body: io.NopCloser(strings.NewReader(body)), + } + usageLogger := &collectingUsageLogger{config: usage.Config{Enabled: true}} + + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/p/anthropic/messages", strings.NewReader(`{}`)) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + info := &core.PassthroughRouteInfo{Provider: "anthropic", RawEndpoint: "messages"} + if err := proxyPassthroughResponse(c, nil, usageLogger, nil, "anthropic", "anthropic", "messages", info, resp); err != nil { + t.Fatalf("proxyPassthroughResponse: %v", err) + } + if rec.Body.String() != body { + t.Fatalf("body not relayed verbatim: %s", rec.Body.String()) + } + if len(usageLogger.entries) != 0 { + t.Fatalf("usage entries = %d, want 0", len(usageLogger.entries)) + } + if rec.Code != tc.status { + t.Fatalf("status = %d, want %d", rec.Code, tc.status) + } + }) + } +}