diff --git a/CHANGELOG.md b/CHANGELOG.md index 8872493..a92b987 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,15 @@ surface is governed by [`COMPATIBILITY.md`](COMPATIBILITY.md). ## [Unreleased] +### Added +- **Tool governance shows up in your traces.** Every governed MCP `tools/call` now + emits an OpenTelemetry span (`execute_tool {tool}`) alongside the model-call spans, + carrying `gen_ai.tool.name`, `riskkernel.tool.side_effect`, and + `riskkernel.tool.status` (`approved`, `blocked`, `denied`, or `timeout`). Allowlist + blocks and approval denials are now visible in whatever OTLP backend you already + run — a refused call is marked with an error span status so it stands out. See + [`api/v1/otel-genai.md`](api/v1/otel-genai.md) and [`examples/otel`](examples/otel). + ## [0.3.0] - 2026-06-06 The crash-resume moat, proven and polished: a real `kill -9` → resume demo, the diff --git a/api/v1/otel-genai.md b/api/v1/otel-genai.md index bc74a33..6a7cf10 100644 --- a/api/v1/otel-genai.md +++ b/api/v1/otel-genai.md @@ -22,6 +22,12 @@ One span per **model call** (`gen_ai.client` kind), nested under one span per **step**, nested under one span per **run**. Span names follow the convention `{gen_ai.operation.name} {gen_ai.request.model}`, e.g. `chat claude-sonnet-4-5`. +One span per **governed MCP tool call** (`gen_ai.operation.name` = `execute_tool`), +named `execute_tool {tool}`, e.g. `execute_tool write_file`. It carries the +governance outcome in `riskkernel.tool.status`, so allowlist blocks and approval +denials are visible alongside model calls — a refused call is marked with an error +span status. + ## Emitted attributes (`gen_ai.*` — standard) | Attribute | Type | Example | Notes | @@ -36,6 +42,7 @@ One span per **model call** (`gen_ai.client` kind), nested under one span per | `gen_ai.usage.output_tokens` | int | `134` | Completion tokens. | | `gen_ai.response.finish_reasons` | string[] | `["stop"]` | | | `gen_ai.response.id` | string | `msg_01...` | Provider response id. | +| `gen_ai.tool.name` | string | `write_file` | On tool-call spans: the tool invoked. | | `error.type` | string | `provider_error` | On failure (standard OTel). | > Prompt/response **content** is NOT emitted by default (privacy + no telemetry @@ -57,10 +64,8 @@ conventions don't model. Names are stable per COMPATIBILITY.md. | `riskkernel.budget.dollars.limit` | double | `5.00` | | | `riskkernel.budget.dollars.remaining` | double | `4.81` | | | `riskkernel.halt.reason` | string | `token_budget_exceeded` | Set on the span where the governor halted the run (see HaltReason in openapi.yaml). | -| `riskkernel.approval.required` | bool | `true` | Side-effecting call gated for HITL. | -| `riskkernel.approval.decision` | string | `approve` | Once resolved. | -| `riskkernel.tool.name` | string | `mcp://shell` | For tool-call spans. | -| `riskkernel.tool.side_effect` | string | `write` | Classified side effect. | +| `riskkernel.tool.side_effect` | string | `write` | On tool-call spans: the classified side effect (empty = read-only). | +| `riskkernel.tool.status` | string | `blocked` | On tool-call spans: `approved`, `blocked` (allowlist), `denied` (approval), or `timeout`. | ## Consumption (ingress) diff --git a/examples/otel/README.md b/examples/otel/README.md index b9b16a9..42adc81 100644 --- a/examples/otel/README.md +++ b/examples/otel/README.md @@ -1,10 +1,11 @@ # Observability — OpenTelemetry GenAI export (Surface 3) -RiskKernel emits one OpenTelemetry span per governed model call, carrying the -attribute set pinned in [`api/v1/otel-genai.md`](../../api/v1/otel-genai.md): -standard `gen_ai.*` (system, model, token usage, finish reason) plus the -`riskkernel.*` governance extension (run id, step, **cost in USD**, **budget -remaining**, **halt reason**). Point it at any OTLP backend you already run. +RiskKernel emits one OpenTelemetry span per governed model call — and one per +governed MCP **tool call** — carrying the attribute set pinned in +[`api/v1/otel-genai.md`](../../api/v1/otel-genai.md): standard `gen_ai.*` (system, +model, token usage, finish reason) plus the `riskkernel.*` governance extension +(run id, step, **cost in USD**, **budget remaining**, **halt reason**, **tool +status**). Point it at any OTLP backend you already run. > **No telemetry by default.** RiskKernel exports *nothing* unless you set > `OTEL_EXPORTER_OTLP_ENDPOINT`. Spans go only to the endpoint you choose. See @@ -36,6 +37,11 @@ span named `chat claude-sonnet-4-5`. Open it to see the attributes — including `riskkernel.cost.usd`, `riskkernel.budget.tokens.remaining`, and (if the run hit a limit) `riskkernel.halt.reason`. +Governed MCP tool calls land the same way, as `execute_tool {tool}` spans carrying +`gen_ai.tool.name` and `riskkernel.tool.status` — a blocked or denied call is marked +with an error status, so policy refusals stand out in the trace UI next to your +model calls. + ## Other backends Same spans, just change the endpoint: @@ -62,3 +68,6 @@ panels directly from spans (e.g. in Grafana over Tempo, or SigNoz): - **Halts** — count of spans where `riskkernel.halt.reason` is present, grouped by reason (`token_budget_exceeded`, `time_budget_exceeded`, …). - **Latency by model** — span duration grouped by `gen_ai.request.model`. +- **Tool refusals** — count of `execute_tool` spans grouped by + `riskkernel.tool.status` (`blocked` by the allowlist, `denied` at the approval + gate, or `timeout` vs `approved`). diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 051ccdf..3a5aea0 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -117,7 +117,7 @@ func Build(cfg *config.Config) (*Deps, error) { var mcpGW *mcp.Gateway if cfg.MCP.Upstream != "" { - mcpGW = mcp.New(cfg.MCP.Upstream, cfg.MCP.Allowlist, cfg.MCP.ReadOnly, gate, mgr, store, + mcpGW = mcp.New(cfg.MCP.Upstream, cfg.MCP.Allowlist, cfg.MCP.ReadOnly, gate, mgr, store, tracer, time.Duration(cfg.MCP.ApprovalTimeoutSeconds)*time.Second, log) log.Info("mcp gateway enabled", "upstream", cfg.MCP.Upstream, "allowlist", len(cfg.MCP.Allowlist), "readonly", len(cfg.MCP.ReadOnly)) diff --git a/internal/mcp/gateway.go b/internal/mcp/gateway.go index 96bf41e..74a5c1f 100644 --- a/internal/mcp/gateway.go +++ b/internal/mcp/gateway.go @@ -21,6 +21,7 @@ import ( "github.com/prashar32/riskkernel/internal/approval" "github.com/prashar32/riskkernel/internal/id" + "github.com/prashar32/riskkernel/internal/otel" "github.com/prashar32/riskkernel/internal/runs" "github.com/prashar32/riskkernel/internal/storage" ) @@ -40,13 +41,14 @@ type Gateway struct { gate *approval.Gate runs *runs.Manager store storage.Store + tracer *otel.Tracer log *slog.Logger approvalTimeout time.Duration } // New constructs an MCP gateway. upstream must be non-empty. func New(upstream string, allowlist, readonly []string, gate *approval.Gate, - mgr *runs.Manager, store storage.Store, approvalTimeout time.Duration, log *slog.Logger) *Gateway { + mgr *runs.Manager, store storage.Store, tracer *otel.Tracer, approvalTimeout time.Duration, log *slog.Logger) *Gateway { ro := make(map[string]bool, len(readonly)) for _, t := range readonly { ro[t] = true @@ -62,6 +64,7 @@ func New(upstream string, allowlist, readonly []string, gate *approval.Gate, gate: gate, runs: mgr, store: store, + tracer: tracer, log: log, approvalTimeout: approvalTimeout, } @@ -89,6 +92,7 @@ type toolsCallParams struct { } func (g *Gateway) handle(w http.ResponseWriter, r *http.Request) { + start := time.Now() body, err := io.ReadAll(io.LimitReader(r.Body, 10<<20)) if err != nil { http.Error(w, "read error", http.StatusBadRequest) @@ -113,7 +117,7 @@ func (g *Gateway) handle(w http.ResponseWriter, r *http.Request) { // tool call is part of the audit trail, not a silent drop. if !g.allowed(tool) { g.log.Warn("mcp tool blocked by allowlist", "tool", tool) - g.recordToolCall(run.ID, stepIdx, tool, "", params.Arguments, "blocked") + g.recordToolCall(r.Context(), start, run.ID, stepIdx, tool, "", params.Arguments, "blocked") writeRPCError(w, req.ID, -32001, "tool not allowed by policy: "+tool) return } @@ -129,19 +133,19 @@ func (g *Gateway) handle(w http.ResponseWriter, r *http.Request) { SideEffect: sideEffect, Arguments: params.Arguments, }) if aerr != nil { - g.recordToolCall(run.ID, stepIdx, tool, sideEffect, params.Arguments, "timeout") + g.recordToolCall(r.Context(), start, run.ID, stepIdx, tool, sideEffect, params.Arguments, "timeout") writeRPCError(w, req.ID, -32002, "approval timed out or run cancelled for tool: "+tool) return } if !decision.Approved { - g.recordToolCall(run.ID, stepIdx, tool, sideEffect, params.Arguments, "denied") + g.recordToolCall(r.Context(), start, run.ID, stepIdx, tool, sideEffect, params.Arguments, "denied") writeRPCError(w, req.ID, -32003, "approval denied for tool: "+tool) return } } // 3) Forward to the real MCP server and record the (approved) call. - g.recordToolCall(run.ID, stepIdx, tool, sideEffect, params.Arguments, "approved") + g.recordToolCall(r.Context(), start, run.ID, stepIdx, tool, sideEffect, params.Arguments, "approved") g.forward(w, r, body) } @@ -177,7 +181,13 @@ func (g *Gateway) resolveRun(r *http.Request) *runs.Run { return g.runs.Create(runs.CreateOptions{Name: "mcp"}) } -func (g *Gateway) recordToolCall(runID string, step int32, tool, sideEffect string, args map[string]any, status string) { +func (g *Gateway) recordToolCall(ctx context.Context, start time.Time, runID string, step int32, tool, sideEffect string, args map[string]any, status string) { + // Emit an OTLP span so the call (and its governance outcome) is visible in the + // user's observability backend, next to the model-call spans. + g.tracer.RecordToolCall(ctx, otel.ToolCall{ + RunID: runID, StepIndex: step, Tool: tool, SideEffect: sideEffect, + Status: status, Start: start, End: time.Now(), + }) if g.store == nil { return } diff --git a/internal/mcp/gateway_test.go b/internal/mcp/gateway_test.go index 780c6eb..f9232e4 100644 --- a/internal/mcp/gateway_test.go +++ b/internal/mcp/gateway_test.go @@ -16,8 +16,11 @@ import ( "github.com/prashar32/riskkernel/internal/approval" "github.com/prashar32/riskkernel/internal/governor" + "github.com/prashar32/riskkernel/internal/otel" "github.com/prashar32/riskkernel/internal/runs" "github.com/prashar32/riskkernel/internal/storage" + + "go.opentelemetry.io/otel/sdk/trace/tracetest" ) type discard struct{} @@ -46,7 +49,7 @@ func newTestGateway(t *testing.T, allowlist, readonly []string) (*Gateway, *int3 gate := approval.NewGate(store, approval.Policy{DefaultSafe: true}, nil, log) mgr := runs.NewManager(governor.Budget{}).WithStore(store, log) - g := New(upstream.URL, allowlist, readonly, gate, mgr, store, 5*time.Second, log) + g := New(upstream.URL, allowlist, readonly, gate, mgr, store, otel.Disabled(), 5*time.Second, log) return g, &hits } @@ -114,7 +117,7 @@ func TestAllowlistBlockIsAudited(t *testing.T) { gate := approval.NewGate(spy, approval.Policy{DefaultSafe: true}, nil, log) mgr := runs.NewManager(governor.Budget{}).WithStore(spy, log) // Upstream is intentionally unreachable: a blocked tool must never be forwarded. - g := New("http://127.0.0.1:1", []string{"safe_*"}, nil, gate, mgr, spy, time.Second, log) + g := New("http://127.0.0.1:1", []string{"safe_*"}, nil, gate, mgr, spy, otel.Disabled(), time.Second, log) w := httptest.NewRecorder() g.handle(w, mcpReq(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"danger_rm","arguments":{"path":"/etc"}}}`)) @@ -207,3 +210,38 @@ func waitPending(t *testing.T, g *Gateway) string { t.Fatal("pending approval never appeared") return "" } + +func TestToolCallEmitsSpan(t *testing.T) { + // A governed tools/call must surface as an OTLP span with its outcome, so tool + // governance is visible in the user's backend next to model calls. + sr := tracetest.NewSpanRecorder() + store, err := storage.OpenSQLite(filepath.Join(t.TempDir(), "span.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = store.Close() }) + log := slog.New(slog.NewTextHandler(discard{}, nil)) + gate := approval.NewGate(store, approval.Policy{DefaultSafe: true}, nil, log) + mgr := runs.NewManager(governor.Budget{}).WithStore(store, log) + tracer := otel.NewWithProcessor(sr, "test") + // A blocked tool never forwards, so the upstream can be unreachable. + g := New("http://127.0.0.1:1", []string{"safe_*"}, nil, gate, mgr, store, tracer, time.Second, log) + + w := httptest.NewRecorder() + g.handle(w, mcpReq(`{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"danger_rm"}}`)) + + spans := sr.Ended() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + if spans[0].Name() != "execute_tool danger_rm" { + t.Errorf("span name = %q", spans[0].Name()) + } + a := map[string]string{} + for _, kv := range spans[0].Attributes() { + a[string(kv.Key)] = kv.Value.Emit() + } + if a["gen_ai.tool.name"] != "danger_rm" || a["riskkernel.tool.status"] != "blocked" { + t.Fatalf("tool span attrs = %v", a) + } +} diff --git a/internal/otel/otel.go b/internal/otel/otel.go index b8c2a1b..77d0f80 100644 --- a/internal/otel/otel.go +++ b/internal/otel/otel.go @@ -49,6 +49,10 @@ const ( attrBudgetDolLimit = "riskkernel.budget.dollars.limit" attrBudgetDolRemain = "riskkernel.budget.dollars.remaining" attrHaltReason = "riskkernel.halt.reason" + + attrGenAIToolName = "gen_ai.tool.name" + attrToolSideEffect = "riskkernel.tool.side_effect" + attrToolStatus = "riskkernel.tool.status" // approved | blocked | denied | timeout ) // Tracer emits governed-call spans. The zero value and Disabled() are safe no-ops. @@ -232,6 +236,47 @@ func (t *Tracer) RecordCall(ctx context.Context, c Call) { span.End(trace.WithTimestamp(orNow(c.End))) } +// ToolCall is a governed MCP tool call to record as a span. +type ToolCall struct { + RunID string + StepIndex int32 + Tool string + SideEffect string // "" for read-only + Status string // approved | blocked | denied | timeout + Start time.Time + End time.Time +} + +// RecordToolCall emits a span for one governed MCP tool call, so tool governance — +// allowlist blocks, approval denials, approved calls — is visible alongside model +// calls in the user's OTLP backend. No-op when disabled. +func (t *Tracer) RecordToolCall(ctx context.Context, tc ToolCall) { + if !t.Enabled() { + return + } + _, span := t.tracer.Start(ctx, "execute_tool "+tc.Tool, + trace.WithSpanKind(trace.SpanKindInternal), + trace.WithTimestamp(orNow(tc.Start)), + ) + attrs := []attribute.KeyValue{ + attribute.String(attrGenAIOperation, "execute_tool"), + attribute.String(attrGenAIToolName, tc.Tool), + attribute.String(attrRunID, tc.RunID), + attribute.Int(attrStepIndex, int(tc.StepIndex)), + attribute.String(attrToolStatus, tc.Status), + } + if tc.SideEffect != "" { + attrs = append(attrs, attribute.String(attrToolSideEffect, tc.SideEffect)) + } + span.SetAttributes(attrs...) + // A refused call (blocked / denied / timeout) is an error status so it stands + // out in the trace UI. + if tc.Status != "approved" { + span.SetStatus(codes.Error, "tool "+tc.Status) + } + span.End(trace.WithTimestamp(orNow(tc.End))) +} + func orNow(t time.Time) time.Time { if t.IsZero() { return time.Now() diff --git a/internal/otel/otel_test.go b/internal/otel/otel_test.go index c1cf0af..fb3766b 100644 --- a/internal/otel/otel_test.go +++ b/internal/otel/otel_test.go @@ -95,6 +95,31 @@ func TestRecordCall_HaltAttribute(t *testing.T) { } } +func TestRecordToolCall_Attributes(t *testing.T) { + tr, sr := newRecording() + tr.RecordToolCall(context.Background(), ToolCall{ + RunID: "r1", StepIndex: 3, Tool: "write_file", SideEffect: "tool", Status: "denied", + }) + spans := sr.Ended() + if len(spans) != 1 { + t.Fatalf("got %d spans, want 1", len(spans)) + } + s := spans[0] + if s.Name() != "execute_tool write_file" { + t.Errorf("span name = %q", s.Name()) + } + a := attrMap(s.Attributes()) + if a[attrGenAIOperation].AsString() != "execute_tool" || a[attrGenAIToolName].AsString() != "write_file" { + t.Errorf("operation/tool = %v / %v", a[attrGenAIOperation], a[attrGenAIToolName]) + } + if a[attrRunID].AsString() != "r1" || a[attrStepIndex].AsInt64() != 3 { + t.Errorf("run/step = %v / %v", a[attrRunID], a[attrStepIndex]) + } + if a[attrToolStatus].AsString() != "denied" || a[attrToolSideEffect].AsString() != "tool" { + t.Errorf("status/side_effect = %v / %v", a[attrToolStatus], a[attrToolSideEffect]) + } +} + func TestDisabled_NoOp(t *testing.T) { d := Disabled() if d.Enabled() {