From aecf9a94872d03f1b6edbc4484704e0254ce4186 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sun, 14 Jun 2026 00:15:36 +0530 Subject: [PATCH] feat(otel): OTLP/HTTP trace ingress to meter apps RiskKernel didn't proxy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the consume side of the OpenTelemetry surface: an OTLP/HTTP trace receiver at POST /v1/traces. Point any OTel exporter at the daemon (OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070) and the GenAI model calls an app already traces — OpenLLMetry, the OpenAI Agents SDK, the Vercel AI SDK — show up against governed runs with tokens and cost metered into the same ledger the proxy uses. This makes spend visible for agents that don't route through the proxy or SDK. The receiver accepts both OTLP encodings (protobuf, the default, and JSON) using the canonical OTLP proto types (promoted from an existing transitive dep — no new module), walks the spans, and meters each one carrying gen_ai.usage.* token counts. It correlates to a run by riskkernel.run.id (on the span, falling back to the resource), pricing via the existing table and recording through the ledger. A GenAI span with no run id is observed and reported as a rejected span in the OTLP partial-success response; spans without usage are ignored. Token counts emitted as a double or numeric string are accepted rather than dropped. Scope is observe + meter: a consumed call is recorded (and marks the run halted if it crosses the budget) but not blocked after the fact, since it already happened — governing consumed spans is a separate step. The receiver is off by default (RISKKERNEL_OTEL_INGRESS_ENABLED) and authenticated like the rest of the API, so an exporter carries the bearer token via OTEL_EXPORTER_OTLP_HEADERS. Tests: protobuf and JSON over HTTP (verbatim decode + metering + partial-success reply), run-id from span and from resource, pricing of a consumed call, lenient numeric attribute types, and skipping of non-GenAI and unattributed spans. Pinned attributes documented in api/v1/otel-genai.md; usage in docs/OTLP_INGRESS.md. --- CHANGELOG.md | 12 ++ api/v1/otel-genai.md | 27 +++- cmd/riskkernel/main.go | 2 +- docs/OTLP_INGRESS.md | 73 ++++++++++ go.mod | 4 +- internal/app/bootstrap.go | 9 ++ internal/config/config.go | 17 ++- internal/httpapi/memory_test.go | 2 +- internal/httpapi/server.go | 14 +- internal/httpapi/server_test.go | 2 +- internal/httpapi/slack_test.go | 2 +- internal/otel/ingress.go | 215 ++++++++++++++++++++++++++++ internal/otel/ingress_test.go | 242 ++++++++++++++++++++++++++++++++ 13 files changed, 604 insertions(+), 17 deletions(-) create mode 100644 docs/OTLP_INGRESS.md create mode 100644 internal/otel/ingress.go create mode 100644 internal/otel/ingress_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 4924658..c9a9367 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,18 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] ### Added +- **OTLP trace ingress.** RiskKernel can now act as an OTLP/HTTP trace endpoint + (`POST /v1/traces`), the consume side of the OpenTelemetry surface — point any + exporter at it (`OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070`) and the GenAI + model calls an app already traces (OpenLLMetry, the OpenAI Agents SDK, the Vercel + AI SDK) show up against governed runs with tokens and cost metered into the + ledger. Correlates by `riskkernel.run.id` (on the span or its resource); accepts + protobuf and JSON and replies with an OTLP `ExportTraceServiceResponse` (a span + without a run id is reported as a rejected partial-success span). Scope is observe + + meter — a consumed call is recorded (and marks the run halted if it crosses the + budget) but isn't blocked after the fact. Off by default; enable with + `RISKKERNEL_OTEL_INGRESS_ENABLED` and authenticated like the rest of the API. See + [`docs/OTLP_INGRESS.md`](docs/OTLP_INGRESS.md). - **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 diff --git a/api/v1/otel-genai.md b/api/v1/otel-genai.md index 6a7cf10..94ad9b9 100644 --- a/api/v1/otel-genai.md +++ b/api/v1/otel-genai.md @@ -69,7 +69,26 @@ conventions don't model. Names are stable per COMPATIBILITY.md. ## Consumption (ingress) -When acting as an OTLP endpoint, RiskKernel reads incoming `gen_ai.usage.*` to feed -the cost ledger and `gen_ai.request.model` / `gen_ai.system` to attribute spend, -correlating by `riskkernel.run.id` when present (or a configured trace→run mapping -otherwise). This lets RiskKernel govern apps it did not directly instrument. +RiskKernel can act as an OTLP/HTTP trace endpoint at **`POST /v1/traces`** — the +standard OTLP path, so any exporter targets it with one env var +(`OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070`). It accepts both protobuf +(`application/x-protobuf`, the OTLP default) and JSON (`application/json`) and +replies with an OTLP `ExportTraceServiceResponse` in the same encoding. + +For each span carrying token usage (`gen_ai.usage.input_tokens` / +`gen_ai.usage.output_tokens`), it reads `gen_ai.response.model` (falling back to +`gen_ai.request.model`) and `gen_ai.system` and meters the call's tokens and cost +into the ledger, correlating to a governed run by `riskkernel.run.id` (taken from +the span, falling back to the resource). A span without a run id is observed and +reported as a rejected span in the OTLP partial-success response, but not metered. +Spans without usage (tool calls, retrieval, framework spans) are ignored. + +This lets RiskKernel make spend visible for apps it did not directly proxy. Scope +is **observe + meter**: a consumed span records against the run's ledger (and marks +the run halted if its budget is crossed) but does not retroactively block a call +that already happened — governing consumed spans is a separate, future step. + +The receiver is **off by default**; no listener is mounted unless +`RISKKERNEL_OTEL_INGRESS_ENABLED` is set. It is authenticated like the rest of the +API, so an exporter must carry the bearer token via +`OTEL_EXPORTER_OTLP_HEADERS=authorization=Bearer ` when a token is set. diff --git a/cmd/riskkernel/main.go b/cmd/riskkernel/main.go index 4296434..5bc3ce3 100644 --- a/cmd/riskkernel/main.go +++ b/cmd/riskkernel/main.go @@ -140,7 +140,7 @@ func runServe(_ []string) error { ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer stop() - srv := httpapi.New(cfg, deps.Gateway, deps.Runs, deps.Approvals, deps.Slack, deps.MCP, deps.Memory, deps.Log) + srv := httpapi.New(cfg, deps.Gateway, deps.Runs, deps.Approvals, deps.Slack, deps.MCP, deps.Memory, deps.Ingress, deps.Log) addr := fmt.Sprintf(":%d", cfg.Port) return srv.Serve(ctx, addr) } diff --git a/docs/OTLP_INGRESS.md b/docs/OTLP_INGRESS.md new file mode 100644 index 0000000..562b0e7 --- /dev/null +++ b/docs/OTLP_INGRESS.md @@ -0,0 +1,73 @@ +# OTLP trace ingress + +RiskKernel emits OpenTelemetry GenAI spans to your observability backend (the +export side of Surface 3). The other half is **ingress**: RiskKernel can also *be* +an OTLP endpoint that consumes GenAI spans from apps already instrumented with +OpenLLMetry, the OpenAI Agents SDK, or the Vercel AI SDK — so it can make spend +visible for agents it never directly proxied. + +Point an existing app's OTLP exporter at RiskKernel and the model calls it already +traces show up against governed runs, with tokens and cost metered into the same +ledger the proxy uses. + +## Enabling it + +The receiver is **off by default** — no listener is mounted unless you turn it on: + +```bash +RISKKERNEL_OTEL_INGRESS_ENABLED=true riskkernel serve +``` + +On the app side, point any OpenTelemetry exporter at the daemon. The endpoint is +the standard OTLP/HTTP traces path, so this is the usual one env var: + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070 +# when RISKKERNEL_API_TOKEN is set, the exporter must carry it: +export OTEL_EXPORTER_OTLP_HEADERS=authorization=Bearer $RISKKERNEL_API_TOKEN +``` + +The exporter appends `/v1/traces` itself. Both OTLP encodings work: protobuf +(`application/x-protobuf`, the SDK default) and JSON (`application/json`). + +## Correlating spans to a run + +Set `riskkernel.run.id` on your spans (or once on the OTel resource) so consumed +calls are attributed to a governed run: + +```python +span.set_attribute("riskkernel.run.id", run_id) +``` + +The run is created lazily under the default budget if it doesn't exist yet, exactly +like the proxy's run-id header. A GenAI usage span with **no** run id is observed +and reported back as a rejected span in the OTLP partial-success response, but is +not metered. + +## What gets metered + +For each span carrying token usage, RiskKernel reads the pinned GenAI attributes +(see [`api/v1/otel-genai.md`](../api/v1/otel-genai.md)): + +| Attribute | Used for | +|---|---| +| `gen_ai.usage.input_tokens` / `gen_ai.usage.output_tokens` | Token counts and cost (priced via the same table the proxy uses). | +| `gen_ai.response.model` (falls back to `gen_ai.request.model`) | The model, for pricing and attribution. | +| `gen_ai.system` | The provider, recorded on the ledger entry. | +| `riskkernel.run.id` | The run to attribute the call to. | + +Token counts emitted as a double or numeric string (some instrumenters do this) +are accepted rather than dropped to zero. Spans without usage — tool calls, +retrieval, framework spans — are ignored. + +## Scope: observe + meter + +Ingested calls already happened in the other app, so the budget can't block them +before the fact. RiskKernel **observes and meters** them: the call is recorded +against the run's ledger, and if the recorded usage crosses the run's budget the +run is marked halted (visible in `riskkernel runs list`, the audit export, and +`GET /v1/runs/{id}`). Actively governing consumed spans — gating or refusing the +*next* call on a run that an external app is driving — is a separate, future step. + +For deterministic, before-the-fact enforcement, route calls through the proxy +(Surface 1) or the SDK (Surface 2) instead. diff --git a/go.mod b/go.mod index 83600ac..3676596 100644 --- a/go.mod +++ b/go.mod @@ -10,6 +10,8 @@ require ( go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.44.0 go.opentelemetry.io/otel/sdk v1.44.0 go.opentelemetry.io/otel/trace v1.44.0 + go.opentelemetry.io/proto/otlp v1.10.0 + google.golang.org/protobuf v1.36.11 gopkg.in/yaml.v3 v3.0.1 modernc.org/sqlite v1.52.0 ) @@ -29,7 +31,6 @@ require ( github.com/sethvargo/go-retry v0.3.0 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect go.opentelemetry.io/otel/metric v1.44.0 // indirect - go.opentelemetry.io/proto/otlp v1.10.0 // indirect go.uber.org/multierr v1.11.0 // indirect golang.org/x/net v0.55.0 // indirect golang.org/x/sync v0.20.0 // indirect @@ -38,7 +39,6 @@ require ( google.golang.org/genproto/googleapis/api v0.0.0-20260526163538-3dc84a4a5aaa // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260526163538-3dc84a4a5aaa // indirect google.golang.org/grpc v1.81.1 // indirect - google.golang.org/protobuf v1.36.11 // indirect modernc.org/libc v1.72.3 // indirect modernc.org/mathutil v1.7.1 // indirect modernc.org/memory v1.11.0 // indirect diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 4e92dad..b241287 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -40,6 +40,7 @@ type Deps struct { Slack *approval.SlackNotifier // nil when the Slack channel isn't configured MCP *mcp.Gateway // nil when no upstream is configured Memory *memory.Reader + Ingress *otel.Ingress // nil unless the OTLP trace receiver is enabled } // Close releases dependencies that hold resources (the tracer's buffered spans, @@ -147,6 +148,13 @@ func Build(cfg *config.Config) (*Deps, error) { log.Warn("RISKKERNEL_MEMORY_EMBEDDINGS is set but embeddings are not implemented in v0.1; using deterministic keyword search") } + // OTLP trace receiver (Surface 3, consume side) — off unless explicitly enabled. + var ingress *otel.Ingress + if cfg.OTel.IngressEnabled { + ingress = otel.NewIngress(mgr, prices, log) + log.Info("otlp trace ingress enabled", "path", "POST /v1/traces") + } + return &Deps{ Config: cfg, Log: log, @@ -160,6 +168,7 @@ func Build(cfg *config.Config) (*Deps, error) { Slack: slackNotifier, MCP: mcpGW, Memory: memReader, + Ingress: ingress, }, nil } diff --git a/internal/config/config.go b/internal/config/config.go index ec2c911..8b3b1a5 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -165,6 +165,12 @@ type OTelConfig struct { // OTEL_EXPORTER_OTLP_HEADERS, as a comma-separated list of key=value pairs. // Carries secrets — never logged. Headers map[string]string + + // IngressEnabled turns on the OTLP/HTTP trace receiver (Surface 3, consume + // side): RiskKernel becomes an OTLP endpoint at POST /v1/traces that meters + // GenAI spans from apps it didn't directly proxy. Off by default — no listener + // is mounted unless RISKKERNEL_OTEL_INGRESS_ENABLED is set truthy. + IngressEnabled bool } // BudgetConfig holds raw budget values (no governor dependency here so config @@ -292,11 +298,12 @@ func loadOTel() OTelConfig { headers = os.Getenv("OTEL_EXPORTER_OTLP_HEADERS") } return OTelConfig{ - Endpoint: endpoint, - Protocol: getenvDefault("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"), - Insecure: insecure, - ServiceName: getenvDefault("OTEL_SERVICE_NAME", "riskkernel"), - Headers: parseOTLPHeaders(headers), + Endpoint: endpoint, + Protocol: getenvDefault("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"), + Insecure: insecure, + ServiceName: getenvDefault("OTEL_SERVICE_NAME", "riskkernel"), + Headers: parseOTLPHeaders(headers), + IngressEnabled: envBoolDefault("RISKKERNEL_OTEL_INGRESS_ENABLED", false), } } diff --git a/internal/httpapi/memory_test.go b/internal/httpapi/memory_test.go index 959182f..321d1f3 100644 --- a/internal/httpapi/memory_test.go +++ b/internal/httpapi/memory_test.go @@ -34,7 +34,7 @@ func newMemoryServer(t *testing.T) http.Handler { t.Cleanup(func() { _ = store.Close() }) log := slog.New(slog.NewTextHandler(io.Discard, nil)) mgr := runs.NewManager(governor.Budget{}).WithStore(store, log) - srv := New(&config.Config{}, nil, mgr, nil, nil, nil, memory.NewReader(memDir), log) + srv := New(&config.Config{}, nil, mgr, nil, nil, nil, memory.NewReader(memDir), nil, log) return srv.Handler() } diff --git a/internal/httpapi/server.go b/internal/httpapi/server.go index e6ddbe6..721ebab 100644 --- a/internal/httpapi/server.go +++ b/internal/httpapi/server.go @@ -19,6 +19,7 @@ import ( "github.com/prashar32/riskkernel/internal/httpx" "github.com/prashar32/riskkernel/internal/mcp" "github.com/prashar32/riskkernel/internal/memory" + "github.com/prashar32/riskkernel/internal/otel" "github.com/prashar32/riskkernel/internal/runs" "github.com/prashar32/riskkernel/internal/storage" "github.com/prashar32/riskkernel/internal/version" @@ -33,13 +34,14 @@ type Server struct { slack *approval.SlackNotifier // nil when the Slack channel isn't configured mcp *mcp.Gateway memory *memory.Reader + ingress *otel.Ingress // nil unless the OTLP trace receiver is enabled log *slog.Logger } // New constructs a Server. func New(cfg *config.Config, gw *gateway.Gateway, mgr *runs.Manager, gate *approval.Gate, - slack *approval.SlackNotifier, mcpGW *mcp.Gateway, mem *memory.Reader, log *slog.Logger) *Server { - return &Server{cfg: cfg, gateway: gw, runs: mgr, approvals: gate, slack: slack, mcp: mcpGW, memory: mem, log: log} + slack *approval.SlackNotifier, mcpGW *mcp.Gateway, mem *memory.Reader, ingress *otel.Ingress, log *slog.Logger) *Server { + return &Server{cfg: cfg, gateway: gw, runs: mgr, approvals: gate, slack: slack, mcp: mcpGW, memory: mem, ingress: ingress, log: log} } // Handler returns the root HTTP handler with all routes mounted. @@ -107,6 +109,14 @@ func (s *Server) Handler() http.Handler { mux.HandleFunc("GET /metrics", s.requireAuth(s.handleMetrics)) } + // Surface 3 (consume side) — the OTLP/HTTP trace receiver. Off unless enabled; + // authenticated like the rest of the API (OTLP exporters carry the bearer token + // via OTEL_EXPORTER_OTLP_HEADERS). Standard OTLP path so any exporter can target + // it with one env var. + if s.ingress != nil { + mux.HandleFunc("POST /v1/traces", s.requireAuth(s.ingress.HandleTraces)) + } + return s.recoverer(mux) } diff --git a/internal/httpapi/server_test.go b/internal/httpapi/server_test.go index 0735864..ac03fa7 100644 --- a/internal/httpapi/server_test.go +++ b/internal/httpapi/server_test.go @@ -30,7 +30,7 @@ func newTestServer(t *testing.T, token string) (*Server, *runs.Manager, *approva log := slog.New(slog.NewTextHandler(io.Discard, nil)) mgr := runs.NewManager(governor.Budget{Tokens: 100000}).WithStore(store, log) gate := approval.NewGate(store, approval.Policy{DefaultSafe: true}, nil, log) - srv := New(&config.Config{APIToken: token}, nil, mgr, gate, nil, nil, memory.NewReader(t.TempDir()), log) + srv := New(&config.Config{APIToken: token}, nil, mgr, gate, nil, nil, memory.NewReader(t.TempDir()), nil, log) return srv, mgr, gate } diff --git a/internal/httpapi/slack_test.go b/internal/httpapi/slack_test.go index 354df9f..e1d17c7 100644 --- a/internal/httpapi/slack_test.go +++ b/internal/httpapi/slack_test.go @@ -40,7 +40,7 @@ func newSlackTestServer(t *testing.T, secret string) (*Server, *runs.Manager, *a mgr := runs.NewManager(governor.Budget{Tokens: 100000}).WithStore(store, log) gate := approval.NewGate(store, approval.Policy{DefaultSafe: true}, nil, log) slack := approval.NewSlackNotifier("xoxb-test", "C123", secret, log) - srv := New(&config.Config{}, nil, mgr, gate, slack, nil, memory.NewReader(t.TempDir()), log) + srv := New(&config.Config{}, nil, mgr, gate, slack, nil, memory.NewReader(t.TempDir()), nil, log) return srv, mgr, gate } diff --git a/internal/otel/ingress.go b/internal/otel/ingress.go new file mode 100644 index 0000000..11ef171 --- /dev/null +++ b/internal/otel/ingress.go @@ -0,0 +1,215 @@ +package otel + +import ( + "io" + "log/slog" + "net/http" + "strconv" + "strings" + + coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1" + commonpb "go.opentelemetry.io/proto/otlp/common/v1" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" + + "github.com/prashar32/riskkernel/internal/pricing" + "github.com/prashar32/riskkernel/internal/runs" +) + +// maxTraceBodyBytes caps a single OTLP export payload (a generous batch of spans). +// The ingress reads the same pinned attribute keys the egress emits (attrRunID, +// attrGenAI*) — defined in otel.go — so consume and emit stay in lockstep. +const maxTraceBodyBytes = 8 << 20 // 8 MiB + +// Ingress is the OTLP/HTTP trace receiver — the consume side of Surface 3. It +// accepts GenAI spans from apps already instrumented (OpenLLMetry, the OpenAI +// Agents SDK, the Vercel AI SDK), correlates each model-call span to a governed +// run by riskkernel.run.id, and meters its token usage and cost into the ledger. +// This lets RiskKernel make spend visible for apps it never directly proxied. +// +// Scope is observe + meter: a consumed span records against the run's ledger but +// does not retroactively block a call that already happened (governing consumed +// spans is a separate, future step). The receiver is off by default and mounted +// only when RISKKERNEL_OTEL_INGRESS_ENABLED is set. +type Ingress struct { + runs *runs.Manager + prices *pricing.Table + log *slog.Logger +} + +// NewIngress constructs the OTLP trace receiver. mgr and prices must be non-nil. +func NewIngress(mgr *runs.Manager, prices *pricing.Table, log *slog.Logger) *Ingress { + return &Ingress{runs: mgr, prices: prices, log: log} +} + +// ingestResult tallies one export for the OTLP partial-success reply and logging. +type ingestResult struct { + metered int // GenAI call spans metered against a run + unattributed int // GenAI call spans with no riskkernel.run.id (observed, not metered) +} + +// HandleTraces implements POST /v1/traces, the standard OTLP/HTTP traces path — +// point any OTLP exporter (OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:7070) at +// it. It accepts protobuf (application/x-protobuf, the OTLP default) or JSON +// (application/json), meters the GenAI model-call spans it recognizes, and replies +// with an OTLP ExportTraceServiceResponse in the request's encoding. +func (in *Ingress) HandleTraces(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, maxTraceBodyBytes)) + if err != nil { + http.Error(w, "could not read request body", http.StatusBadRequest) + return + } + + isJSON := strings.HasPrefix(r.Header.Get("Content-Type"), "application/json") + var req coltracepb.ExportTraceServiceRequest + if isJSON { + if err := protojson.Unmarshal(body, &req); err != nil { + http.Error(w, "invalid OTLP JSON: "+err.Error(), http.StatusBadRequest) + return + } + } else { + // An unset or application/x-protobuf content type decodes as protobuf, the + // OTLP default. + if err := proto.Unmarshal(body, &req); err != nil { + http.Error(w, "invalid OTLP protobuf: "+err.Error(), http.StatusBadRequest) + return + } + } + + res := in.ingest(&req) + if in.log != nil && (res.metered > 0 || res.unattributed > 0) { + in.log.Debug("otlp trace ingress", "metered", res.metered, "unattributed", res.unattributed) + } + + // OTLP partial success: spans we couldn't attribute to a run are still a 200 + // (partial success is not an error), but reported as rejected so the caller can + // see they were observed and not metered. + resp := &coltracepb.ExportTraceServiceResponse{} + if res.unattributed > 0 { + resp.PartialSuccess = &coltracepb.ExportTracePartialSuccess{ + RejectedSpans: int64(res.unattributed), + ErrorMessage: "GenAI spans without riskkernel.run.id were observed but not metered to a run", + } + } + writeOTLPResponse(w, isJSON, resp) +} + +// ingest walks the export and meters each recognized GenAI model-call span against +// its run. The run id is taken from the span, falling back to its resource (some +// instrumenters set riskkernel.run.id once on the resource). +func (in *Ingress) ingest(req *coltracepb.ExportTraceServiceRequest) ingestResult { + var res ingestResult + for _, rs := range req.GetResourceSpans() { + resAttrs := otlpAttrMap(rs.GetResource().GetAttributes()) + for _, ss := range rs.GetScopeSpans() { + for _, span := range ss.GetSpans() { + attrs := otlpAttrMap(span.GetAttributes()) + // A GenAI model call is identified by carrying token usage; spans + // without usage (tool calls, retrieval, framework spans) are ignored + // for metering. + if _, hasIn := attrs[attrGenAIInputTokens]; !hasIn { + if _, hasOut := attrs[attrGenAIOutputTokens]; !hasOut { + continue + } + } + runID := otlpStringAttr(attrs, attrRunID) + if runID == "" { + runID = otlpStringAttr(resAttrs, attrRunID) + } + if runID == "" { + res.unattributed++ + continue + } + in.meter(runID, attrs) + res.metered++ + } + } + } + return res +} + +// meter records one consumed model call against its run's ledger. The run is +// created lazily under the default budget if unknown — mirroring how the proxy +// resolves a run from the run-id header — so spend shows up against a run either +// way. A budget crossed by the recorded usage marks the run halted (visible in +// runs list / audit); it cannot un-spend a call that already happened. +func (in *Ingress) meter(runID string, attrs map[string]*commonpb.AnyValue) { + model := otlpStringAttr(attrs, attrGenAIResponseModel) + if model == "" { + model = otlpStringAttr(attrs, attrGenAIRequestModel) + } + inTok := otlpIntAttr(attrs, attrGenAIInputTokens) + outTok := otlpIntAttr(attrs, attrGenAIOutputTokens) + cost, priced := in.prices.Cost(model, inTok, outTok) + + run := in.runs.GetOrCreate(runID) + _ = run.RecordCall(runs.Call{ + Provider: otlpStringAttr(attrs, attrGenAISystem), + Model: model, + PromptTokens: inTok, + CompletionTokens: outTok, + Dollars: cost, + Priced: priced, + ResponseID: otlpStringAttr(attrs, attrGenAIResponseID), + }) +} + +// writeOTLPResponse replies with an ExportTraceServiceResponse in the request's +// encoding, as the OTLP/HTTP spec requires. +func writeOTLPResponse(w http.ResponseWriter, isJSON bool, resp *coltracepb.ExportTraceServiceResponse) { + var ( + out []byte + err error + ct string + ) + if isJSON { + out, err = protojson.Marshal(resp) + ct = "application/json" + } else { + out, err = proto.Marshal(resp) + ct = "application/x-protobuf" + } + if err != nil { + http.Error(w, "encoding OTLP response", http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", ct) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(out) +} + +// otlpAttrMap indexes a span's (or resource's) attributes by key for lookup. +func otlpAttrMap(kvs []*commonpb.KeyValue) map[string]*commonpb.AnyValue { + m := make(map[string]*commonpb.AnyValue, len(kvs)) + for _, kv := range kvs { + if kv != nil { + m[kv.GetKey()] = kv.GetValue() + } + } + return m +} + +// otlpStringAttr returns a string attribute, or "" if absent or not a string. +func otlpStringAttr(m map[string]*commonpb.AnyValue, key string) string { + if v, ok := m[key].GetValue().(*commonpb.AnyValue_StringValue); ok { + return v.StringValue + } + return "" +} + +// otlpIntAttr returns an integer attribute. Token counts are ints per the convention, +// but some instrumenters emit them as doubles or numeric strings, so accept those +// too rather than silently dropping the usage to zero. +func otlpIntAttr(m map[string]*commonpb.AnyValue, key string) int64 { + switch v := m[key].GetValue().(type) { + case *commonpb.AnyValue_IntValue: + return v.IntValue + case *commonpb.AnyValue_DoubleValue: + return int64(v.DoubleValue) + case *commonpb.AnyValue_StringValue: + n, _ := strconv.ParseInt(strings.TrimSpace(v.StringValue), 10, 64) + return n + default: + return 0 + } +} diff --git a/internal/otel/ingress_test.go b/internal/otel/ingress_test.go new file mode 100644 index 0000000..cc4c7e6 --- /dev/null +++ b/internal/otel/ingress_test.go @@ -0,0 +1,242 @@ +package otel + +import ( + "bytes" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "testing" + + coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1" + commonpb "go.opentelemetry.io/proto/otlp/common/v1" + resourcepb "go.opentelemetry.io/proto/otlp/resource/v1" + tracepb "go.opentelemetry.io/proto/otlp/trace/v1" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/proto" + + "github.com/prashar32/riskkernel/internal/governor" + "github.com/prashar32/riskkernel/internal/pricing" + "github.com/prashar32/riskkernel/internal/runs" +) + +// --- helpers to build OTLP spans --- + +func kvStr(k, v string) *commonpb.KeyValue { + return &commonpb.KeyValue{Key: k, Value: &commonpb.AnyValue{Value: &commonpb.AnyValue_StringValue{StringValue: v}}} +} +func kvInt(k string, v int64) *commonpb.KeyValue { + return &commonpb.KeyValue{Key: k, Value: &commonpb.AnyValue{Value: &commonpb.AnyValue_IntValue{IntValue: v}}} +} +func kvDouble(k string, v float64) *commonpb.KeyValue { + return &commonpb.KeyValue{Key: k, Value: &commonpb.AnyValue{Value: &commonpb.AnyValue_DoubleValue{DoubleValue: v}}} +} + +// export wraps spans into an ExportTraceServiceRequest, optionally with +// resource-level attributes. +func export(resAttrs []*commonpb.KeyValue, spans ...*tracepb.Span) *coltracepb.ExportTraceServiceRequest { + var res *resourcepb.Resource + if resAttrs != nil { + res = &resourcepb.Resource{Attributes: resAttrs} + } + return &coltracepb.ExportTraceServiceRequest{ + ResourceSpans: []*tracepb.ResourceSpans{{ + Resource: res, + ScopeSpans: []*tracepb.ScopeSpans{{Spans: spans}}, + }}, + } +} + +func newIngress(t *testing.T, budget governor.Budget, prices *pricing.Table) (*Ingress, *runs.Manager) { + t.Helper() + mgr := runs.NewManager(budget) + return NewIngress(mgr, prices, slog.New(slog.NewTextHandler(io.Discard, nil))), mgr +} + +// --- tests --- + +func TestIngress_MetersGenAISpanToRun(t *testing.T) { + in, mgr := newIngress(t, governor.Budget{Tokens: 100000}, pricing.NewTable(nil)) + + span := &tracepb.Span{ + Name: "chat claude-sonnet-4-5", + Attributes: []*commonpb.KeyValue{ + kvStr(attrRunID, "run-1"), + kvStr(attrGenAISystem, "anthropic"), + kvStr(attrGenAIResponseModel, "claude-sonnet-4-5-20250101"), + kvInt(attrGenAIInputTokens, 100), + kvInt(attrGenAIOutputTokens, 50), + }, + } + res := in.ingest(export(nil, span)) + if res.metered != 1 || res.unattributed != 0 { + t.Fatalf("ingest result = %+v, want metered=1", res) + } + + run, ok := mgr.Get("run-1") + if !ok { + t.Fatal("run-1 was not created/metered") + } + v := run.View() + if v.Usage.PromptTokens != 100 || v.Usage.CompletionTokens != 50 { + t.Fatalf("usage = %+v, want 100/50", v.Usage) + } + if v.Usage.Tokens() != 150 { + t.Fatalf("total tokens = %d, want 150", v.Usage.Tokens()) + } +} + +func TestIngress_PricesConsumedCall(t *testing.T) { + // A model with a known rate must produce a non-zero, priced cost in the ledger. + prices := pricing.NewTable(map[string]pricing.Rate{ + "gpt-4o": {InputPerM: 5, OutputPerM: 15}, // $ per 1M tokens + }) + in, mgr := newIngress(t, governor.Budget{}, prices) + + span := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrRunID, "r"), + kvStr(attrGenAISystem, "openai"), + kvStr(attrGenAIRequestModel, "gpt-4o"), + kvInt(attrGenAIInputTokens, 1_000_000), + kvInt(attrGenAIOutputTokens, 0), + }} + in.ingest(export(nil, span)) + + run, _ := mgr.Get("r") + if got := run.View().Usage.Dollars; got != 5.0 { + t.Fatalf("metered dollars = %v, want 5.0", got) + } +} + +func TestIngress_RunIDFromResourceAttributes(t *testing.T) { + // Some instrumenters set riskkernel.run.id once on the resource, not per span. + in, mgr := newIngress(t, governor.Budget{}, pricing.NewTable(nil)) + span := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrGenAIRequestModel, "claude-sonnet-4-5"), + kvInt(attrGenAIInputTokens, 10), + kvInt(attrGenAIOutputTokens, 5), + }} + res := in.ingest(export([]*commonpb.KeyValue{kvStr(attrRunID, "res-run")}, span)) + if res.metered != 1 { + t.Fatalf("metered = %d, want 1 (run id from resource)", res.metered) + } + if _, ok := mgr.Get("res-run"); !ok { + t.Fatal("run from resource attribute was not metered") + } +} + +func TestIngress_SkipsNonGenAISpansAndUnattributed(t *testing.T) { + in, mgr := newIngress(t, governor.Budget{}, pricing.NewTable(nil)) + + // A non-GenAI span (no usage) — ignored entirely. + toolSpan := &tracepb.Span{Name: "execute_tool write_file", Attributes: []*commonpb.KeyValue{ + kvStr(attrRunID, "r"), kvStr("gen_ai.tool.name", "write_file"), + }} + // A GenAI span with usage but NO run id — observed but not metered. + orphan := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrGenAIRequestModel, "gpt-4o"), + kvInt(attrGenAIInputTokens, 10), kvInt(attrGenAIOutputTokens, 2), + }} + res := in.ingest(export(nil, toolSpan, orphan)) + if res.metered != 0 { + t.Fatalf("metered = %d, want 0", res.metered) + } + if res.unattributed != 1 { + t.Fatalf("unattributed = %d, want 1", res.unattributed) + } + if _, ok := mgr.Get("r"); ok { + t.Fatal("the tool span must not create a run") + } +} + +func TestIngress_LenientNumericAttrTypes(t *testing.T) { + // Token counts emitted as a double or a numeric string must still be metered. + in, mgr := newIngress(t, governor.Budget{}, pricing.NewTable(nil)) + span := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrRunID, "r"), + kvStr(attrGenAIRequestModel, "m"), + kvDouble(attrGenAIInputTokens, 42), + kvStr(attrGenAIOutputTokens, "8"), + }} + in.ingest(export(nil, span)) + v, _ := mgr.Get("r") + if u := v.View().Usage; u.PromptTokens != 42 || u.CompletionTokens != 8 { + t.Fatalf("usage = %+v, want 42/8 (double + string parsed)", u) + } +} + +func TestIngress_HandleTraces_Protobuf(t *testing.T) { + in, mgr := newIngress(t, governor.Budget{Tokens: 1000}, pricing.NewTable(nil)) + span := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrRunID, "http-run"), + kvStr(attrGenAIResponseModel, "claude"), + kvInt(attrGenAIInputTokens, 20), kvInt(attrGenAIOutputTokens, 5), + }} + body, err := proto.Marshal(export(nil, span)) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest(http.MethodPost, "/v1/traces", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/x-protobuf") + w := httptest.NewRecorder() + in.HandleTraces(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 != "application/x-protobuf" { + t.Errorf("response content-type = %q", ct) + } + var resp coltracepb.ExportTraceServiceResponse + if err := proto.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("response is not a valid OTLP protobuf: %v", err) + } + if resp.GetPartialSuccess().GetRejectedSpans() != 0 { + t.Errorf("unexpected rejected spans: %d", resp.GetPartialSuccess().GetRejectedSpans()) + } + if _, ok := mgr.Get("http-run"); !ok { + t.Fatal("span POSTed over HTTP was not metered") + } +} + +func TestIngress_HandleTraces_JSONPartialSuccess(t *testing.T) { + in, _ := newIngress(t, governor.Budget{}, pricing.NewTable(nil)) + // A GenAI usage span with no run id → reported as a rejected (unattributed) span. + orphan := &tracepb.Span{Attributes: []*commonpb.KeyValue{ + kvStr(attrGenAIRequestModel, "m"), + kvInt(attrGenAIInputTokens, 1), kvInt(attrGenAIOutputTokens, 1), + }} + body, err := protojson.Marshal(export(nil, orphan)) + if err != nil { + t.Fatal(err) + } + r := httptest.NewRequest(http.MethodPost, "/v1/traces", bytes.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + in.HandleTraces(w, r) + + if w.Code != http.StatusOK { + t.Fatalf("status = %d", w.Code) + } + if ct := w.Header().Get("Content-Type"); ct != "application/json" { + t.Errorf("response content-type = %q", ct) + } + var resp coltracepb.ExportTraceServiceResponse + if err := protojson.Unmarshal(w.Body.Bytes(), &resp); err != nil { + t.Fatalf("response is not valid OTLP JSON: %v", err) + } + if resp.GetPartialSuccess().GetRejectedSpans() != 1 { + t.Errorf("rejected spans = %d, want 1", resp.GetPartialSuccess().GetRejectedSpans()) + } +} + +func TestIngress_HandleTraces_BadBody(t *testing.T) { + in, _ := newIngress(t, governor.Budget{}, pricing.NewTable(nil)) + r := httptest.NewRequest(http.MethodPost, "/v1/traces", bytes.NewReader([]byte("not protobuf"))) + r.Header.Set("Content-Type", "application/x-protobuf") + w := httptest.NewRecorder() + in.HandleTraces(w, r) + if w.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400", w.Code) + } +}