diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c83ab6..1c43cde 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). - **CrewAI adapter (Python SDK).** `from riskkernel.adapters.crewai import RiskKernelStepCallback` — a `step_callback` you wire onto a CrewAI `Agent` or the whole `Crew` to bind it to a governed run with no other code change. One agent step 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) + } +}