From 0bfab1d1955fb5010addcf4541b3c82f37c809d5 Mon Sep 17 00:00:00 2001 From: Alex Wilkerson John Date: Thu, 18 Jun 2026 20:08:59 -0400 Subject: [PATCH 1/2] feat(mcp): HTTP transport, Devin pack, verified fix, sampling/roots, streaming Extend the MCP server beyond the stdio tools into a more complete, agent-native surface (transport-agnostic core shared by both transports): Transport & integration - Streamable HTTP transport (serve --mcp --http): JSON for sync methods, SSE for tools/call, static bearer auth, /healthz, body/concurrency caps, graceful shutdown, session registry with eviction. - Devin integration pack (examples/hooks/devin): HTTP + stdio onboarding, config examples, run/setup scripts. Capabilities - resources (codeguard://rules, codeguard://config, rules/{id} template), prompts (review-diff, triage-findings, explain-rule), logging. - tool annotations (readOnly/destructive hints) + output schemas. Verified auto-fix tools - verify_fix (caller diff), propose_fix (generate via sampling or configured AI provider, then verify), apply_fix (verify then write the working tree, the one destructive tool; confirms via elicitation when supported). - Failures return isError + structuredContent (attempted diff, remaining findings). Verification stays fail-closed. Server->client requests (stdio + HTTP GET-SSE) - sampling/createMessage, roots/list, elicitation/create via a shared serverRequester; client capabilities captured at initialize. - roots feed config_path confinement and are cached per connection (invalidated on notifications/roots/list_changed). - richer sampling request (includeContext + modelPreferences). Streaming & security - scan streams a per-section progress notification (core.ScanOptions OnSectionComplete fired in FinalizeSection). - caller-supplied config_path is confined to the config dir, cwd, and client roots; out-of-tree paths rejected with a generic error (closes a remote arbitrary-file-read vector). Tests cover stdio + HTTP transports, resources/prompts, streaming, verify/apply fail-closed, and bidirectional sampling/roots round-trips. Co-Authored-By: Claude Opus 4.8 (1M context) --- .claude/knowledge/architecture-boundaries.md | 5 + .claude/knowledge/testing-patterns.md | 2 + docs/agent-native.md | 36 +- docs/integrations.md | 123 +++++- examples/hooks/README.md | 9 +- examples/hooks/devin/README.md | 72 ++++ examples/hooks/devin/mcp-http.json.example | 11 + examples/hooks/devin/mcp-stdio.json.example | 16 + examples/hooks/devin/run-http.sh | 36 ++ examples/hooks/devin/setup-snapshot.sh | 37 ++ internal/cli/mcp_client.go | 360 ++++++++++++++++ internal/cli/mcp_dispatch.go | 131 ++++++ internal/cli/mcp_fix.go | 355 ++++++++++++++++ internal/cli/mcp_http.go | 385 ++++++++++++++++++ internal/cli/mcp_http_session.go | 118 ++++++ internal/cli/mcp_prompts.go | 105 +++++ internal/cli/mcp_protocol.go | 143 ++++++- internal/cli/mcp_requests.go | 14 + internal/cli/mcp_resources.go | 133 ++++++ internal/cli/mcp_response.go | 28 +- internal/cli/mcp_run.go | 111 ++++- internal/cli/mcp_tools.go | 172 +++++++- internal/cli/mcp_types.go | 4 + .../codeguard/core/rule_scan_pack_types.go | 5 + internal/codeguard/runner/support/findings.go | 3 + internal/codeguard/runner/support/patch.go | 17 + tests/mcp/http_test.go | 269 ++++++++++++ tests/mcp/sampling_test.go | 332 +++++++++++++++ tests/mcp/smoke_assertions_test.go | 169 +++++++- tests/mcp/smoke_helpers_test.go | 39 +- tests/mcp/smoke_test.go | 30 ++ .../transcripts/apply_fix_failclosed.jsonl | 3 + .../transcripts/prompts_discovery.jsonl | 4 + .../transcripts/resources_discovery.jsonl | 4 + .../testdata/transcripts/scan_streaming.jsonl | 3 + .../transcripts/verify_fix_failclosed.jsonl | 3 + 36 files changed, 3218 insertions(+), 69 deletions(-) create mode 100644 examples/hooks/devin/README.md create mode 100644 examples/hooks/devin/mcp-http.json.example create mode 100644 examples/hooks/devin/mcp-stdio.json.example create mode 100755 examples/hooks/devin/run-http.sh create mode 100755 examples/hooks/devin/setup-snapshot.sh create mode 100644 internal/cli/mcp_client.go create mode 100644 internal/cli/mcp_dispatch.go create mode 100644 internal/cli/mcp_fix.go create mode 100644 internal/cli/mcp_http.go create mode 100644 internal/cli/mcp_http_session.go create mode 100644 internal/cli/mcp_prompts.go create mode 100644 internal/cli/mcp_resources.go create mode 100644 tests/mcp/http_test.go create mode 100644 tests/mcp/sampling_test.go create mode 100644 tests/mcp/testdata/transcripts/apply_fix_failclosed.jsonl create mode 100644 tests/mcp/testdata/transcripts/prompts_discovery.jsonl create mode 100644 tests/mcp/testdata/transcripts/resources_discovery.jsonl create mode 100644 tests/mcp/testdata/transcripts/scan_streaming.jsonl create mode 100644 tests/mcp/testdata/transcripts/verify_fix_failclosed.jsonl diff --git a/.claude/knowledge/architecture-boundaries.md b/.claude/knowledge/architecture-boundaries.md index efddbb5..c669b06 100644 --- a/.claude/knowledge/architecture-boundaries.md +++ b/.claude/knowledge/architecture-boundaries.md @@ -8,4 +8,9 @@ Key architectural decisions, service boundaries, data flow, integration points, - **OWASP categories are baked into the rule catalog at var-init time**, not via `init()`. The merged `catalog` var calls `withSecurityOWASP(mergeRuleCatalogs(...))` in `rules/catalog.go`; an `init()` would run *after* the `catalog` var initializer copied stale entries. The mapping lives in `rules/catalog_security_owasp.go`. All read paths (`Catalog`, `RuleCatalogForConfig`, SARIF, CLI) flow through `rules.Catalog()`. - Config-controlled artifact paths (`baseline.path`, `cache.path`, `ai.cache.path`) are resolved relative to the config dir and contained within it by `config.containConfigArtifactPaths` (in `config/io.go`, `LoadFile`). Paths escaping the config directory are rejected. - The language-agnostic OWASP-gap line rules (CORS, debug, bind-all, weak-hash/cipher, deserialization, Dockerfile root) live in `checks/security/security_owasp_extra.go` and run per-line via `findingsForFile`. NOTE: `findingsForFile` is the **non-TypeScript** path — `securityTargetFindings` routes TS/JS targets to a separate `typeScriptTargetFindings` pipeline, so these rules do not run on TS/JS targets. String-literal signals (CORS, bind-all) match the raw line; call/identifier signals match the masked line. +- **The MCP server is hand-rolled (no SDK) and transport-agnostic.** Shared JSON-RPC message builders + the synchronous method router live in `internal/cli/mcp_dispatch.go` (`buildResultMessage`/`buildErrorMessage`/`buildProgressMessage`, `serverCapabilities`, `dispatchSyncMethod`). Both transports reuse them. stdio (`mcp_run.go`/`mcp_requests.go`) keeps an async goroutine + cancel-map for `tools/call` (cancellation via `notifications/cancelled`); HTTP (`mcp_http.go`) runs `tools/call` synchronously per-request and streams progress over SSE, cancellation driven by the request context. When adding a capability, wire it into `dispatchSyncMethod` (HTTP picks it up automatically) AND the stdio method switch in `mcp_run.go`. +- **MCP server→client requests (sampling/roots)** flow through `internal/cli/mcp_client.go`: `clientCaller`/`clientBridge` + a `serverRequester` (pending-id→chan registry). The transport supplies a `send` closure and feeds inbound responses (classified by `isResponseMessage`: id present, no method) back via `serverRequester.deliver`. stdio demuxes responses in the read loop (`mcp_run.go handleLine`); HTTP needs the client to open the `GET {mcp-path}` SSE stream — server-initiated requests are written there and answered on later POSTs, correlated per session in `internal/cli/mcp_http_session.go`. Client capabilities are captured at `initialize` (`parseClientCapabilities`). The `clientCaller`/progress emitter reach tools via context (`withClientCaller`/`withProgress`), not signatures. +- **MCP fix tools** (`internal/cli/mcp_fix.go`): `verify_fix` wraps `service.VerifyFix` (caller diff); `propose_fix` wraps `service.GenerateVerifiedFix` with a generator resolved as sampling-first (`samplingGenerator` calls the client LLM) then `internalfix.NewAIGenerator(cfg.AI)`; `apply_fix` verifies then writes the tree via `runnersupport.ApplyUnifiedDiff` (the one `destructiveHint` tool), confirming via `clientCaller.elicit` (`elicitation/create`) when supported. Verification test execution is trust-gated, so it needs `CODEGUARD_ALLOW_CONFIG_COMMANDS=1` or it fails closed. Fix failures return `toolErrorResultData` (isError + structuredContent with attempted diff + remaining findings). +- **MCP `config_path` confinement**: caller-supplied `config_path` is confined (`confinePath` in `mcp_tools.go`) to the server config dir + cwd + client roots; the server's own `--config` default is never confined. This closed a remote arbitrary-file-read finding. Streaming: `core.ScanOptions.OnSectionComplete` (tagged `json:"-"` — marshaling a non-nil func errors) fires in `runner/support.FinalizeSection`. +- **`codeguard serve --mcp --http`** serves Streamable HTTP for remote hosts (Devin): `POST {mcp-path}` returns `application/json` for sync methods and `text/event-stream` for `tools/call`; optional static-bearer auth (`--auth-token`/`$CODEGUARD_MCP_AUTH_TOKEN`, constant-time compared), `GET /healthz`, body/concurrency caps, SIGINT/SIGTERM graceful shutdown. No OAuth. MCP resources (`codeguard://rules`, `codeguard://config`, `codeguard://rules/{rule_id}`) and prompts (`review-diff`, `triage-findings`, `explain-rule`) wrap the same SDK accessors the tools use. Devin onboarding pack: `examples/hooks/devin/`. - The taint engine emits a single rule ID per language, chosen by sink: `goTaintRuleID`/`pyTaintRuleID` return `security.ssrf.{go,python}` when the sink is an outbound-HTTP callee (`goHTTPSinkArgIndex`/`pyHTTPSinkArgIndex`), else the generic `security.taint.{go,python}`. To add an SSRF sink, add it to the arg-index map; the rule routing and OWASP A10 mapping follow automatically. diff --git a/.claude/knowledge/testing-patterns.md b/.claude/knowledge/testing-patterns.md index fbd1b1d..e69f597 100644 --- a/.claude/knowledge/testing-patterns.md +++ b/.claude/knowledge/testing-patterns.md @@ -5,3 +5,5 @@ Testing strategies, test infrastructure quirks, how to run/debug specific test s - The trust policy (`internal/codeguard/trust`) is a process-global. Tests that exercise command-driven checks or local (127.0.0.1) AI endpoints must enable it. `tests/checks` and `tests/codeguard` each have a `TestMain` (`trust_main_test.go`) that sets `Policy{AllowConfigCommands: true, AllowConfigAIEndpoints: true}`. If you add a new test package that runs config commands or hits a loopback test server, add a similar `TestMain` or the tests will fail with "refusing to run config-supplied command" / "ssrf guard" errors. - `tests/security` deliberately has NO trust `TestMain`: it verifies the secure *default*. Tests there set/restore the policy locally with `trust.Set(...)` + `t.Cleanup(trust.ResetFromEnv)`. Do not add a package-wide opt-in there. - All tests live under `tests/` as external (`_test`) packages; there are no white-box tests inside `internal/`/`pkg/`. Internal packages are still importable from `tests/` because they are within the module. +- The MCP server smoke tests (`tests/mcp/`) drive the **real binary in a subprocess**, not the in-process handler, to keep tests external. Pattern: a `Test...HelperProcess` func gated by an env var (`GO_WANT_MCP_HELPER_PROCESS` for stdio, `GO_WANT_MCP_HTTP_HELPER_PROCESS` for HTTP) re-execs `os.Args[0]` and calls `cli.Run(...)`. stdio replays NDJSON transcripts over stdin; HTTP reserves a free `127.0.0.1:0` port, launches `serve --mcp --http`, polls `/healthz` for readiness, then issues real HTTP requests. Add new MCP behavior as a transcript + assertion (stdio) or a subtest in `http_test.go` (HTTP). +- **Bidirectional (server→client) MCP tests** live in `tests/mcp/sampling_test.go`: the test acts as the MCP client, advertises `sampling`/`roots` at `initialize`, and answers the server's server-initiated requests. stdio uses interactive `StdinPipe`/`StdoutPipe` (not the replay harness). HTTP opens the `GET /mcp` SSE stream (waits for the `: ready` comment to avoid the attach race), reads the request off the stream, and POSTs the response with the matching `Mcp-Session-Id`. propose_fix verification is expected to fail on the throwaway diff — assert the round trip fired, not a verified patch. The HTTP helper passes `-config` via `CODEGUARD_TEST_HTTP_CONFIG`. diff --git a/docs/agent-native.md b/docs/agent-native.md index e6e4559..36f8720 100644 --- a/docs/agent-native.md +++ b/docs/agent-native.md @@ -9,9 +9,37 @@ This document is a short status brief for `codeguard` features aimed at AI agent - `codeguard serve --mcp` - exposes MCP tools for `scan`, `validate_patch`, and `explain` - also exposes `validate_config` and `list_rules` + - exposes verified auto-fix tools `verify_fix` (verify a caller-supplied diff), `propose_fix` (generate a fix, then verify it), and `apply_fix` (verify then write to the working tree) — see below + - all tools carry read-only / non-destructive annotations and output schemas + - exposes `resources`: `codeguard://rules`, `codeguard://config`, and the `codeguard://rules/{rule_id}` template + - exposes `prompts`: `review-diff`, `triage-findings`, `explain-rule` + - declares a `logging` capability (accepts `logging/setLevel`) + - consumes the client's `sampling` and `roots` capabilities when advertised (server→client requests; see below) + - `scan` streams a progress notification per check section as it completes - supports `initialize`, `tools/list`, `tools/call`, `ping`, progress notifications, and cancellation - covered by CLI compatibility tests and host-shaped smoke tests in `tests/cli/` and `tests/mcp/` +- Verified auto-fix + - `verify_fix` applies a caller-supplied candidate diff in an isolated workspace, re-scans the changed lines, runs the nearest inferred tests, and returns the result only if it passes (fails closed) + - `propose_fix` first generates the candidate, then runs the same verification. The generator is the client's LLM via MCP `sampling` when the client supports it (no API key needed), otherwise a configured AI provider + - `verify_fix`/`propose_fix` wrap `service.VerifyFix` / `service.GenerateVerifiedFix`; verification test execution requires the server to run with command trust enabled (`CODEGUARD_ALLOW_CONFIG_COMMANDS=1`), otherwise it fails closed + - `verify_fix`/`propose_fix` do not mutate the working tree; on verification failure they return `isError` with `structuredContent` (attempted diff + remaining findings) so an agent can iterate + - `apply_fix` verifies first and only then writes the diff to the working tree (the one destructive tool); when the client supports `elicitation` it asks the user to confirm before writing + +- Server→client requests (`sampling`, `roots`, `elicitation`) + - the server issues `sampling/createMessage` (fix generation), `roots/list` (workspace folders), and `elicitation/create` (confirm before `apply_fix` writes) over both transports — on HTTP via the `GET` SSE stream and a session registry + - client roots are cached per connection and invalidated on `notifications/roots/list_changed` + - advertised client `roots` are added to the allowed set for caller-supplied `config_path` confinement; a caller-supplied `config_path` is always confined to the server's config dir, the working directory, and any client roots, and is rejected otherwise + +- `codeguard serve --mcp --http` + - serves the same MCP server over Streamable HTTP for remote / cloud-hosted hosts (e.g. Devin) + - `tools/call` streams progress over SSE; other methods return a single JSON response + - optional static bearer auth (`--auth-token` / `$CODEGUARD_MCP_AUTH_TOKEN`, `--auth-header`), `GET /healthz`, request-size and concurrency limits, and graceful shutdown + - covered by `internal/cli/mcp_http_test.go` + +- Devin integration pack + - HTTP and stdio MCP configs plus host scripts under [examples/hooks/devin](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/README.md:1) + - Patch validation API - `codeguard validate-patch` accepts a unified diff on stdin - `codeguard.RunPatch(ctx, cfg, diffText)` is available in the public Go SDK @@ -49,7 +77,12 @@ This document is a short status brief for `codeguard` features aimed at AI agent ## File map -- MCP server: [internal/cli/mcp_run.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_run.go:1), [internal/cli/mcp_tools.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_tools.go:1), [internal/cli/mcp_protocol.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_protocol.go:1) +- MCP server core: [internal/cli/mcp_run.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_run.go:1), [internal/cli/mcp_dispatch.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_dispatch.go:1), [internal/cli/mcp_tools.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_tools.go:1), [internal/cli/mcp_protocol.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_protocol.go:1) +- MCP HTTP transport: [internal/cli/mcp_http.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_http.go:1), session/SSE registry [internal/cli/mcp_http_session.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_http_session.go:1) +- MCP resources and prompts: [internal/cli/mcp_resources.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_resources.go:1), [internal/cli/mcp_prompts.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_prompts.go:1) +- MCP verified auto-fix tools: [internal/cli/mcp_fix.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_fix.go:1) +- MCP server→client requests (sampling/roots), progress + path confinement: [internal/cli/mcp_client.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/mcp_client.go:1) +- Devin pack: [examples/hooks/devin/README.md](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/README.md:1) - Patch validation CLI: [internal/cli/commands.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/commands.go:101) - Patch validation SDK: [pkg/codeguard/sdk_run.go](/Users/alex/Documents/GitHub/codeguard/pkg/codeguard/sdk_run.go:29) - Agent explain output: [internal/cli/info.go](/Users/alex/Documents/GitHub/codeguard/internal/cli/info.go:33) @@ -62,3 +95,4 @@ This document is a short status brief for `codeguard` features aimed at AI agent 1. Broaden agent-config governance from pattern matching into richer contradictory-instruction and permission-model analysis. 2. Add end-to-end CI coverage for the composite action flow, not just unit coverage for the comment publisher and markdown formatter. +3. Add OAuth resource-server support to the HTTP transport (currently static bearer only) if hosts require Devin's OAuth mode. diff --git a/docs/integrations.md b/docs/integrations.md index c2c9654..b0196b9 100644 --- a/docs/integrations.md +++ b/docs/integrations.md @@ -60,12 +60,18 @@ Included packs: - [examples/hooks/cursor/before-apply.sh](/Users/alex/Documents/GitHub/codeguard/examples/hooks/cursor/before-apply.sh:1) - [examples/hooks/cursor/after-edit.sh](/Users/alex/Documents/GitHub/codeguard/examples/hooks/cursor/after-edit.sh:1) - [examples/hooks/cursor/mcp.json.example](/Users/alex/Documents/GitHub/codeguard/examples/hooks/cursor/mcp.json.example:1) +- Devin + - [examples/hooks/devin/mcp-http.json.example](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/mcp-http.json.example:1) + - [examples/hooks/devin/mcp-stdio.json.example](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/mcp-stdio.json.example:1) + - [examples/hooks/devin/run-http.sh](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/run-http.sh:1) + - [examples/hooks/devin/setup-snapshot.sh](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/setup-snapshot.sh:1) The packs are script-first on purpose: - pre-write hooks call `codeguard validate-patch` - post-edit hooks call `codeguard scan -mode diff` - the Cursor pack also includes an MCP server example for `codeguard serve --mcp` +- the Devin pack provides HTTP and stdio MCP configs plus host scripts (see below) Start with: @@ -77,6 +83,107 @@ export CODEGUARD_BASE_REF=origin/main Then wire the scripts into your host's hook or workflow settings. +## MCP Server + +`codeguard serve --mcp` runs an MCP server (hand-rolled JSON-RPC 2.0, no extra +dependencies). It advertises `tools`, `resources`, `prompts`, and `logging` +capabilities and negotiates protocol versions `2025-11-25` (current) and +`2025-06-18` (compatibility). + +### Tools, resources, prompts + +- Tools: `scan`, `validate_patch`, `validate_config`, `explain`, `list_rules`, + `verify_fix`, `propose_fix` (read-only), and `apply_fix` (destructive — writes + the working tree). `scan` streams a progress notification per check section. +- Resources: `codeguard://rules` (rule catalog), `codeguard://config` (active + configuration), and the `codeguard://rules/{rule_id}` template (per-rule + explanation). +- Prompts: `review-diff`, `triage-findings`, `explain-rule`. + +### Verified auto-fix + +- `verify_fix` takes a caller-supplied candidate `diff` plus the `finding` it + addresses, applies it in an isolated workspace, re-scans the changed lines, + runs the nearest inferred tests, and returns the result only if it passes. +- `propose_fix` generates the candidate first — via the client's LLM through MCP + `sampling` when the client supports it (no API key needed), otherwise a + configured AI provider — then runs the same verification. +- `apply_fix` verifies and, only on success, writes the diff to the working + tree; when the client supports `elicitation` it asks the user to confirm + first. It is the one destructive tool. +- All fail closed. Verification test execution is trust-gated, so run the + server with `CODEGUARD_ALLOW_CONFIG_COMMANDS=1` to let the inferred tests run; + otherwise verification fails closed. On failure, `verify_fix`/`propose_fix` + return `structuredContent` with the attempted diff and remaining findings. + +### Server→client requests (sampling, roots, elicitation) + +The server consumes the client's `sampling`, `roots`, and `elicitation` +capabilities when the client advertises them at `initialize`. These are +server-initiated JSON-RPC requests: over stdio they interleave on the existing +streams; over Streamable HTTP the client must open the `GET {mcp-path}` SSE +stream (using the `Mcp-Session-Id` returned by `initialize`) over which the +server sends `sampling/createMessage`, `roots/list`, and `elicitation/create`, +and answer them on subsequent POSTs. Advertised `roots` widen the allowed set +for `config_path` confinement and are cached per connection (invalidated on +`notifications/roots/list_changed`); `elicitation` is used by `apply_fix` to +confirm before writing. + +### Security + +A caller-supplied `config_path` is confined to the server's config directory, +the working directory, and any client-advertised roots; paths outside are +rejected with a generic error so the HTTP transport is not a filesystem oracle. + +### Transports + +**stdio (default)** — for local subprocess hosts (Claude Code, Cursor): + +```bash +codeguard serve --mcp -config codeguard.yaml -profile ai-safe +``` + +**Streamable HTTP** — for remote / cloud-hosted hosts (e.g. Devin): + +```bash +codeguard serve --mcp --http \ + --addr 0.0.0.0:8080 \ + --mcp-path /mcp \ + --auth-token "$CODEGUARD_MCP_AUTH_TOKEN" \ + --auth-header Authorization \ + -config codeguard.yaml -profile ai-safe +``` + +The HTTP endpoint: + +- accepts JSON-RPC over `POST {mcp-path}`; `tools/call` streams progress as + `text/event-stream` (SSE), other methods return a single `application/json` + response, and notifications return `202`. +- enforces an optional static bearer token (constant-time compared); a blank + token disables auth and should only be used behind a private network. The + token also falls back to `$CODEGUARD_MCP_AUTH_TOKEN`. +- exposes `GET /healthz` for health checks, caps request body size, limits + concurrent tool executions, and drains in-flight requests on SIGINT/SIGTERM. +- mints an `Mcp-Session-Id` on `initialize`; `GET {mcp-path}` returns `405` + (no server-initiated stream), `DELETE` acknowledges session termination. + +### Devin + +Devin connects to a custom MCP server over HTTP (recommended), SSE, or stdio. +See the [Devin pack README](/Users/alex/Documents/GitHub/codeguard/examples/hooks/devin/README.md:1) +for the full walkthrough. In short: + +1. Host the server with `examples/hooks/devin/run-http.sh` behind a reachable + URL (TLS/ingress/tunnel). +2. In Devin → Settings → MCP Marketplace → Add a custom MCP: Transport `HTTP`, + Server URL `https://…/mcp`, Authentication `Auth Header` → + `Authorization: Bearer `. +3. Click **Test listing tools** — Devin runs `initialize` + `tools/list` and + should discover the five tools. + +For a binary-in-snapshot setup instead, use `setup-snapshot.sh` and the stdio +config in `mcp-stdio.json.example`. + ## MCP Smoke Harness `codeguard serve --mcp` is covered by a host-shaped smoke harness in `tests/mcp/testdata/transcripts/` and `tests/mcp/smoke_test.go`. @@ -87,6 +194,15 @@ The harness launches the local MCP server, replays NDJSON transcripts that model - `editor-compat` - `review-agent` - `scan-agent` +- `resource-agent` +- `prompt-agent` +- `streaming-agent` (per-section progress) +- `verify-fix-agent` (verified fix fail-closed) + +Bidirectional server→client flows (sampling, roots) and the HTTP transport are +covered by `tests/mcp/sampling_test.go` and `tests/mcp/http_test.go`, which act +as an MCP client and answer the server's `sampling/createMessage` and +`roots/list` requests. Run it with: @@ -96,11 +212,12 @@ GOROOT=/opt/homebrew/opt/go/libexec GOCACHE=/private/tmp/codeguard-go-cache go t Current scope: -- validates host-like `initialize`, `tools/list`, `tools/call`, `ping`, and config/patch/explain flows -- validates the server's current newline-delimited stdio JSON-RPC transport +- validates host-like `initialize`, `tools/list`, `tools/call`, `ping`, config/patch/explain, and `resources`/`prompts` flows +- validates the server's newline-delimited stdio JSON-RPC transport +- the Streamable-HTTP transport is covered by `internal/cli/mcp_http_test.go` (initialize, auth, SSE tool calls, resources/prompts, health) Out of scope: - automating the actual desktop/editor hosts themselves - `Content-Length` framed stdio compatibility -- prompts/resources/task APIs +- OAuth-based HTTP authentication (static bearer only) diff --git a/examples/hooks/README.md b/examples/hooks/README.md index c5a5aac..5dd01f7 100644 --- a/examples/hooks/README.md +++ b/examples/hooks/README.md @@ -11,14 +11,19 @@ Current packs: - `before-apply.sh` - `after-edit.sh` - `mcp.json.example` +- `devin/` + - `mcp-http.json.example` + - `mcp-stdio.json.example` + - `run-http.sh` + - `setup-snapshot.sh` - `lib/` - - shared shell helpers used by both packs + - shared shell helpers used by the packs Each script is designed to work with the existing `codeguard` CLI: - `codeguard validate-patch` - `codeguard scan -mode diff` -- `codeguard serve --mcp` +- `codeguard serve --mcp` (stdio) or `codeguard serve --mcp --http` (Streamable HTTP, e.g. for Devin) - `codeguard explain -format agent` See the per-pack READMEs for install examples and expected environment variables. diff --git a/examples/hooks/devin/README.md b/examples/hooks/devin/README.md new file mode 100644 index 0000000..87c15ec --- /dev/null +++ b/examples/hooks/devin/README.md @@ -0,0 +1,72 @@ +# Devin Integration Pack + +This pack connects [Devin](https://devin.ai) to the `codeguard` MCP server so a +Devin session can call `scan`, `validate_patch`, `explain`, `validate_config`, +and `list_rules` directly, and read the `codeguard://rules` / `codeguard://config` +resources and the `review-diff` / `triage-findings` / `explain-rule` prompts. + +Devin adds custom MCP servers from **Settings → MCP Marketplace → Add a custom +MCP** (requires the *Manage MCP Servers* permission). It supports three +transports: **HTTP (Streamable HTTP, recommended)**, **SSE (legacy)**, and +**STDIO**. After saving, click **Test listing tools** — Devin runs `initialize` ++ `tools/list` against the server and should discover the five tools. + +There are two ways to wire codeguard in. Pick one. + +## Option A — HTTP (recommended) + +Because Devin runs in the cloud, the simplest robust setup is to host the +codeguard MCP server at a URL Devin can reach and point Devin at it. + +1. Host the server (see [`run-http.sh`](run-http.sh)): + + ```bash + export CODEGUARD_MCP_AUTH_TOKEN="$(openssl rand -hex 24)" + export CODEGUARD_CONFIG=codeguard.yaml + export CODEGUARD_PROFILE=ai-safe + ./run-http.sh # listens on 0.0.0.0:8080/mcp by default + ``` + + Put it behind TLS / an ingress / a tunnel so Devin's cloud can reach it + (e.g. `https://codeguard.your-org.dev/mcp`). The endpoint also exposes + `GET /healthz` for load-balancer health checks. + +2. In Devin, add a custom MCP server: + + - **Transport:** HTTP + - **Server URL:** `https://codeguard.your-org.dev/mcp` + - **Authentication:** Auth Header + - Header key: `Authorization` + - Header value: `Bearer ` + + See [`mcp-http.json.example`](mcp-http.json.example) for the equivalent + config payload. **Never commit the real token** — the examples use + placeholders so codeguard's own prompt-governance checks don't flag them. + +3. Click **Test listing tools** to confirm discovery. + +> Set no auth token only if the endpoint is reachable solely over a private +> network; otherwise always set `CODEGUARD_MCP_AUTH_TOKEN`. + +## Option B — STDIO (binary baked into Devin's machine) + +If you prefer no hosted endpoint, install the `codeguard` binary into Devin's +machine snapshot and configure a STDIO server. + +1. Add [`setup-snapshot.sh`](setup-snapshot.sh) to your Devin machine setup so + the `codeguard` binary is on `PATH` in every session. +2. Add a custom MCP server with **Transport: STDIO** using + [`mcp-stdio.json.example`](mcp-stdio.json.example): + + - Command: `codeguard` + - Args: `serve --mcp -config codeguard.yaml -profile ai-safe` + +The MCP server writes only valid JSON-RPC to stdout (diagnostics go to stderr), +which is what Devin's STDIO transport requires. + +## Files + +- `mcp-http.json.example` — custom MCP config for the HTTP transport +- `mcp-stdio.json.example` — custom MCP config for the STDIO transport +- `run-http.sh` — launch the codeguard MCP server over HTTP +- `setup-snapshot.sh` — install the codeguard binary for STDIO mode diff --git a/examples/hooks/devin/mcp-http.json.example b/examples/hooks/devin/mcp-http.json.example new file mode 100644 index 0000000..2bc532e --- /dev/null +++ b/examples/hooks/devin/mcp-http.json.example @@ -0,0 +1,11 @@ +{ + "mcpServers": { + "codeguard": { + "type": "http", + "url": "https://codeguard.your-org.dev/mcp", + "headers": { + "Authorization": "Bearer ${CODEGUARD_MCP_AUTH_TOKEN}" + } + } + } +} diff --git a/examples/hooks/devin/mcp-stdio.json.example b/examples/hooks/devin/mcp-stdio.json.example new file mode 100644 index 0000000..25d91d5 --- /dev/null +++ b/examples/hooks/devin/mcp-stdio.json.example @@ -0,0 +1,16 @@ +{ + "mcpServers": { + "codeguard": { + "type": "stdio", + "command": "codeguard", + "args": [ + "serve", + "--mcp", + "-config", + "codeguard.yaml", + "-profile", + "ai-safe" + ] + } + } +} diff --git a/examples/hooks/devin/run-http.sh b/examples/hooks/devin/run-http.sh new file mode 100755 index 0000000..e2889a1 --- /dev/null +++ b/examples/hooks/devin/run-http.sh @@ -0,0 +1,36 @@ +#!/bin/sh + +# Launch the codeguard MCP server over Streamable HTTP for Devin (Option A). +# +# Environment: +# CODEGUARD_BIN override the codeguard binary (default: codeguard) +# CODEGUARD_MCP_ADDR listen address (default: 0.0.0.0:8080) +# CODEGUARD_MCP_PATH MCP endpoint path (default: /mcp) +# CODEGUARD_MCP_AUTH_TOKEN static bearer token; if set, required on requests +# CODEGUARD_MCP_AUTH_HEADER header carrying the token (default: Authorization) +# CODEGUARD_CONFIG policy config path (passed as -config) +# CODEGUARD_PROFILE policy profile (passed as -profile) + +set -eu + +SCRIPT_DIR="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)" +# shellcheck source=../lib/codeguard-hook-lib.sh +. "$SCRIPT_DIR/../lib/codeguard-hook-lib.sh" + +require_codeguard + +addr="${CODEGUARD_MCP_ADDR:-0.0.0.0:8080}" +path="${CODEGUARD_MCP_PATH:-/mcp}" +auth_header="${CODEGUARD_MCP_AUTH_HEADER:-Authorization}" + +if [ -z "${CODEGUARD_MCP_AUTH_TOKEN:-}" ]; then + echo "run-http.sh: warning: CODEGUARD_MCP_AUTH_TOKEN is unset; the endpoint will accept unauthenticated requests" >&2 +fi + +# shellcheck disable=SC2086 +exec "$(codeguard_bin)" serve --mcp --http \ + --addr "$addr" \ + --mcp-path "$path" \ + --auth-token "${CODEGUARD_MCP_AUTH_TOKEN:-}" \ + --auth-header "$auth_header" \ + $(codeguard_args) diff --git a/examples/hooks/devin/setup-snapshot.sh b/examples/hooks/devin/setup-snapshot.sh new file mode 100755 index 0000000..7880ec3 --- /dev/null +++ b/examples/hooks/devin/setup-snapshot.sh @@ -0,0 +1,37 @@ +#!/bin/sh + +# Install the codeguard binary into a Devin machine snapshot for STDIO mode +# (Option B). Run this from Devin's machine setup so `codeguard` is on PATH in +# every session. +# +# Environment: +# CODEGUARD_VERSION version/ref to install (default: latest) +# GOBIN install target for `go install` (default: Go's default) + +set -eu + +version="${CODEGUARD_VERSION:-latest}" + +if command -v codeguard >/dev/null 2>&1; then + echo "setup-snapshot.sh: codeguard already on PATH ($(command -v codeguard))" + codeguard version || true + exit 0 +fi + +if command -v go >/dev/null 2>&1; then + echo "setup-snapshot.sh: installing codeguard@${version} via go install" + go install "github.com/devr-tools/codeguard/cmd/codeguard@${version}" + # Ensure the Go bin dir is on PATH for subsequent steps. + gobin="${GOBIN:-$(go env GOPATH)/bin}" + case ":$PATH:" in + *":$gobin:"*) ;; + *) echo "setup-snapshot.sh: add $gobin to PATH (e.g. echo 'export PATH=\"$gobin:\$PATH\"' >> ~/.profile)" >&2 ;; + esac + exit 0 +fi + +echo "setup-snapshot.sh: Go toolchain not found." >&2 +echo "Install Go, or download a release binary from" >&2 +echo " https://github.com/devr-tools/codeguard/releases" >&2 +echo "and place 'codeguard' on PATH." >&2 +exit 1 diff --git a/internal/cli/mcp_client.go b/internal/cli/mcp_client.go new file mode 100644 index 0000000..fbf9132 --- /dev/null +++ b/internal/cli/mcp_client.go @@ -0,0 +1,360 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +// serverRequestTimeout bounds how long the server waits for a client to answer +// a server-initiated request (sampling/createMessage, roots/list). +const serverRequestTimeout = 120 * time.Second + +// clientCaller is the server→client capability surface used by tools: it reports +// the client's advertised capabilities and issues server-initiated requests +// (sampling, roots). Implemented per transport by clientBridge. +type clientCaller interface { + supports(capability string) bool + sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) + listRoots(ctx context.Context) ([]mcpRoot, error) + elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) +} + +// elicitResult is the client's answer to an elicitation/create request. +type elicitResult struct { + Action string `json:"action"` // "accept" | "decline" | "cancel" + Content json.RawMessage `json:"content"` +} + +func (e elicitResult) accepted() bool { return e.Action == "accept" } + +type mcpRoot struct { + URI string `json:"uri"` + Name string `json:"name,omitempty"` +} + +type clientCallerCtxKey struct{} + +func withClientCaller(ctx context.Context, c clientCaller) context.Context { + if c == nil { + return ctx + } + return context.WithValue(ctx, clientCallerCtxKey{}, c) +} + +func clientCallerFrom(ctx context.Context) clientCaller { + c, _ := ctx.Value(clientCallerCtxKey{}).(clientCaller) + return c +} + +// serverRequester correlates server-initiated requests with the client's +// responses. The transport writes the outbound request (via the send closure +// passed to call) and feeds inbound responses back through deliver. +type serverRequester struct { + mu sync.Mutex + pending map[string]chan json.RawMessage + counter int +} + +func newServerRequester() *serverRequester { + return &serverRequester{pending: map[string]chan json.RawMessage{}} +} + +// call issues a server-initiated request: it allocates an id, lets send write +// the request to the transport, and waits for the matching response (or a +// timeout / context cancellation). +func (r *serverRequester) call(ctx context.Context, send func(id string) error) (json.RawMessage, error) { + r.mu.Lock() + r.counter++ + id := fmt.Sprintf("srv-%d", r.counter) + ch := make(chan json.RawMessage, 1) + r.pending[id] = ch + r.mu.Unlock() + defer func() { + r.mu.Lock() + delete(r.pending, id) + r.mu.Unlock() + }() + + if err := send(id); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, serverRequestTimeout) + defer cancel() + select { + case resp := <-ch: + return parseServerResponse(resp) + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// deliver routes an inbound response to a waiting call. It returns false when no +// server-initiated request is awaiting that id (so the caller can dispatch the +// message normally). +func (r *serverRequester) deliver(id string, raw json.RawMessage) bool { + r.mu.Lock() + ch, ok := r.pending[id] + r.mu.Unlock() + if !ok { + return false + } + select { + case ch <- raw: + default: + } + return true +} + +func parseServerResponse(raw json.RawMessage) (json.RawMessage, error) { + var env struct { + Result json.RawMessage `json:"result"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(raw, &env); err != nil { + return nil, err + } + if env.Error != nil { + return nil, fmt.Errorf("client error %d: %s", env.Error.Code, env.Error.Message) + } + return env.Result, nil +} + +// isResponseMessage reports whether an inbound JSON-RPC message is a response to +// a server-initiated request: it carries an id and no method. +func isResponseMessage(req mcpRequest) bool { + return len(req.ID) > 0 && strings.TrimSpace(req.Method) == "" +} + +// decodeIDKey canonicalizes a JSON-RPC id into the string key used by +// serverRequester (which issues string ids like "srv-1"). +func decodeIDKey(raw json.RawMessage) string { + var s string + if json.Unmarshal(raw, &s) == nil { + return s + } + return strings.TrimSpace(string(raw)) +} + +// parseClientCapabilities extracts the capabilities object the client advertised +// during initialize, so the server knows whether sampling/roots are available. +func parseClientCapabilities(raw json.RawMessage) map[string]any { + var params struct { + Capabilities map[string]any `json:"capabilities"` + } + if err := json.Unmarshal(raw, ¶ms); err != nil { + return nil + } + return params.Capabilities +} + +// rootsCache memoizes the client's roots per connection so a config-path check +// does not issue a fresh roots/list round trip every time. It is invalidated on +// notifications/roots/list_changed. +type rootsCache struct { + mu sync.Mutex + loaded bool + roots []mcpRoot +} + +func (c *rootsCache) invalidate() { + c.mu.Lock() + c.loaded = false + c.roots = nil + c.mu.Unlock() +} + +// load returns the cached roots, fetching once via fetch on a miss. The lock is +// held across fetch so concurrent callers coalesce into a single round trip. +func (c *rootsCache) load(fetch func() ([]mcpRoot, error)) ([]mcpRoot, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.loaded { + return c.roots, nil + } + roots, err := fetch() + if err != nil { + return nil, err + } + c.roots = roots + c.loaded = true + return roots, nil +} + +// clientBridge is the transport-agnostic clientCaller. Each transport supplies +// the client capabilities, a serverRequester, a send closure that writes a +// server→client message over that transport, and a per-connection roots cache. +type clientBridge struct { + caps map[string]any + requester *serverRequester + send func(payload any) error + roots *rootsCache +} + +func (b *clientBridge) supports(capability string) bool { + if b == nil || b.caps == nil { + return false + } + _, ok := b.caps[capability] + return ok +} + +func (b *clientBridge) sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) { + if !b.supports("sampling") { + return nil, fmt.Errorf("client does not support sampling") + } + return b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "sampling/createMessage", + "params": params, + }) + }) +} + +func (b *clientBridge) listRoots(ctx context.Context) ([]mcpRoot, error) { + if !b.supports("roots") { + return nil, fmt.Errorf("client does not support roots") + } + if b.roots != nil { + return b.roots.load(func() ([]mcpRoot, error) { return b.fetchRoots(ctx) }) + } + return b.fetchRoots(ctx) +} + +func (b *clientBridge) elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) { + if !b.supports("elicitation") { + return elicitResult{}, fmt.Errorf("client does not support elicitation") + } + raw, err := b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "elicitation/create", + "params": map[string]any{ + "message": message, + "requestedSchema": schema, + }, + }) + }) + if err != nil { + return elicitResult{}, err + } + var result elicitResult + if err := json.Unmarshal(raw, &result); err != nil { + return elicitResult{}, err + } + return result, nil +} + +func (b *clientBridge) fetchRoots(ctx context.Context) ([]mcpRoot, error) { + raw, err := b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "roots/list", + "params": map[string]any{}, + }) + }) + if err != nil { + return nil, err + } + var result struct { + Roots []mcpRoot `json:"roots"` + } + if err := json.Unmarshal(raw, &result); err != nil { + return nil, err + } + return result.Roots, nil +} + +// mcp_client.go holds the per-call transport callbacks that the tool layer +// reaches through context: a progress emitter (for streaming partial findings) +// and, added in a later phase, a client caller (for sampling/roots). Threading +// them via context keeps the many tool method signatures unchanged, and both +// are nil-safe so tools work when a transport does not supply them. + +type progressFunc func(progress float64, total float64, message string) + +type progressCtxKey struct{} + +func withProgress(ctx context.Context, fn progressFunc) context.Context { + if fn == nil { + return ctx + } + return context.WithValue(ctx, progressCtxKey{}, fn) +} + +func progressFrom(ctx context.Context) progressFunc { + fn, _ := ctx.Value(progressCtxKey{}).(progressFunc) + return fn +} + +type clientRootsCtxKey struct{} + +// withClientRoots attaches the filesystem roots the connected client advertised +// (via the roots capability) so config-path confinement can permit them. +func withClientRoots(ctx context.Context, roots []string) context.Context { + if len(roots) == 0 { + return ctx + } + return context.WithValue(ctx, clientRootsCtxKey{}, roots) +} + +func clientRootsFrom(ctx context.Context) []string { + roots, _ := ctx.Value(clientRootsCtxKey{}).([]string) + return roots +} + +// countEnabledSections returns a best-effort count of the scan sections that +// will run for a config, used as the `total` for per-section progress. It is a +// hint, not a guarantee — if it drifts from the runner the progress bar simply +// may not land exactly on 100%. +func countEnabledSections(cfg service.Config, mode service.ScanMode) float64 { + count := 0 + if cfg.Checks.Quality { + count++ + } + if cfg.Checks.Design { + count++ + } + if cfg.Checks.Security { + count++ + } + if cfg.Checks.Prompts { + count++ + } + if cfg.Checks.CI { + count++ + } + if cfg.Checks.SupplyChain { + count++ + } + if cfg.Checks.Contracts != nil { + if *cfg.Checks.Contracts { + count++ + } + } else if mode == service.ScanModeDiff { + count++ + } + for _, pack := range cfg.RulePacks { + if len(pack.Rules) > 0 { + count++ + break + } + } + if count == 0 { + count = 1 + } + return float64(count) +} diff --git a/internal/cli/mcp_dispatch.go b/internal/cli/mcp_dispatch.go new file mode 100644 index 0000000..9b16d59 --- /dev/null +++ b/internal/cli/mcp_dispatch.go @@ -0,0 +1,131 @@ +package cli + +import ( + "encoding/json" + + "github.com/devr-tools/codeguard/internal/version" +) + +// This file holds the transport-neutral core shared by the stdio (mcp_run.go) +// and HTTP (mcp_http.go) transports: JSON-RPC message builders, the advertised +// capability set, and a synchronous method router for everything except +// tools/call (which each transport drives with its own concurrency model). + +const mcpServerInstructions = "Use validate_patch before writing files to disk when you want policy feedback on a proposed diff. Read codeguard://rules for the rule catalog and prompts/get for review workflows." + +// buildResultMessage constructs a JSON-RPC result envelope. Callers that handle +// notifications (empty id) must skip writing the message themselves. +func buildResultMessage(id json.RawMessage, result any) map[string]any { + return map[string]any{ + "jsonrpc": "2.0", + "id": json.RawMessage(id), + "result": result, + } +} + +// buildErrorMessage constructs a JSON-RPC error envelope. A nil id serializes to +// a null id, matching the JSON-RPC spec for unparseable requests. +func buildErrorMessage(id *json.RawMessage, code int, message string) map[string]any { + payload := map[string]any{ + "jsonrpc": "2.0", + "error": mcpError{Code: code, Message: message}, + } + if id != nil { + payload["id"] = json.RawMessage(*id) + } else { + payload["id"] = nil + } + return payload +} + +// buildProgressMessage constructs a notifications/progress envelope. +func buildProgressMessage(token json.RawMessage, progress float64, total float64, message string) map[string]any { + return map[string]any{ + "jsonrpc": "2.0", + "method": "notifications/progress", + "params": map[string]any{ + "progressToken": json.RawMessage(token), + "progress": progress, + "total": total, + "message": message, + }, + } +} + +// serverCapabilities is the capability set advertised during initialize. It is +// shared so both transports report identical capabilities. +func serverCapabilities() map[string]any { + return map[string]any{ + "tools": map[string]any{}, + "resources": map[string]any{}, + "prompts": map[string]any{}, + "logging": map[string]any{}, + } +} + +// buildInitializeResult builds the initialize response payload, negotiating the +// protocol version from the client's requested version. +func buildInitializeResult(params json.RawMessage) map[string]any { + return map[string]any{ + "protocolVersion": negotiateMCPProtocolVersion(params), + "capabilities": serverCapabilities(), + "serverInfo": map[string]any{ + "name": "codeguard", + "title": "CodeGuard MCP Server", + "version": version.Number, + }, + "instructions": mcpServerInstructions, + } +} + +// dispatchSyncMethod handles every request method that produces a single, +// synchronous response (i.e. everything except tools/call, which streams +// progress). It returns the JSON-RPC result envelope and true when it handled +// the method; handled is false for unknown methods so the caller can emit a +// method-not-found error. tools/call must be routed by the transport before +// calling this. +// +// Notifications (notifications/initialized, notifications/cancelled) are not +// handled here — they carry no id and are transport-specific. +func (s *mcpToolService) dispatchSyncMethod(method string, id json.RawMessage, params json.RawMessage) (map[string]any, bool) { + switch method { + case "initialize": + return buildResultMessage(id, buildInitializeResult(params)), true + case "ping": + return buildResultMessage(id, map[string]any{}), true + case "tools/list": + return buildResultMessage(id, map[string]any{"tools": mcpTools()}), true + case "resources/list": + return buildResultMessage(id, map[string]any{"resources": s.resourcesList()}), true + case "resources/templates/list": + return buildResultMessage(id, map[string]any{"resourceTemplates": resourceTemplates()}), true + case "resources/read": + result, errMsg := s.readResource(params) + if errMsg != "" { + return buildErrorMessage(ptrID(id), -32602, errMsg), true + } + return buildResultMessage(id, result), true + case "prompts/list": + return buildResultMessage(id, map[string]any{"prompts": mcpPrompts()}), true + case "prompts/get": + result, errMsg := getPrompt(params) + if errMsg != "" { + return buildErrorMessage(ptrID(id), -32602, errMsg), true + } + return buildResultMessage(id, result), true + case "logging/setLevel": + return buildResultMessage(id, map[string]any{}), true + default: + return nil, false + } +} + +// ptrID returns a pointer to a copy of id, or nil when id is empty, for use with +// buildErrorMessage. +func ptrID(id json.RawMessage) *json.RawMessage { + if len(id) == 0 { + return nil + } + cp := id + return &cp +} diff --git a/internal/cli/mcp_fix.go b/internal/cli/mcp_fix.go new file mode 100644 index 0000000..8c1e8f7 --- /dev/null +++ b/internal/cli/mcp_fix.go @@ -0,0 +1,355 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + internalfix "github.com/devr-tools/codeguard/internal/codeguard/ai/fix" + runnersupport "github.com/devr-tools/codeguard/internal/codeguard/runner/support" + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +// mcp_fix.go exposes codeguard's verified auto-fix flow as MCP tools: +// - verify_fix: the caller supplies a candidate diff; the server applies it +// in a throwaway worktree, re-scans the changed lines, runs the nearest +// inferred tests, and returns the result only if it verifies (fails closed). +// - propose_fix: the server first generates the candidate diff (via the +// client's LLM when the client supports MCP sampling, else a configured AI +// provider) and then runs the same verification. +// +// Neither tool mutates the working tree — verification happens in a temp copy +// and the verified diff is returned for the agent to apply — so both are +// annotated read-only / non-destructive. + +type fixToolArgs struct { + ConfigPath string `json:"config_path"` + Profile string `json:"profile"` + Finding service.Finding `json:"finding"` + Diff string `json:"diff"` + BaseRef string `json:"base_ref"` + MaxNearestTests int `json:"max_nearest_tests"` + TestCommands []service.FixVerificationCommand `json:"test_commands"` +} + +func (a fixToolArgs) options() service.FixOptions { + return service.FixOptions{ + BaseRef: strings.TrimSpace(a.BaseRef), + MaxNearestTests: a.MaxNearestTests, + TestCommands: a.TestCommands, + } +} + +func (s *mcpToolService) callVerifyFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { + var args fixToolArgs + if err := json.Unmarshal(raw, &args); err != nil { + return nil, fmt.Errorf("invalid verify_fix arguments") + } + if strings.TrimSpace(args.Diff) == "" { + return toolErrorResult("verify_fix requires a unified diff"), nil + } + + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) + if err != nil { + return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + } + + result, err := service.VerifyFix(ctx, cfg, args.Finding, service.FixCandidate{Diff: args.Diff}, args.options()) + if err != nil { + data := map[string]any{ + "verified": false, + "error": err.Error(), + "attempted_diff": args.Diff, + } + // Surface what still fails so an agent can iterate, by re-evaluating the + // diff against policy (does not mutate the working tree). + if report, perr := service.RunPatch(ctx, cfg, args.Diff); perr == nil { + data["remaining_findings"] = report + } + return toolErrorResultData(fmt.Sprintf("fix did not verify: %v", err), data), nil + } + return toolSuccessResult(result), nil +} + +func (s *mcpToolService) callProposeFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { + var args fixToolArgs + if err := json.Unmarshal(raw, &args); err != nil { + return nil, fmt.Errorf("invalid propose_fix arguments") + } + if strings.TrimSpace(args.Finding.RuleID) == "" && strings.TrimSpace(args.Finding.Message) == "" { + return toolErrorResult("propose_fix requires a finding to fix"), nil + } + + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) + if err != nil { + return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + } + + generator, err := s.resolveFixGenerator(ctx, cfg) + if err != nil { + return toolErrorResult(fmt.Sprintf("initialize fix generator: %v", err)), nil + } + if generator == nil { + return toolErrorResult("no fix generator available: the client does not support sampling and no AI provider is configured"), nil + } + + result, err := service.GenerateVerifiedFix(ctx, service.FixGenerateRequest{ + Config: cfg, + Finding: args.Finding, + Analysis: firstNonEmpty(args.Finding.Why, args.Finding.Message), + Generator: generator, + Options: args.options(), + }) + if err != nil { + return toolErrorResultData(fmt.Sprintf("generate verified fix: %v", err), map[string]any{ + "verified": false, + "error": err.Error(), + }), nil + } + return toolSuccessResult(result), nil +} + +// callApplyFix verifies a candidate diff and, only if it passes, writes it to +// the working tree. When the client supports elicitation it first asks the user +// to confirm. This is the one tool that mutates the repo (destructiveHint). +func (s *mcpToolService) callApplyFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { + var args fixToolArgs + if err := json.Unmarshal(raw, &args); err != nil { + return nil, fmt.Errorf("invalid apply_fix arguments") + } + if strings.TrimSpace(args.Diff) == "" { + return toolErrorResult("apply_fix requires a unified diff"), nil + } + + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) + if err != nil { + return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + } + + result, err := service.VerifyFix(ctx, cfg, args.Finding, service.FixCandidate{Diff: args.Diff}, args.options()) + if err != nil { + data := map[string]any{"verified": false, "applied": false, "error": err.Error(), "attempted_diff": args.Diff} + if report, perr := service.RunPatch(ctx, cfg, args.Diff); perr == nil { + data["remaining_findings"] = report + } + return toolErrorResultData(fmt.Sprintf("fix did not verify, not applied: %v", err), data), nil + } + + // Confirm with the user before writing, when the client supports it. + if caller := clientCallerFrom(ctx); caller != nil && caller.supports("elicitation") { + accepted, err := confirmApply(ctx, caller, result.ChangedFiles) + if err != nil { + return toolErrorResult(fmt.Sprintf("confirmation failed: %v", err)), nil + } + if !accepted { + return toolSuccessResult(map[string]any{ + "applied": false, + "declined": true, + "diff": result.Diff, + "changed_files": result.ChangedFiles, + }), nil + } + } + + if err := runnersupport.ApplyUnifiedDiff(cfg, result.Diff); err != nil { + return toolErrorResult(fmt.Sprintf("verified fix failed to apply to the working tree: %v", err)), nil + } + return toolSuccessResult(map[string]any{ + "applied": true, + "diff": result.Diff, + "summary": result.Summary, + "changed_files": result.ChangedFiles, + "report": result.Report, + "test_results": result.TestResults, + }), nil +} + +func confirmApply(ctx context.Context, caller clientCaller, changedFiles []string) (bool, error) { + message := fmt.Sprintf("Apply the verified fix to %d file(s): %s?", len(changedFiles), strings.Join(changedFiles, ", ")) + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "apply": map[string]any{"type": "boolean", "description": "Apply the verified fix to the working tree"}, + }, + "required": []string{"apply"}, + } + res, err := caller.elicit(ctx, message, schema) + if err != nil { + return false, err + } + if !res.accepted() { + return false, nil + } + // If the host returned content, honor an explicit apply:false. + var content struct { + Apply *bool `json:"apply"` + } + if json.Unmarshal(res.Content, &content) == nil && content.Apply != nil { + return *content.Apply, nil + } + return true, nil +} + +// resolveFixGenerator picks a fix generator: the connected client's LLM via MCP +// sampling when the client supports it (no API key needed), otherwise a +// configured AI provider. Returns (nil, nil) when neither is available. +func (s *mcpToolService) resolveFixGenerator(ctx context.Context, cfg service.Config) (internalfix.Generator, error) { + if caller := clientCallerFrom(ctx); caller != nil && caller.supports("sampling") { + return samplingGenerator{caller: caller}, nil + } + generator, available, err := internalfix.NewAIGenerator(cfg.AI) + if err != nil { + return nil, err + } + if !available { + return nil, nil + } + return generator, nil +} + +// samplingGenerator generates a fix candidate by asking the connected client's +// LLM through MCP sampling, so no server-side AI provider or API key is needed. +type samplingGenerator struct { + caller clientCaller +} + +func (g samplingGenerator) GenerateFix(ctx context.Context, input internalfix.GenerateInput) (internalfix.Candidate, error) { + params := map[string]any{ + "systemPrompt": "You are a precise code-fixing assistant. Return ONLY a valid unified diff (git format) that fixes the reported finding and changes nothing unrelated. Do not include any prose, explanation, or markdown fences.", + "messages": []map[string]any{{ + "role": "user", + "content": map[string]any{ + "type": "text", + "text": buildSamplingFixPrompt(input), + }, + }}, + "maxTokens": 2048, + // Ask the host to attach this server's context (the repo it is scanning) + // so the model can read the affected files, and bias model selection + // toward capability over speed for a correct patch. + "includeContext": "thisServer", + "modelPreferences": map[string]any{ + "intelligencePriority": 0.8, + "speedPriority": 0.3, + }, + } + raw, err := g.caller.sampleMessage(ctx, params) + if err != nil { + return internalfix.Candidate{}, err + } + text, err := samplingResultText(raw) + if err != nil { + return internalfix.Candidate{}, err + } + candidate := candidateFromText(text) + if strings.TrimSpace(candidate.Diff) == "" { + return internalfix.Candidate{}, fmt.Errorf("sampling response did not contain a diff") + } + return candidate, nil +} + +func buildSamplingFixPrompt(input internalfix.GenerateInput) string { + var b strings.Builder + b.WriteString("Fix this codeguard finding by producing a minimal unified diff.\n\n") + if input.Finding.RuleID != "" { + fmt.Fprintf(&b, "Rule: %s\n", input.Finding.RuleID) + } + if input.Finding.Path != "" { + fmt.Fprintf(&b, "File: %s", input.Finding.Path) + if input.Finding.Line > 0 { + fmt.Fprintf(&b, ":%d", input.Finding.Line) + } + b.WriteString("\n") + } + if input.Finding.Message != "" { + fmt.Fprintf(&b, "Issue: %s\n", input.Finding.Message) + } + if analysis := firstNonEmpty(input.Analysis, input.Finding.Why); analysis != "" { + fmt.Fprintf(&b, "Why: %s\n", analysis) + } + if input.Finding.HowToFix != "" { + fmt.Fprintf(&b, "How to fix: %s\n", input.Finding.HowToFix) + } + if input.Instructions != "" { + fmt.Fprintf(&b, "\n%s\n", input.Instructions) + } + b.WriteString("\nReturn only the unified diff.") + return b.String() +} + +// samplingResultText extracts the assistant text from a sampling/createMessage +// result, tolerating both a single content block and a content array. +func samplingResultText(raw json.RawMessage) (string, error) { + var single struct { + Content struct { + Text string `json:"text"` + } `json:"content"` + } + if err := json.Unmarshal(raw, &single); err == nil && strings.TrimSpace(single.Content.Text) != "" { + return single.Content.Text, nil + } + var multi struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } + if err := json.Unmarshal(raw, &multi); err == nil { + for _, block := range multi.Content { + if strings.TrimSpace(block.Text) != "" { + return block.Text, nil + } + } + } + return "", fmt.Errorf("sampling response had no text content") +} + +// candidateFromText turns a model reply into a fix candidate, accepting a raw +// diff, a {"summary","diff"} JSON object, or a fenced code block. +func candidateFromText(text string) internalfix.Candidate { + trimmed := strings.TrimSpace(text) + var asJSON internalfix.Candidate + if json.Unmarshal([]byte(trimmed), &asJSON) == nil && strings.TrimSpace(asJSON.Diff) != "" { + return asJSON + } + if fenced := extractFencedBlock(trimmed); fenced != "" { + return internalfix.Candidate{Diff: fenced} + } + if looksLikeDiff(trimmed) { + return internalfix.Candidate{Diff: trimmed} + } + return internalfix.Candidate{} +} + +func extractFencedBlock(text string) string { + start := strings.Index(text, "```") + if start < 0 { + return "" + } + rest := text[start+3:] + if nl := strings.IndexByte(rest, '\n'); nl >= 0 { + rest = rest[nl+1:] // drop the optional language tag line + } + end := strings.Index(rest, "```") + if end < 0 { + return strings.TrimSpace(rest) + } + return strings.TrimSpace(rest[:end]) +} + +func looksLikeDiff(text string) bool { + return strings.Contains(text, "@@") || strings.HasPrefix(text, "diff ") || strings.HasPrefix(text, "--- ") +} diff --git a/internal/cli/mcp_http.go b/internal/cli/mcp_http.go new file mode 100644 index 0000000..8d0822a --- /dev/null +++ b/internal/cli/mcp_http.go @@ -0,0 +1,385 @@ +package cli + +import ( + "context" + "crypto/rand" + "crypto/subtle" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync" +) + +// mcp_http.go implements a Streamable-HTTP transport for the MCP server so +// remote, cloud-hosted hosts (e.g. Devin) can reach codeguard over a URL. It +// reuses the transport-neutral core in mcp_dispatch.go: synchronous methods go +// through dispatchSyncMethod and tools/call streams progress over SSE. + +const ( + mcpHTTPMaxBodyBytes = 4 << 20 // 4 MiB request cap + mcpHTTPMaxConcurrentTools = 8 // concurrent tools/call executions + mcpSessionHeader = "Mcp-Session-Id" + mcpDefaultAuthHeader = "Authorization" + mcpDefaultHTTPPath = "/mcp" + mcpHealthPath = "/healthz" + contentTypeJSON = "application/json" + contentTypeEventStream = "text/event-stream" +) + +// mcpAuthConfig configures optional static-bearer auth for the HTTP transport. +// A blank token disables auth (suitable only behind a private network). +type mcpAuthConfig struct { + token string + header string +} + +func (a mcpAuthConfig) enabled() bool { return strings.TrimSpace(a.token) != "" } + +func (a mcpAuthConfig) headerName() string { + if strings.TrimSpace(a.header) == "" { + return mcpDefaultAuthHeader + } + return a.header +} + +// authorize reports whether the request carries the expected credential. For +// the Authorization header it accepts an optional "Bearer " scheme prefix. +func (a mcpAuthConfig) authorize(r *http.Request) bool { + if !a.enabled() { + return true + } + got := strings.TrimSpace(r.Header.Get(a.headerName())) + if strings.EqualFold(a.headerName(), mcpDefaultAuthHeader) { + if rest := strings.TrimSpace(strings.TrimPrefix(got, "Bearer ")); len(rest) != len(got) { + got = rest + } else if rest := strings.TrimSpace(strings.TrimPrefix(got, "bearer ")); len(rest) != len(got) { + got = rest + } + } + want := strings.TrimSpace(a.token) + if len(got) != len(want) { + return false + } + return subtle.ConstantTimeCompare([]byte(got), []byte(want)) == 1 +} + +type mcpHTTPHandler struct { + tools *mcpToolService + auth mcpAuthConfig + path string + sem chan struct{} + sessions *sessionRegistry +} + +// newMCPHTTPHandler builds the HTTP handler for the MCP server. It is split out +// from the listener so tests can mount it on httptest.NewServer. +func newMCPHTTPHandler(tools *mcpToolService, auth mcpAuthConfig, path string) http.Handler { + if strings.TrimSpace(path) == "" { + path = mcpDefaultHTTPPath + } + h := &mcpHTTPHandler{ + tools: tools, + auth: auth, + path: path, + sem: make(chan struct{}, mcpHTTPMaxConcurrentTools), + sessions: newSessionRegistry(), + } + mux := http.NewServeMux() + mux.HandleFunc(mcpHealthPath, h.handleHealth) + mux.HandleFunc(path, h.handleMCP) + return mux +} + +// callerForRequest returns the client caller bound to the request's session, or +// nil when no session/header is present. +func (h *mcpHTTPHandler) callerForRequest(r *http.Request) clientCaller { + sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) + if !ok { + return nil + } + return sess.caller() +} + +func (h *mcpHTTPHandler) handleHealth(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/plain; charset=utf-8") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "ok\n") +} + +func (h *mcpHTTPHandler) handleMCP(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + h.handlePost(w, r) + case http.MethodGet: + h.handleGetStream(w, r) + case http.MethodDelete: + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + return + } + h.sessions.remove(r.Header.Get(mcpSessionHeader)) + w.WriteHeader(http.StatusOK) + default: + w.Header().Set("Allow", "GET, POST, DELETE") + w.WriteHeader(http.StatusMethodNotAllowed) + } +} + +// handleGetStream opens the server→client SSE stream for a session. The server +// writes sampling/createMessage and roots/list requests over this stream; the +// client answers them on subsequent POSTs. +func (h *mcpHTTPHandler) handleGetStream(w http.ResponseWriter, r *http.Request) { + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + return + } + sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) + if !ok { + w.WriteHeader(http.StatusNotFound) + return + } + flusher, ok := w.(http.Flusher) + if !ok { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", contentTypeEventStream) + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher.Flush() + + sess.attachStream(w, flusher) + defer sess.detachStream() + // Emit an SSE comment so the client knows the stream is attached and the + // server can now deliver server-initiated requests over it. + _, _ = io.WriteString(w, ": ready\n\n") + flusher.Flush() + <-r.Context().Done() + // The client disconnected; drop the session so the map does not retain dead + // sessions (the common client keeps one stream open for the session's life). + h.sessions.remove(sess.id) +} + +func (h *mcpHTTPHandler) handlePost(w http.ResponseWriter, r *http.Request) { + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + writeJSONStatus(w, http.StatusUnauthorized, buildErrorMessage(nil, -32001, "unauthorized")) + return + } + + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, mcpHTTPMaxBodyBytes)) + if err != nil { + writeJSONStatus(w, http.StatusRequestEntityTooLarge, buildErrorMessage(nil, -32600, "request too large")) + return + } + trimmed := strings.TrimSpace(string(body)) + if trimmed == "" { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + + // JSON-RPC batch (array) — process each message and return an array of the + // non-notification responses. Batched requests do not stream progress. + if strings.HasPrefix(trimmed, "[") { + h.handleBatch(w, r, []byte(trimmed)) + return + } + + var req mcpRequest + if err := json.Unmarshal([]byte(trimmed), &req); err != nil { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + if req.JSONRPC != "2.0" { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(req.idPtr(), -32600, "invalid request")) + return + } + + // A response (id, no method) is the client answering a server-initiated + // request; route it to the session's pending caller. + if isResponseMessage(req) { + if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { + sess.requester.deliver(decodeIDKey(req.ID), json.RawMessage(trimmed)) + } + w.WriteHeader(http.StatusAccepted) + return + } + + if req.Method == "notifications/roots/list_changed" { + if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { + sess.roots.invalidate() + } + w.WriteHeader(http.StatusAccepted) + return + } + + // Notifications carry no id and expect no response body. + if len(req.ID) == 0 { + w.WriteHeader(http.StatusAccepted) + return + } + + if req.Method == "initialize" { + sess := h.sessions.create(parseClientCapabilities(req.Params)) + w.Header().Set(mcpSessionHeader, sess.id) + } + + if req.Method == "tools/call" { + h.streamToolCall(w, r, req) + return + } + + msg, handled := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) + if !handled { + writeJSONStatus(w, http.StatusOK, buildErrorMessage(req.idPtr(), -32601, "method not found")) + return + } + writeJSONStatus(w, http.StatusOK, msg) +} + +func (h *mcpHTTPHandler) handleBatch(w http.ResponseWriter, r *http.Request, body []byte) { + var reqs []mcpRequest + if err := json.Unmarshal(body, &reqs); err != nil { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + ctx := withClientCaller(r.Context(), h.callerForRequest(r)) + responses := make([]map[string]any, 0, len(reqs)) + for _, req := range reqs { + if req.JSONRPC != "2.0" { + responses = append(responses, buildErrorMessage(req.idPtr(), -32600, "invalid request")) + continue + } + if len(req.ID) == 0 { + continue // notification — no response + } + responses = append(responses, h.produceMessage(ctx, req)) + } + if len(responses) == 0 { + w.WriteHeader(http.StatusAccepted) + return + } + writeJSONStatus(w, http.StatusOK, responses) +} + +// produceMessage resolves a single request to its response envelope, running +// tools/call synchronously (no progress streaming). Used for batched requests. +func (h *mcpHTTPHandler) produceMessage(ctx context.Context, req mcpRequest) map[string]any { + if req.Method == "tools/call" { + if err := h.acquire(ctx); err != nil { + return buildErrorMessage(req.idPtr(), -32603, "server busy") + } + defer h.release() + result, err := h.tools.callToolWithContext(ctx, req.Params) + if err != nil { + return buildErrorMessage(req.idPtr(), -32602, err.Error()) + } + return buildResultMessage(req.ID, result) + } + msg, handled := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) + if !handled { + return buildErrorMessage(req.idPtr(), -32601, "method not found") + } + return msg +} + +// streamToolCall runs a tools/call and streams progress + result as SSE so the +// HTTP transport matches the stdio transport's progress behavior. Cancellation +// is driven by the request context (client disconnect). +func (h *mcpHTTPHandler) streamToolCall(w http.ResponseWriter, r *http.Request, req mcpRequest) { + flusher, ok := w.(http.Flusher) + if !ok { + // No streaming support — fall back to a single JSON response. + writeJSONStatus(w, http.StatusOK, h.produceMessage(r.Context(), req)) + return + } + + if err := h.acquire(r.Context()); err != nil { + writeJSONStatus(w, http.StatusServiceUnavailable, buildErrorMessage(req.idPtr(), -32603, "server busy")) + return + } + defer h.release() + + w.Header().Set("Content-Type", contentTypeEventStream) + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + + ctx := withClientCaller(r.Context(), h.callerForRequest(r)) + progressToken := progressTokenFromParams(req.Params) + if progressToken != nil { + _ = writeSSE(w, buildProgressMessage(*progressToken, 0, 1, "Started")) + flusher.Flush() + var progressMu sync.Mutex + ctx = withProgress(ctx, func(progress float64, total float64, message string) { + progressMu.Lock() + defer progressMu.Unlock() + _ = writeSSE(w, buildProgressMessage(*progressToken, progress, total, message)) + flusher.Flush() + }) + } + + result, err := h.tools.callToolWithContext(ctx, req.Params) + + if progressToken != nil { + message := "Completed" + if err != nil { + message = "Stopped" + } + _ = writeSSE(w, buildProgressMessage(*progressToken, 1, 1, message)) + flusher.Flush() + } + + var msg map[string]any + if err != nil { + msg = buildErrorMessage(req.idPtr(), -32602, err.Error()) + } else { + msg = buildResultMessage(req.ID, result) + } + _ = writeSSE(w, msg) + flusher.Flush() +} + +func (h *mcpHTTPHandler) acquire(ctx context.Context) error { + select { + case h.sem <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (h *mcpHTTPHandler) release() { <-h.sem } + +func writeJSONStatus(w http.ResponseWriter, status int, payload any) { + w.Header().Set("Content-Type", contentTypeJSON) + w.WriteHeader(status) + data, err := json.Marshal(payload) + if err != nil { + return + } + _, _ = w.Write(data) + _, _ = w.Write([]byte("\n")) +} + +func writeSSE(w io.Writer, payload any) error { + data, err := json.Marshal(payload) + if err != nil { + return err + } + _, err = fmt.Fprintf(w, "data: %s\n\n", data) + return err +} + +// newSessionID returns a random hex session identifier. +func newSessionID() string { + buf := make([]byte, 16) + if _, err := rand.Read(buf); err != nil { + return "codeguard-session" + } + return hex.EncodeToString(buf) +} diff --git a/internal/cli/mcp_http_session.go b/internal/cli/mcp_http_session.go new file mode 100644 index 0000000..f43b428 --- /dev/null +++ b/internal/cli/mcp_http_session.go @@ -0,0 +1,118 @@ +package cli + +import ( + "errors" + "net/http" + "strings" + "sync" +) + +// mcp_http_session.go gives the Streamable-HTTP transport the state needed for +// server→client requests. A session is created on initialize (its id returned +// via Mcp-Session-Id); the client opens a GET SSE stream for that session over +// which the server writes sampling/roots requests, and answers them on later +// POSTs which are routed back through the session's serverRequester. + +var errNoClientStream = errors.New("client SSE stream is not connected") + +// maxHTTPSessions bounds the session map so streamless clients that never send +// DELETE cannot grow it without limit; the oldest session is evicted on create. +const maxHTTPSessions = 512 + +type httpSession struct { + id string + caps map[string]any + requester *serverRequester + roots *rootsCache + + streamMu sync.Mutex + stream http.ResponseWriter + flusher http.Flusher + attached bool +} + +func (sess *httpSession) attachStream(w http.ResponseWriter, f http.Flusher) { + sess.streamMu.Lock() + sess.stream = w + sess.flusher = f + sess.attached = true + sess.streamMu.Unlock() +} + +func (sess *httpSession) detachStream() { + sess.streamMu.Lock() + sess.stream = nil + sess.flusher = nil + sess.attached = false + sess.streamMu.Unlock() +} + +// sendToClient writes a server→client message over the session's SSE stream. +func (sess *httpSession) sendToClient(payload any) error { + sess.streamMu.Lock() + defer sess.streamMu.Unlock() + if !sess.attached || sess.stream == nil { + return errNoClientStream + } + if err := writeSSE(sess.stream, payload); err != nil { + return err + } + sess.flusher.Flush() + return nil +} + +func (sess *httpSession) caller() clientCaller { + return &clientBridge{caps: sess.caps, requester: sess.requester, send: sess.sendToClient, roots: sess.roots} +} + +type sessionRegistry struct { + mu sync.Mutex + sessions map[string]*httpSession + order []string + max int +} + +func newSessionRegistry() *sessionRegistry { + return &sessionRegistry{sessions: map[string]*httpSession{}, max: maxHTTPSessions} +} + +func (r *sessionRegistry) create(caps map[string]any) *httpSession { + sess := &httpSession{id: newSessionID(), caps: caps, requester: newServerRequester(), roots: &rootsCache{}} + r.mu.Lock() + if r.max > 0 && len(r.order) >= r.max { + oldest := r.order[0] + r.order = r.order[1:] + delete(r.sessions, oldest) + } + r.sessions[sess.id] = sess + r.order = append(r.order, sess.id) + r.mu.Unlock() + return sess +} + +func (r *sessionRegistry) get(id string) (*httpSession, bool) { + if strings.TrimSpace(id) == "" { + return nil, false + } + r.mu.Lock() + sess, ok := r.sessions[id] + r.mu.Unlock() + return sess, ok +} + +func (r *sessionRegistry) remove(id string) { + if strings.TrimSpace(id) == "" { + return + } + r.mu.Lock() + if _, ok := r.sessions[id]; ok { + delete(r.sessions, id) + for i, existing := range r.order { + if existing == id { + r.order = append(r.order[:i], r.order[i+1:]...) + break + } + } + } + r.mu.Unlock() +} diff --git a/internal/cli/mcp_prompts.go b/internal/cli/mcp_prompts.go new file mode 100644 index 0000000..a77e759 --- /dev/null +++ b/internal/cli/mcp_prompts.go @@ -0,0 +1,105 @@ +package cli + +import ( + "encoding/json" + "fmt" + "strings" +) + +// MCP prompts ship reusable workflows that steer an agent toward the codeguard +// tools for common review tasks. Each prompt renders to one or more user +// messages that name the tool to call and the arguments to pass. + +type mcpPrompt struct { + Name string `json:"name"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + Arguments []mcpPromptArgument `json:"arguments,omitempty"` +} + +type mcpPromptArgument struct { + Name string `json:"name"` + Description string `json:"description,omitempty"` + Required bool `json:"required,omitempty"` +} + +func mcpPrompts() []mcpPrompt { + return []mcpPrompt{ + { + Name: "review-diff", + Title: "Review a diff against policy", + Description: "Validate a proposed unified diff with the validate_patch tool before writing it to disk.", + Arguments: []mcpPromptArgument{ + {Name: "diff", Description: "The unified diff to validate.", Required: true}, + }, + }, + { + Name: "triage-findings", + Title: "Triage repository findings", + Description: "Run a full scan and summarize the failing sections with concrete next steps.", + Arguments: []mcpPromptArgument{}, + }, + { + Name: "explain-rule", + Title: "Explain a rule", + Description: "Fetch machine-first explanation metadata for a rule and restate how to fix it.", + Arguments: []mcpPromptArgument{ + {Name: "rule_id", Description: "The rule id to explain (e.g. security.hardcoded-secret).", Required: true}, + }, + }, + } +} + +// getPrompt resolves a prompts/get request. The second return value is a +// non-empty error message when the prompt could not be rendered. +func getPrompt(params json.RawMessage) (map[string]any, string) { + var args struct { + Name string `json:"name"` + Arguments map[string]string `json:"arguments"` + } + if err := json.Unmarshal(params, &args); err != nil { + return nil, "invalid prompts/get params" + } + + switch strings.TrimSpace(args.Name) { + case "review-diff": + diff := strings.TrimSpace(args.Arguments["diff"]) + if diff == "" { + return nil, "prompt review-diff requires a diff argument" + } + return promptResult( + "Validate a proposed diff against codeguard policy.", + fmt.Sprintf("Call the validate_patch tool with this unified diff and report any failing findings before applying it. Do not write the diff to disk if validation fails.\n\n```diff\n%s\n```", diff), + ), "" + case "triage-findings": + return promptResult( + "Triage codeguard findings for the repository.", + "Call the scan tool with mode \"full\". Group the resulting findings by section, list every failing section first, and for each failing finding give the rule id, the file:line, and a one-line fix suggestion drawn from the finding's how_to_fix field.", + ), "" + case "explain-rule": + ruleID := strings.TrimSpace(args.Arguments["rule_id"]) + if ruleID == "" { + return nil, "prompt explain-rule requires a rule_id argument" + } + return promptResult( + fmt.Sprintf("Explain the %s rule.", ruleID), + fmt.Sprintf("Call the explain tool with rule_id %q (or read the codeguard://rules/%s resource). Summarize what the rule checks, why it matters, and the concrete steps to fix a violation.", ruleID, ruleID), + ), "" + default: + return nil, fmt.Sprintf("unknown prompt %q", args.Name) + } +} + +// promptResult builds an MCP prompts/get result with a single user message. +func promptResult(description string, text string) map[string]any { + return map[string]any{ + "description": description, + "messages": []map[string]any{{ + "role": "user", + "content": map[string]any{ + "type": "text", + "text": text, + }, + }}, + } +} diff --git a/internal/cli/mcp_protocol.go b/internal/cli/mcp_protocol.go index 2d0e93c..17a51b3 100644 --- a/internal/cli/mcp_protocol.go +++ b/internal/cli/mcp_protocol.go @@ -27,6 +27,17 @@ func toolErrorResult(message string) map[string]any { } } +// toolErrorResultData is toolErrorResult with machine-readable structuredContent +// so callers can act on the failure (e.g. a fix that did not verify carries the +// attempted diff and remaining findings). +func toolErrorResultData(message string, data any) map[string]any { + result := toolErrorResult(message) + if data != nil { + result["structuredContent"] = data + } + return result +} + func mustJSON(value any) string { data, err := json.Marshal(value) if err != nil { @@ -36,7 +47,7 @@ func mustJSON(value any) string { } func mcpTools() []mcpTool { - return []mcpTool{ + tools := []mcpTool{ { Name: "scan", Title: "Scan Repository", @@ -50,6 +61,7 @@ func mcpTools() []mcpTool { "base_ref": map[string]any{"type": "string"}, }, }, + OutputSchema: reportOutputSchema(), }, { Name: "validate_config", @@ -62,6 +74,7 @@ func mcpTools() []mcpTool { "profile": map[string]any{"type": "string"}, }, }, + OutputSchema: objectOutputSchema(), }, { Name: "validate_patch", @@ -76,6 +89,7 @@ func mcpTools() []mcpTool { }, "required": []string{"diff"}, }, + OutputSchema: reportOutputSchema(), }, { Name: "explain", @@ -90,6 +104,7 @@ func mcpTools() []mcpTool { }, "required": []string{"rule_id"}, }, + OutputSchema: objectOutputSchema(), }, { Name: "list_rules", @@ -102,6 +117,132 @@ func mcpTools() []mcpTool { "profile": map[string]any{"type": "string"}, }, }, + OutputSchema: objectOutputSchema(), + }, + { + Name: "verify_fix", + Title: "Verify Fix", + Description: "Verify a candidate unified diff against a finding: apply it in an isolated workspace, re-scan the changed lines, run the nearest inferred tests, and return the result only if it passes. Does not modify the working tree.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "finding": fixFindingSchema(), + "diff": map[string]any{"type": "string", "description": "Candidate unified diff to verify."}, + "base_ref": map[string]any{"type": "string"}, + "max_nearest_tests": map[string]any{"type": "integer"}, + "test_commands": map[string]any{"type": "array"}, + }, + "required": []string{"diff"}, + }, + OutputSchema: objectOutputSchema(), + }, + { + Name: "propose_fix", + Title: "Propose Fix", + Description: "Generate a candidate fix for a finding (via the client's LLM when MCP sampling is supported, else a configured AI provider), verify it in an isolated workspace with re-scan and nearest tests, and return it only if it passes. Does not modify the working tree.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "finding": fixFindingSchema(), + "base_ref": map[string]any{"type": "string"}, + "max_nearest_tests": map[string]any{"type": "integer"}, + "test_commands": map[string]any{"type": "array"}, + }, + "required": []string{"finding"}, + }, + OutputSchema: objectOutputSchema(), + }, + { + Name: "apply_fix", + Title: "Apply Fix", + Description: "Verify a candidate unified diff and, only if it passes, write it to the working tree. Asks the user to confirm first when the client supports elicitation. This tool modifies files on disk.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "finding": fixFindingSchema(), + "diff": map[string]any{"type": "string", "description": "Candidate unified diff to verify and apply."}, + "base_ref": map[string]any{"type": "string"}, + "max_nearest_tests": map[string]any{"type": "integer"}, + "test_commands": map[string]any{"type": "array"}, + }, + "required": []string{"diff"}, + }, + OutputSchema: objectOutputSchema(), + Annotations: writeToolAnnotations("Apply Fix"), + }, + } + + // Most codeguard tools only read the repository; tools that already declared + // annotations (e.g. the destructive apply_fix) keep them. + for i := range tools { + if tools[i].Annotations == nil { + tools[i].Annotations = readOnlyToolAnnotations(tools[i].Title) + } + } + return tools +} + +// writeToolAnnotations marks a tool that mutates the working tree. +func writeToolAnnotations(title string) map[string]any { + return map[string]any{ + "title": title, + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": false, + } +} + +// fixFindingSchema describes the finding object accepted by the fix tools. +func fixFindingSchema() map[string]any { + return map[string]any{ + "type": "object", + "description": "The finding to fix, as returned in a scan/validate_patch report.", + "properties": map[string]any{ + "rule_id": map[string]any{"type": "string"}, + "path": map[string]any{"type": "string"}, + "line": map[string]any{"type": "integer"}, + "message": map[string]any{"type": "string"}, + "why": map[string]any{"type": "string"}, + "how_to_fix": map[string]any{"type": "string"}, + }, + } +} + +// readOnlyToolAnnotations returns MCP tool annotations marking a tool as +// read-only and non-destructive. +func readOnlyToolAnnotations(title string) map[string]any { + return map[string]any{ + "title": title, + "readOnlyHint": true, + "destructiveHint": false, + "idempotentHint": true, + "openWorldHint": false, + } +} + +// objectOutputSchema is a permissive output schema for tools that return an +// arbitrary JSON object in structuredContent. +func objectOutputSchema() map[string]any { + return map[string]any{"type": "object"} +} + +// reportOutputSchema describes the codeguard report object returned by scan and +// validate_patch in structuredContent. +func reportOutputSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "sections": map[string]any{"type": "array"}, + "summary": map[string]any{"type": "object"}, }, } } diff --git a/internal/cli/mcp_requests.go b/internal/cli/mcp_requests.go index 2548a63..95f64bd 100644 --- a/internal/cli/mcp_requests.go +++ b/internal/cli/mcp_requests.go @@ -35,14 +35,28 @@ func (s *mcpServer) handleToolCall(req mcpRequest, stdout io.Writer) error { s.mu.Lock() s.active[key] = cancel delete(s.cancelled, key) + caps := s.clientCaps s.mu.Unlock() + bridge := &clientBridge{ + caps: caps, + requester: s.requester, + send: func(payload any) error { + return s.responder.writeMessage(stdout, payload) + }, + roots: s.rootsCache, + } + ctx = withClientCaller(ctx, bridge) + s.wg.Add(1) go func() { defer s.wg.Done() defer s.finishRequest(key) if progressToken != nil { _ = s.responder.writeProgress(stdout, *progressToken, 0, 1, "Started") + ctx = withProgress(ctx, func(progress float64, total float64, message string) { + _ = s.responder.writeProgress(stdout, *progressToken, progress, total, message) + }) } result, err := s.tools.callToolWithContext(ctx, req.Params) if progressToken != nil { diff --git a/internal/cli/mcp_resources.go b/internal/cli/mcp_resources.go new file mode 100644 index 0000000..64ed03f --- /dev/null +++ b/internal/cli/mcp_resources.go @@ -0,0 +1,133 @@ +package cli + +import ( + "encoding/json" + "fmt" + "strings" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +// MCP resources expose codeguard's static, read-only knowledge (the rule +// catalog and the active configuration) so agents can pull context without a +// tool round-trip. They wrap the same SDK accessors used by the tools. + +const ( + resourceURIRules = "codeguard://rules" + resourceURIConfig = "codeguard://config" + resourceRulePrefix = "codeguard://rules/" + resourceRuleURITmpl = "codeguard://rules/{rule_id}" + resourceMIMEType = "application/json" + resourceContentsField = "contents" +) + +type mcpResource struct { + URI string `json:"uri"` + Name string `json:"name"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + MIMEType string `json:"mimeType,omitempty"` +} + +type mcpResourceTemplate struct { + URITemplate string `json:"uriTemplate"` + Name string `json:"name"` + Title string `json:"title,omitempty"` + Description string `json:"description,omitempty"` + MIMEType string `json:"mimeType,omitempty"` +} + +// resourcesList returns the concrete (non-templated) resources. +func (s *mcpToolService) resourcesList() []mcpResource { + return []mcpResource{ + { + URI: resourceURIRules, + Name: "rules", + Title: "Rule Catalog", + Description: "The full codeguard rule catalog that applies to the active configuration.", + MIMEType: resourceMIMEType, + }, + { + URI: resourceURIConfig, + Name: "config", + Title: "Active Configuration", + Description: "The resolved codeguard policy configuration the server was started with.", + MIMEType: resourceMIMEType, + }, + } +} + +// resourceTemplates returns the parameterized resources. +func resourceTemplates() []mcpResourceTemplate { + return []mcpResourceTemplate{ + { + URITemplate: resourceRuleURITmpl, + Name: "rule", + Title: "Rule Explanation", + Description: "Machine-first explanation metadata for a single rule, addressed by rule id.", + MIMEType: resourceMIMEType, + }, + } +} + +// readResource resolves a resources/read request. The second return value is a +// non-empty error message when the resource could not be served. +func (s *mcpToolService) readResource(params json.RawMessage) (map[string]any, string) { + var args struct { + URI string `json:"uri"` + } + if err := json.Unmarshal(params, &args); err != nil { + return nil, "invalid resources/read params" + } + uri := strings.TrimSpace(args.URI) + switch { + case uri == resourceURIRules: + return resourceResult(uri, map[string]any{"rules": s.rulesForDefaults()}), "" + case uri == resourceURIConfig: + cfg, err := s.loadConfig("", "") + if err != nil { + return nil, fmt.Sprintf("load config: %v", err) + } + return resourceResult(uri, cfg), "" + case strings.HasPrefix(uri, resourceRulePrefix): + ruleID := strings.TrimPrefix(uri, resourceRulePrefix) + if strings.TrimSpace(ruleID) == "" { + return nil, "resource uri missing rule id" + } + rule, ok, err := s.resolveExplainRule(s.defaultConfigPath, s.defaultProfile, ruleID) + if err != nil { + return nil, fmt.Sprintf("load config: %v", err) + } + if !ok { + return nil, fmt.Sprintf("unknown rule %q", ruleID) + } + return resourceResult(uri, buildExplainAgentOutput(rule)), "" + default: + return nil, fmt.Sprintf("unknown resource %q", uri) + } +} + +// rulesForDefaults returns the rule catalog scoped to the server's default +// config when it loads, falling back to the built-in catalog otherwise. +func (s *mcpToolService) rulesForDefaults() []service.RuleMetadata { + if strings.TrimSpace(s.defaultConfigPath) == "" { + return service.Rules() + } + cfg, err := s.loadConfig("", "") + if err != nil { + return service.Rules() + } + return service.RulesForConfig(cfg) +} + +// resourceResult builds an MCP resources/read result with a single JSON text +// content block. +func resourceResult(uri string, payload any) map[string]any { + return map[string]any{ + resourceContentsField: []map[string]any{{ + "uri": uri, + "mimeType": resourceMIMEType, + "text": mustJSON(payload), + }}, + } +} diff --git a/internal/cli/mcp_response.go b/internal/cli/mcp_response.go index 4fadbee..f666cf9 100644 --- a/internal/cli/mcp_response.go +++ b/internal/cli/mcp_response.go @@ -10,37 +10,15 @@ func (s *mcpResponder) writeResult(stdout io.Writer, id json.RawMessage, result if len(id) == 0 { return nil } - return s.writeMessage(stdout, map[string]any{ - "jsonrpc": "2.0", - "id": json.RawMessage(id), - "result": result, - }) + return s.writeMessage(stdout, buildResultMessage(id, result)) } func (s *mcpResponder) writeError(stdout io.Writer, id *json.RawMessage, code int, message string) error { - payload := map[string]any{ - "jsonrpc": "2.0", - "error": mcpError{Code: code, Message: message}, - } - if id != nil { - payload["id"] = json.RawMessage(*id) - } else { - payload["id"] = nil - } - return s.writeMessage(stdout, payload) + return s.writeMessage(stdout, buildErrorMessage(id, code, message)) } func (s *mcpResponder) writeProgress(stdout io.Writer, token json.RawMessage, progress float64, total float64, message string) error { - return s.writeMessage(stdout, map[string]any{ - "jsonrpc": "2.0", - "method": "notifications/progress", - "params": map[string]any{ - "progressToken": json.RawMessage(token), - "progress": progress, - "total": total, - "message": message, - }, - }) + return s.writeMessage(stdout, buildProgressMessage(token, progress, total, message)) } func (s *mcpResponder) writeMessage(stdout io.Writer, payload any) error { diff --git a/internal/cli/mcp_run.go b/internal/cli/mcp_run.go index a271317..4321b69 100644 --- a/internal/cli/mcp_run.go +++ b/internal/cli/mcp_run.go @@ -4,19 +4,29 @@ import ( "bufio" "context" "encoding/json" + "errors" "flag" "fmt" "io" + "net/http" + "os" + "os/signal" "strings" + "syscall" + "time" - "github.com/devr-tools/codeguard/internal/version" service "github.com/devr-tools/codeguard/pkg/codeguard" ) func runServe(args []string, stdin io.Reader, stdout io.Writer, stderr io.Writer) int { fs := flag.NewFlagSet("serve", flag.ContinueOnError) fs.SetOutput(stderr) - mcpMode := fs.Bool("mcp", false, "serve an MCP server over stdio") + mcpMode := fs.Bool("mcp", false, "serve an MCP server") + httpMode := fs.Bool("http", false, "serve MCP over Streamable HTTP instead of stdio") + addr := fs.String("addr", "127.0.0.1:8080", "HTTP listen address (with --http)") + mcpPath := fs.String("mcp-path", mcpDefaultHTTPPath, "HTTP path for the MCP endpoint (with --http)") + authToken := fs.String("auth-token", "", "optional static bearer token required on HTTP requests; falls back to $CODEGUARD_MCP_AUTH_TOKEN") + authHeader := fs.String("auth-header", mcpDefaultAuthHeader, "HTTP header carrying the auth token (with --http)") configPath := fs.String("config", service.DefaultConfigPath(), "default config file or directory path") profile := fs.String("profile", "", "optional default policy profile override") if err := fs.Parse(args); err != nil { @@ -27,16 +37,33 @@ func runServe(args []string, stdin io.Reader, stdout io.Writer, stderr io.Writer return 1 } + tools := &mcpToolService{ + defaultConfigPath: *configPath, + defaultProfile: *profile, + } + + if *httpMode { + token := *authToken + if strings.TrimSpace(token) == "" { + token = strings.TrimSpace(os.Getenv("CODEGUARD_MCP_AUTH_TOKEN")) + } + auth := mcpAuthConfig{token: token, header: *authHeader} + if err := serveMCPHTTP(*addr, *mcpPath, tools, auth, stderr); err != nil { + _, _ = fmt.Fprintf(stderr, "mcp http server failed: %v\n", err) + return 1 + } + return 0 + } + server := mcpServer{ defaultConfigPath: *configPath, defaultProfile: *profile, active: map[string]context.CancelFunc{}, cancelled: map[string]bool{}, responder: &mcpResponder{}, - tools: &mcpToolService{ - defaultConfigPath: *configPath, - defaultProfile: *profile, - }, + tools: tools, + requester: newServerRequester(), + rootsCache: &rootsCache{}, } if err := server.serve(stdin, stdout); err != nil { _, _ = fmt.Fprintf(stderr, "mcp server failed: %v\n", err) @@ -45,6 +72,38 @@ func runServe(args []string, stdin io.Reader, stdout io.Writer, stderr io.Writer return 0 } +// serveMCPHTTP runs the Streamable-HTTP transport until interrupted, then drains +// in-flight requests gracefully. +func serveMCPHTTP(addr string, path string, tools *mcpToolService, auth mcpAuthConfig, stderr io.Writer) error { + handler := newMCPHTTPHandler(tools, auth, path) + srv := &http.Server{Addr: addr, Handler: handler} + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + errCh := make(chan error, 1) + go func() { + authNote := "auth disabled" + if auth.enabled() { + authNote = "auth required via " + auth.headerName() + } + _, _ = fmt.Fprintf(stderr, "codeguard MCP HTTP server listening on %s%s (%s)\n", addr, path, authNote) + errCh <- srv.ListenAndServe() + }() + + select { + case err := <-errCh: + if err != nil && !errors.Is(err, http.ErrServerClosed) { + return err + } + return nil + case <-ctx.Done(): + shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + return srv.Shutdown(shutdownCtx) + } +} + func (s *mcpServer) serve(stdin io.Reader, stdout io.Writer) error { scanner := bufio.NewScanner(stdin) scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) @@ -69,6 +128,12 @@ func (s *mcpServer) handleLine(line string, stdout io.Writer) error { if req.JSONRPC != "2.0" { return s.responder.writeError(stdout, req.idPtr(), -32600, "invalid request") } + // Responses to server-initiated requests (sampling/createMessage, + // roots/list) carry an id and no method; route them to the waiting caller. + if isResponseMessage(req) { + s.requester.deliver(decodeIDKey(req.ID), json.RawMessage(line)) + return nil + } return s.handleRequestMethod(req, stdout) } @@ -81,12 +146,18 @@ func (s *mcpServer) handleRequestMethod(req mcpRequest, stdout io.Writer) error case "notifications/cancelled": s.handleCancelledNotification(req.Params) return nil + case "notifications/roots/list_changed": + s.rootsCache.invalidate() + return nil case "ping": return s.responder.writeResult(stdout, req.ID, map[string]any{}) case "tools/list": return s.handleToolsList(req, stdout) case "tools/call": return s.handleToolsCallRequest(req, stdout) + case "resources/list", "resources/templates/list", "resources/read", + "prompts/list", "prompts/get", "logging/setLevel": + return s.handleSyncMethod(req, stdout) default: if len(req.ID) == 0 { return nil @@ -95,21 +166,25 @@ func (s *mcpServer) handleRequestMethod(req mcpRequest, stdout io.Writer) error } } +// handleSyncMethod serves the request methods that produce a single synchronous +// response (resources/*, prompts/*, logging/*) by delegating to the shared +// transport-neutral router and writing its envelope to stdout. +func (s *mcpServer) handleSyncMethod(req mcpRequest, stdout io.Writer) error { + if !s.isInitialized() { + return s.responder.writeError(stdout, req.idPtr(), -32002, "server not initialized") + } + msg, handled := s.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) + if !handled { + return s.responder.writeError(stdout, req.idPtr(), -32601, "method not found") + } + return s.responder.writeMessage(stdout, msg) +} + func (s *mcpServer) handleInitializeResponse(req mcpRequest, stdout io.Writer) error { s.mu.Lock() s.initializeSeen = true + s.clientCaps = parseClientCapabilities(req.Params) s.mu.Unlock() - return s.responder.writeResult(stdout, req.ID, map[string]any{ - "protocolVersion": negotiateMCPProtocolVersion(req.Params), - "capabilities": map[string]any{ - "tools": map[string]any{}, - }, - "serverInfo": map[string]any{ - "name": "codeguard", - "title": "CodeGuard MCP Server", - "version": version.Number, - }, - "instructions": "Use validate_patch before writing files to disk when you want policy feedback on a proposed diff.", - }) + return s.responder.writeResult(stdout, req.ID, buildInitializeResult(req.Params)) } diff --git a/internal/cli/mcp_tools.go b/internal/cli/mcp_tools.go index 8ce1310..380b9e6 100644 --- a/internal/cli/mcp_tools.go +++ b/internal/cli/mcp_tools.go @@ -3,12 +3,22 @@ package cli import ( "context" "encoding/json" + "errors" "fmt" + "os" + "path/filepath" "strings" + "sync" + "time" service "github.com/devr-tools/codeguard/pkg/codeguard" ) +// errConfigPathNotPermitted is returned when a caller-supplied config_path +// resolves outside the allowed roots. The message is intentionally generic so +// the HTTP transport does not become a filesystem oracle. +var errConfigPathNotPermitted = errors.New("config_path is not within an allowed root") + func (s *mcpToolService) callToolWithContext(ctx context.Context, raw json.RawMessage) (map[string]any, error) { var call struct { Name string `json:"name"` @@ -22,13 +32,19 @@ func (s *mcpToolService) callToolWithContext(ctx context.Context, raw json.RawMe case "scan": return s.callScan(ctx, normalizeMCPArguments(call.Arguments)) case "validate_config": - return s.callValidateConfig(normalizeMCPArguments(call.Arguments)) + return s.callValidateConfig(ctx, normalizeMCPArguments(call.Arguments)) case "validate_patch": return s.callValidatePatch(ctx, normalizeMCPArguments(call.Arguments)) case "explain": - return s.callExplain(normalizeMCPArguments(call.Arguments)) + return s.callExplain(ctx, normalizeMCPArguments(call.Arguments)) case "list_rules": - return s.callListRules(normalizeMCPArguments(call.Arguments)) + return s.callListRules(ctx, normalizeMCPArguments(call.Arguments)) + case "verify_fix": + return s.callVerifyFix(ctx, normalizeMCPArguments(call.Arguments)) + case "propose_fix": + return s.callProposeFix(ctx, normalizeMCPArguments(call.Arguments)) + case "apply_fix": + return s.callApplyFix(ctx, normalizeMCPArguments(call.Arguments)) default: return nil, fmt.Errorf("unknown tool: %s", call.Name) } @@ -45,7 +61,11 @@ func (s *mcpToolService) callScan(ctx context.Context, raw json.RawMessage) (map return nil, fmt.Errorf("invalid scan arguments") } - cfg, err := s.loadConfig(args.ConfigPath, args.Profile) + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) if err != nil { return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil } @@ -61,14 +81,28 @@ func (s *mcpToolService) callScan(ctx context.Context, raw json.RawMessage) (map baseRef = "main" } - report, err := service.RunWithOptions(ctx, cfg, service.ScanOptions{Mode: mode, BaseRef: baseRef}) + opts := service.ScanOptions{Mode: mode, BaseRef: baseRef} + if emit := progressFrom(ctx); emit != nil { + total := countEnabledSections(cfg, mode) + var mu sync.Mutex + var done float64 + opts.OnSectionComplete = func(section service.SectionResult) { + mu.Lock() + done++ + progress := done + mu.Unlock() + emit(progress, total, fmt.Sprintf("%s: %s (%d findings)", section.Name, section.Status, len(section.Findings))) + } + } + + report, err := service.RunWithOptions(ctx, cfg, opts) if err != nil { return toolErrorResult(fmt.Sprintf("scan failed: %v", err)), nil } return toolSuccessResult(report), nil } -func (s *mcpToolService) callValidateConfig(raw json.RawMessage) (map[string]any, error) { +func (s *mcpToolService) callValidateConfig(ctx context.Context, raw json.RawMessage) (map[string]any, error) { var args struct { ConfigPath string `json:"config_path"` Profile string `json:"profile"` @@ -77,7 +111,11 @@ func (s *mcpToolService) callValidateConfig(raw json.RawMessage) (map[string]any return nil, fmt.Errorf("invalid validate_config arguments") } - cfg, err := s.loadConfig(args.ConfigPath, args.Profile) + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) if err != nil { return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil } @@ -104,7 +142,11 @@ func (s *mcpToolService) callValidatePatch(ctx context.Context, raw json.RawMess return toolErrorResult("validate_patch requires a unified diff"), nil } - cfg, err := s.loadConfig(args.ConfigPath, args.Profile) + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) if err != nil { return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil } @@ -115,7 +157,7 @@ func (s *mcpToolService) callValidatePatch(ctx context.Context, raw json.RawMess return toolSuccessResult(report), nil } -func (s *mcpToolService) callExplain(raw json.RawMessage) (map[string]any, error) { +func (s *mcpToolService) callExplain(ctx context.Context, raw json.RawMessage) (map[string]any, error) { var args struct { ConfigPath string `json:"config_path"` Profile string `json:"profile"` @@ -128,7 +170,11 @@ func (s *mcpToolService) callExplain(raw json.RawMessage) (map[string]any, error return toolErrorResult("explain requires rule_id"), nil } - rule, ok, err := s.resolveExplainRule(args.ConfigPath, args.Profile, args.RuleID) + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + rule, ok, err := s.resolveExplainRule(confinedPath, args.Profile, args.RuleID) if err != nil { return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil } @@ -138,7 +184,7 @@ func (s *mcpToolService) callExplain(raw json.RawMessage) (map[string]any, error return toolSuccessResult(buildExplainAgentOutput(rule)), nil } -func (s *mcpToolService) callListRules(raw json.RawMessage) (map[string]any, error) { +func (s *mcpToolService) callListRules(ctx context.Context, raw json.RawMessage) (map[string]any, error) { var args struct { ConfigPath string `json:"config_path"` Profile string `json:"profile"` @@ -150,7 +196,11 @@ func (s *mcpToolService) callListRules(raw json.RawMessage) (map[string]any, err if strings.TrimSpace(args.ConfigPath) == "" && strings.TrimSpace(args.Profile) == "" { return toolSuccessResult(map[string]any{"rules": service.Rules()}), nil } - cfg, err := s.loadConfig(args.ConfigPath, args.Profile) + confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + if err != nil { + return toolErrorResult(err.Error()), nil + } + cfg, err := s.loadConfig(confinedPath, args.Profile) if err != nil { return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil } @@ -169,6 +219,104 @@ func (s *mcpToolService) loadConfig(configPath string, profile string) (service. return loadConfigWithProfile(path, overrideProfile) } +// confineConfigArg validates a caller-supplied config_path against the allowed +// roots and returns it resolved to an absolute path. An empty argument returns +// "" so the server's trusted default config is used. The server's own default +// config path is never passed through here, so a trusted out-of-tree --config +// keeps working. +func (s *mcpToolService) confineConfigArg(ctx context.Context, configPath string) (string, error) { + candidate := strings.TrimSpace(configPath) + if candidate == "" { + return "", nil + } + return confinePath(s.allowedRoots(ctx), candidate) +} + +// rootsFetchTimeout bounds the server→client roots/list round trip done while +// resolving a config path, so a config load never blocks on a slow client. +const rootsFetchTimeout = 10 * time.Second + +// allowedRoots is the set of directories a caller-supplied config_path may live +// under: the directory of the server's default config, the working directory, +// any roots advertised by the connected client (via the roots capability), and +// any roots injected explicitly (tests). +func (s *mcpToolService) allowedRoots(ctx context.Context) []string { + roots := make([]string, 0, 4) + if def := strings.TrimSpace(s.defaultConfigPath); def != "" { + roots = append(roots, configDirOf(def)) + } + if wd, err := os.Getwd(); err == nil { + roots = append(roots, wd) + } + if caller := clientCallerFrom(ctx); caller != nil && caller.supports("roots") { + rctx, cancel := context.WithTimeout(ctx, rootsFetchTimeout) + if clientRoots, err := caller.listRoots(rctx); err == nil { + for _, root := range clientRoots { + if p := rootURIToPath(root.URI); p != "" { + roots = append(roots, p) + } + } + } + cancel() + } + roots = append(roots, clientRootsFrom(ctx)...) + return roots +} + +// rootURIToPath converts a roots entry to a filesystem path. MCP roots are +// file:// URIs; a bare path is accepted as-is. +func rootURIToPath(uri string) string { + uri = strings.TrimSpace(uri) + if uri == "" { + return "" + } + if strings.HasPrefix(uri, "file://") { + return strings.TrimPrefix(uri, "file://") + } + if strings.Contains(uri, "://") { + return "" // non-file scheme is not a local path + } + return uri +} + +// configDirOf returns the directory that should anchor confinement for a config +// path: the path itself if it is a directory, otherwise its parent. +func configDirOf(path string) string { + if info, err := os.Stat(path); err == nil && info.IsDir() { + return path + } + return filepath.Dir(path) +} + +// confinePath resolves candidate to an absolute, cleaned path and requires it to +// sit within one of allowedRoots. Symlinks are resolved best-effort to defeat +// link escapes; on resolution failure it falls back to the cleaned path. +func confinePath(allowedRoots []string, candidate string) (string, error) { + abs := resolvePath(candidate) + for _, root := range allowedRoots { + if strings.TrimSpace(root) == "" { + continue + } + rootAbs := resolvePath(root) + if abs == rootAbs || strings.HasPrefix(abs, rootAbs+string(os.PathSeparator)) { + return abs, nil + } + } + return "", errConfigPathNotPermitted +} + +func resolvePath(path string) string { + abs, err := filepath.Abs(path) + if err != nil { + return filepath.Clean(path) + } + abs = filepath.Clean(abs) + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + return resolved + } + return abs +} + func (s *mcpToolService) resolveExplainRule(configPath string, profile string, ruleID string) (service.RuleMetadata, bool, error) { if strings.TrimSpace(configPath) != "" { cfg, err := s.loadConfig(configPath, profile) diff --git a/internal/cli/mcp_types.go b/internal/cli/mcp_types.go index c48451f..7258e83 100644 --- a/internal/cli/mcp_types.go +++ b/internal/cli/mcp_types.go @@ -15,12 +15,15 @@ type mcpServer struct { defaultConfigPath string defaultProfile string initializeSeen bool + clientCaps map[string]any mu sync.Mutex active map[string]context.CancelFunc cancelled map[string]bool wg sync.WaitGroup responder *mcpResponder tools *mcpToolService + requester *serverRequester + rootsCache *rootsCache } type mcpResponder struct { @@ -50,4 +53,5 @@ type mcpTool struct { Description string `json:"description,omitempty"` InputSchema map[string]any `json:"inputSchema"` OutputSchema map[string]any `json:"outputSchema,omitempty"` + Annotations map[string]any `json:"annotations,omitempty"` } diff --git a/internal/codeguard/core/rule_scan_pack_types.go b/internal/codeguard/core/rule_scan_pack_types.go index c681e8e..5a2e715 100644 --- a/internal/codeguard/core/rule_scan_pack_types.go +++ b/internal/codeguard/core/rule_scan_pack_types.go @@ -13,6 +13,11 @@ type ScanOptions struct { DiffText string EnableAI bool EnableFix bool + // OnSectionComplete, when set, is invoked once per section as soon as that + // section finishes, enabling callers (e.g. the MCP server) to stream + // partial results. It is never serialized — json.Marshal errors on a + // non-nil func field, so the json:"-" tag is required. + OnSectionComplete func(SectionResult) `json:"-"` } type RulePackConfig struct { diff --git a/internal/codeguard/runner/support/findings.go b/internal/codeguard/runner/support/findings.go index 3d5b428..5db0a7d 100644 --- a/internal/codeguard/runner/support/findings.go +++ b/internal/codeguard/runner/support/findings.go @@ -122,6 +122,9 @@ func FinalizeSection(sc Context, id string, name string, findings []core.Finding } } section.Findings = active + if sc.Opts.OnSectionComplete != nil { + sc.Opts.OnSectionComplete(section) + } return section } diff --git a/internal/codeguard/runner/support/patch.go b/internal/codeguard/runner/support/patch.go index 9458d82..67d4b2f 100644 --- a/internal/codeguard/runner/support/patch.go +++ b/internal/codeguard/runner/support/patch.go @@ -64,6 +64,23 @@ func MaterializePatchedTargets(cfg core.Config, diffText string) (core.Config, m return patched, diffCommand, cleanup, nil } +// ApplyUnifiedDiff applies diffText to each configured target in place, rebasing +// per target the same way MaterializePatchedTargets does for verification. It is +// used to write a verified fix to the working tree. Targets whose rebased diff +// is empty are skipped. +func ApplyUnifiedDiff(cfg core.Config, diffText string) error { + for _, target := range cfg.Targets { + targetDiff := strings.TrimSpace(RebaseUnifiedDiff(diffText, DiffPrefixForTarget(target.Path))) + if targetDiff == "" { + continue + } + if err := applyUnifiedDiff(target.Path, targetDiff+"\n"); err != nil { + return fmt.Errorf("apply patch for target %q: %w", target.Name, err) + } + } + return nil +} + func applyUnifiedDiff(dir string, diffText string) error { cmd := exec.Command("git", "apply", "--recount", "--whitespace=nowarn") cmd.Dir = dir diff --git a/tests/mcp/http_test.go b/tests/mcp/http_test.go new file mode 100644 index 0000000..53f7d15 --- /dev/null +++ b/tests/mcp/http_test.go @@ -0,0 +1,269 @@ +package mcp_test + +import ( + "encoding/json" + "io" + "net" + "net/http" + "os" + "os/exec" + "strings" + "testing" + "time" + + "github.com/devr-tools/codeguard/internal/cli" +) + +// The HTTP transport is exercised end-to-end by launching the real `serve --mcp +// --http` binary in a subprocess (mirroring the stdio smoke harness) and +// driving it over HTTP, keeping all tests in the external mcp_test package. + +const httpTestToken = "http-smoke-token" + +func TestMCPServeHTTPHelperProcess(t *testing.T) { + if os.Getenv("GO_WANT_MCP_HTTP_HELPER_PROCESS") != "1" { + return + } + args := []string{"serve", "--mcp", "--http", "--addr", os.Getenv("CODEGUARD_TEST_HTTP_ADDR")} + if token := os.Getenv("CODEGUARD_TEST_HTTP_TOKEN"); token != "" { + args = append(args, "--auth-token", token) + } + if cfg := os.Getenv("CODEGUARD_TEST_HTTP_CONFIG"); cfg != "" { + args = append(args, "-config", cfg) + } + os.Exit(cli.Run(args, os.Stdin, os.Stdout, os.Stderr)) +} + +func TestMCPServeHTTP(t *testing.T) { + openBase := startHTTPServer(t, "") + authBase := startHTTPServer(t, httpTestToken) + + t.Run("initialize-capabilities", func(t *testing.T) { + resp, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}`) + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", resp.StatusCode, body) + } + if resp.Header.Get("Mcp-Session-Id") == "" { + t.Fatalf("expected Mcp-Session-Id header on initialize") + } + var out struct { + Result struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]any `json:"capabilities"` + } `json:"result"` + } + decodeLine(t, body, &out) + if out.Result.ProtocolVersion != "2025-06-18" { + t.Fatalf("unexpected protocol version: %s", out.Result.ProtocolVersion) + } + for _, capability := range []string{"tools", "resources", "prompts", "logging"} { + if _, ok := out.Result.Capabilities[capability]; !ok { + t.Fatalf("missing capability %q: %#v", capability, out.Result.Capabilities) + } + } + }) + + t.Run("tools-list-annotations", func(t *testing.T) { + _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`) + var out struct { + Result struct { + Tools []struct { + Name string `json:"name"` + Annotations struct { + ReadOnlyHint bool `json:"readOnlyHint"` + } `json:"annotations"` + } `json:"tools"` + } `json:"result"` + } + decodeLine(t, body, &out) + if len(out.Result.Tools) != 8 { + t.Fatalf("expected 8 tools, got %d", len(out.Result.Tools)) + } + for _, tool := range out.Result.Tools { + if tool.Name == "apply_fix" { + continue // apply_fix is the one destructive tool + } + if !tool.Annotations.ReadOnlyHint { + t.Fatalf("tool %q missing readOnlyHint", tool.Name) + } + } + }) + + t.Run("tool-call-streams-sse", func(t *testing.T) { + resp, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"list_rules","arguments":{},"_meta":{"progressToken":"p1"}}}`) + if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") { + t.Fatalf("expected SSE content type, got %q", ct) + } + var sawProgress, sawResult bool + for _, frame := range strings.Split(body, "\n\n") { + frame = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(frame), "data:")) + if frame == "" { + continue + } + var msg struct { + Method string `json:"method"` + Result struct { + IsError bool `json:"isError"` + } `json:"result"` + } + decodeLine(t, frame, &msg) + if msg.Method == "notifications/progress" { + sawProgress = true + } + if msg.Method == "" && !msg.Result.IsError { + sawResult = true + } + } + if !sawProgress || !sawResult { + t.Fatalf("expected progress + result over SSE (progress=%v result=%v): %s", sawProgress, sawResult, body) + } + }) + + t.Run("resource-read", func(t *testing.T) { + _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"codeguard://rules"}}`) + var out struct { + Result struct { + Contents []struct { + URI string `json:"uri"` + Text string `json:"text"` + } `json:"contents"` + } `json:"result"` + } + decodeLine(t, body, &out) + if len(out.Result.Contents) == 0 || out.Result.Contents[0].URI != "codeguard://rules" { + t.Fatalf("unexpected resources/read response: %s", body) + } + var payload struct { + Rules []json.RawMessage `json:"rules"` + } + if err := json.Unmarshal([]byte(out.Result.Contents[0].Text), &payload); err != nil || len(payload.Rules) == 0 { + t.Fatalf("expected rule catalog in resource text, got err=%v len=%d", err, len(payload.Rules)) + } + }) + + t.Run("prompt-get", func(t *testing.T) { + _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":5,"method":"prompts/get","params":{"name":"review-diff","arguments":{"diff":"--- a\n+++ b\n"}}}`) + if !strings.Contains(body, "validate_patch") { + t.Fatalf("expected review-diff prompt to mention validate_patch: %s", body) + } + }) + + t.Run("health-and-get-stream-needs-session", func(t *testing.T) { + resp, body := httpGet(t, openBase+"/healthz") + if resp.StatusCode != http.StatusOK || !strings.Contains(body, "ok") { + t.Fatalf("unexpected health response: %d %s", resp.StatusCode, body) + } + // GET /mcp is now the server→client SSE stream; without a session id it + // has nothing to attach to and returns 404. + resp, _ = httpGet(t, openBase+"/mcp") + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("expected 404 for GET /mcp without session, got %d", resp.StatusCode) + } + }) + + t.Run("auth-enforcement", func(t *testing.T) { + initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}` + cases := []struct { + name string + header map[string]string + want int + }{ + {"missing", nil, http.StatusUnauthorized}, + {"wrong", map[string]string{"Authorization": "Bearer nope"}, http.StatusUnauthorized}, + {"correct", map[string]string{"Authorization": "Bearer " + httpTestToken}, http.StatusOK}, + {"correct-no-scheme", map[string]string{"Authorization": httpTestToken}, http.StatusOK}, + } + for _, tc := range cases { + resp, body := mcpPost(t, authBase, tc.header, initBody) + if resp.StatusCode != tc.want { + t.Fatalf("%s: expected %d, got %d body=%s", tc.name, tc.want, resp.StatusCode, body) + } + } + }) +} + +// startHTTPServer launches the serve --mcp --http binary on a free port and +// waits for /healthz, returning the base URL. The subprocess is killed on test +// cleanup. +func startHTTPServer(t *testing.T, token string) string { + t.Helper() + return startHTTPServerWithConfig(t, token, "") +} + +func startHTTPServerWithConfig(t *testing.T, token string, configPath string) string { + t.Helper() + addr := freeTCPAddr(t) + cmd := exec.Command(os.Args[0], "-test.run=TestMCPServeHTTPHelperProcess") + cmd.Env = append(os.Environ(), + "GO_WANT_MCP_HTTP_HELPER_PROCESS=1", + "CODEGUARD_TEST_HTTP_ADDR="+addr, + "CODEGUARD_TEST_HTTP_TOKEN="+token, + "CODEGUARD_TEST_HTTP_CONFIG="+configPath, + ) + if err := cmd.Start(); err != nil { + t.Fatalf("start http helper: %v", err) + } + t.Cleanup(func() { + _ = cmd.Process.Kill() + _ = cmd.Wait() + }) + + base := "http://" + addr + deadline := time.Now().Add(10 * time.Second) + for time.Now().Before(deadline) { + resp, err := http.Get(base + "/healthz") + if err == nil { + _ = resp.Body.Close() + if resp.StatusCode == http.StatusOK { + return base + } + } + time.Sleep(25 * time.Millisecond) + } + t.Fatalf("http server at %s did not become ready", base) + return "" +} + +func freeTCPAddr(t *testing.T) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("reserve port: %v", err) + } + addr := ln.Addr().String() + _ = ln.Close() + return addr +} + +func mcpPost(t *testing.T, base string, header map[string]string, body string) (*http.Response, string) { + t.Helper() + req, err := http.NewRequest(http.MethodPost, base+"/mcp", strings.NewReader(body)) + if err != nil { + t.Fatalf("new request: %v", err) + } + req.Header.Set("Content-Type", "application/json") + for k, v := range header { + req.Header.Set(k, v) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("do request: %v", err) + } + data, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + if err != nil { + t.Fatalf("read body: %v", err) + } + return resp, string(data) +} + +func httpGet(t *testing.T, url string) (*http.Response, string) { + t.Helper() + resp, err := http.Get(url) + if err != nil { + t.Fatalf("get %s: %v", url, err) + } + data, _ := io.ReadAll(resp.Body) + _ = resp.Body.Close() + return resp, string(data) +} diff --git a/tests/mcp/sampling_test.go b/tests/mcp/sampling_test.go new file mode 100644 index 0000000..006e176 --- /dev/null +++ b/tests/mcp/sampling_test.go @@ -0,0 +1,332 @@ +package mcp_test + +import ( + "bufio" + "encoding/json" + "io" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" +) + +// These tests exercise server→client requests end-to-end: the test acts as the +// MCP client, advertising sampling/roots and answering the server's +// server-initiated requests. propose_fix's verification is expected to fail on +// the throwaway diff — the point is to prove the bidirectional machinery and the +// sampling generator fire, not to land a verified patch. + +const sampleDiff = "--- a/main.go\n+++ b/main.go\n@@ -1 +1 @@\n-bad\n+good\n" + +func writeFixtureConfig(t *testing.T, dir string) string { + t.Helper() + configPath := filepath.Join(dir, "codeguard.json") + body := `{ + "name": "mcp-sampling-test", + "targets": [{"name": "repo", "path": "` + dir + `", "language": "go"}], + "checks": {"quality": true, "design": false, "security": false, "prompts": false, "ci": false}, + "output": {"format": "json"} +}` + if err := os.WriteFile(configPath, []byte(body), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + return configPath +} + +func samplingResponse(t *testing.T, id json.RawMessage) string { + t.Helper() + resp, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "result": map[string]any{ + "role": "assistant", + "content": map[string]any{"type": "text", "text": sampleDiff}, + "model": "test-model", + }, + }) + if err != nil { + t.Fatalf("marshal sampling response: %v", err) + } + return string(resp) +} + +// TestMCPStdioSampling drives propose_fix over stdio and answers the server's +// sampling/createMessage request. +func TestMCPStdioSampling(t *testing.T) { + dir := t.TempDir() + cfg := writeFixtureConfig(t, dir) + + cmd := exec.Command(os.Args[0], "-test.run=TestMCPServeHelperProcess", "--", cfg) + cmd.Env = append(os.Environ(), "GO_WANT_MCP_HELPER_PROCESS=1") + stdin, err := cmd.StdinPipe() + if err != nil { + t.Fatalf("stdin pipe: %v", err) + } + stdout, err := cmd.StdoutPipe() + if err != nil { + t.Fatalf("stdout pipe: %v", err) + } + if err := cmd.Start(); err != nil { + t.Fatalf("start: %v", err) + } + t.Cleanup(func() { + _ = stdin.Close() + _ = cmd.Process.Kill() + _ = cmd.Wait() + }) + + write := func(line string) { + if _, err := io.WriteString(stdin, line+"\n"); err != nil { + t.Fatalf("write: %v", err) + } + } + write(`{"jsonrpc":"2.0","id":"init","method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{"sampling":{}}}}`) + write(`{"jsonrpc":"2.0","method":"notifications/initialized"}`) + write(`{"jsonrpc":"2.0","id":"fix","method":"tools/call","params":{"name":"propose_fix","arguments":{"finding":{"rule_id":"demo.rule","message":"bad value","path":"main.go","line":1}}}}`) + + scanner := bufio.NewScanner(stdout) + scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) + + type outcome struct { + sawSampling bool + sawFix bool + } + result := make(chan outcome, 1) + go func() { + var out outcome + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + var msg struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal([]byte(line), &msg) != nil { + continue + } + if msg.Method == "sampling/createMessage" { + out.sawSampling = true + write(samplingResponse(t, msg.ID)) + continue + } + if strings.TrimSpace(string(msg.ID)) == `"fix"` { + out.sawFix = true + result <- out + return + } + } + result <- out + }() + + select { + case out := <-result: + if !out.sawSampling { + t.Fatalf("server did not issue sampling/createMessage") + } + if !out.sawFix { + t.Fatalf("propose_fix never returned a result") + } + case <-time.After(30 * time.Second): + t.Fatalf("timed out waiting for sampling round trip") + } +} + +// TestMCPHTTPSampling drives propose_fix over HTTP and answers the server's +// sampling request over the GET SSE stream. +func TestMCPHTTPSampling(t *testing.T) { + dir := t.TempDir() + cfg := writeFixtureConfig(t, dir) + base := startHTTPServerWithConfig(t, "", cfg) + + // initialize, advertising sampling, to obtain a session id. + resp, _ := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{"sampling":{}}}}`) + session := resp.Header.Get("Mcp-Session-Id") + if session == "" { + t.Fatalf("no session id returned") + } + + // Open the server→client SSE stream and wait for the readiness comment. + streamReq, _ := http.NewRequest(http.MethodGet, base+"/mcp", nil) + streamReq.Header.Set("Mcp-Session-Id", session) + streamResp, err := http.DefaultClient.Do(streamReq) + if err != nil { + t.Fatalf("open stream: %v", err) + } + t.Cleanup(func() { _ = streamResp.Body.Close() }) + streamReader := bufio.NewReader(streamResp.Body) + if _, err := streamReader.ReadString('\n'); err != nil { // ": ready" + t.Fatalf("read stream readiness: %v", err) + } + + // Fire propose_fix; it blocks server-side until we answer the sampling request. + fixDone := make(chan string, 1) + go func() { + _, body := mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, + `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"propose_fix","arguments":{"finding":{"rule_id":"demo.rule","message":"bad value","path":"main.go","line":1}}}}`) + fixDone <- body + }() + + // Read the sampling request off the stream and answer it on a POST. + sawSampling := false + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) && !sawSampling { + line, err := streamReader.ReadString('\n') + if err != nil { + t.Fatalf("read stream: %v", err) + } + line = strings.TrimSpace(line) + data, ok := strings.CutPrefix(line, "data:") + if !ok { + continue + } + var msg struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) != nil { + continue + } + if msg.Method == "sampling/createMessage" { + sawSampling = true + mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, samplingResponse(t, msg.ID)) + } + } + if !sawSampling { + t.Fatalf("server did not issue sampling/createMessage over the stream") + } + + select { + case <-fixDone: + // propose_fix returned (verification may fail; the round trip is what we assert). + case <-time.After(30 * time.Second): + t.Fatalf("propose_fix did not complete after sampling answer") + } +} + +// TestMCPHTTPRootsConfinement proves the server fetches client roots and uses +// them to permit a config_path that is otherwise outside the allowed roots. +func TestMCPHTTPRootsConfinement(t *testing.T) { + dir := t.TempDir() + cfg := writeFixtureConfig(t, dir) + base := startHTTPServer(t, "") // default config; cfg lives in an out-of-tree temp dir + + resp, _ := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{"roots":{}}}}`) + session := resp.Header.Get("Mcp-Session-Id") + if session == "" { + t.Fatalf("no session id returned") + } + + streamReq, _ := http.NewRequest(http.MethodGet, base+"/mcp", nil) + streamReq.Header.Set("Mcp-Session-Id", session) + streamResp, err := http.DefaultClient.Do(streamReq) + if err != nil { + t.Fatalf("open stream: %v", err) + } + t.Cleanup(func() { _ = streamResp.Body.Close() }) + streamReader := bufio.NewReader(streamResp.Body) + if _, err := streamReader.ReadString('\n'); err != nil { + t.Fatalf("read readiness: %v", err) + } + + // validate_config with an out-of-tree config_path; the server must fetch + // roots to decide whether it is permitted. + validateDone := make(chan string, 1) + go func() { + req := `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"validate_config","arguments":{"config_path":"` + cfg + `"}}}` + _, body := mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, req) + validateDone <- body + }() + + // Answer the server's roots/list request, advertising the temp dir as a root. + sawRoots := false + deadline := time.Now().Add(20 * time.Second) + for time.Now().Before(deadline) && !sawRoots { + line, err := streamReader.ReadString('\n') + if err != nil { + t.Fatalf("read stream: %v", err) + } + data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") + if !ok { + continue + } + var msg struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) != nil { + continue + } + if msg.Method == "roots/list" { + sawRoots = true + rootsResp, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": msg.ID, + "result": map[string]any{"roots": []map[string]any{{"uri": "file://" + dir, "name": "temp"}}}, + }) + mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, string(rootsResp)) + } + } + if !sawRoots { + t.Fatalf("server did not issue roots/list") + } + + select { + case body := <-validateDone: + // With the temp dir advertised as a root, confinement must permit the + // path — so the response must not be the "not within an allowed root" error. + if strings.Contains(body, "not within an allowed root") { + t.Fatalf("config_path was rejected despite being an advertised root: %s", body) + } + case <-time.After(20 * time.Second): + t.Fatalf("validate_config did not complete after roots answer") + } + + // A second config_path call must reuse the cached roots — no new roots/list. + secondDone := make(chan string, 1) + go func() { + req := `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"validate_config","arguments":{"config_path":"` + cfg + `"}}}` + _, body := mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, req) + secondDone <- body + }() + + secondRoots := make(chan struct{}, 1) + go func() { + for { + line, err := streamReader.ReadString('\n') + if err != nil { + return + } + data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") + if !ok { + continue + } + var msg struct { + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) == nil && msg.Method == "roots/list" { + secondRoots <- struct{}{} + return + } + } + }() + + select { + case body := <-secondDone: + if strings.Contains(body, "not within an allowed root") { + t.Fatalf("cached-roots config_path was rejected: %s", body) + } + case <-time.After(20 * time.Second): + t.Fatalf("second validate_config did not complete") + } + select { + case <-secondRoots: + t.Fatalf("server issued a second roots/list instead of using the cache") + case <-time.After(500 * time.Millisecond): + // no second roots/list — cache hit, as expected + } +} diff --git a/tests/mcp/smoke_assertions_test.go b/tests/mcp/smoke_assertions_test.go index f3d6987..9596dac 100644 --- a/tests/mcp/smoke_assertions_test.go +++ b/tests/mcp/smoke_assertions_test.go @@ -2,6 +2,7 @@ package mcp_test import ( "encoding/json" + "strings" "testing" ) @@ -27,14 +28,25 @@ func assertCurrentDiscovery(t *testing.T, lines []string) { var listResp struct { Result struct { Tools []struct { - Name string `json:"name"` + Name string `json:"name"` + Annotations struct { + ReadOnlyHint bool `json:"readOnlyHint"` + } `json:"annotations"` } `json:"tools"` } `json:"result"` } decodeLine(t, toolsLine, &listResp) - if !containsTool(listResp.Result.Tools, "list_rules") || !containsTool(listResp.Result.Tools, "validate_patch") { + if !containsTool(toolNames(listResp.Result.Tools), "list_rules") || !containsTool(toolNames(listResp.Result.Tools), "validate_patch") { t.Fatalf("unexpected tool catalog: %#v", listResp.Result.Tools) } + for _, tool := range listResp.Result.Tools { + if tool.Name == "apply_fix" { + continue // apply_fix is the one destructive tool + } + if !tool.Annotations.ReadOnlyHint { + t.Fatalf("expected tool %q to carry readOnlyHint annotation: %#v", tool.Name, listResp.Result.Tools) + } + } var explainResp struct { Result struct { @@ -127,6 +139,159 @@ func assertScanProgressCancel(t *testing.T, lines []string) { } } +func assertResourcesDiscovery(t *testing.T, lines []string) { + t.Helper() + if len(lines) != 3 { + t.Fatalf("expected 3 responses, got %d: %q", len(lines), lines) + } + listLine := findResponseLineByID(t, lines, `"res-list"`) + readLine := findResponseLineByID(t, lines, `"res-read"`) + + var listResp struct { + Result struct { + Resources []struct { + URI string `json:"uri"` + } `json:"resources"` + } `json:"result"` + } + decodeLine(t, listLine, &listResp) + foundRules := false + for _, res := range listResp.Result.Resources { + if res.URI == "codeguard://rules" { + foundRules = true + } + } + if !foundRules { + t.Fatalf("expected codeguard://rules in resources/list: %#v", listResp.Result.Resources) + } + + var readResp struct { + Result struct { + Contents []struct { + URI string `json:"uri"` + MIMEType string `json:"mimeType"` + Text string `json:"text"` + } `json:"contents"` + } `json:"result"` + } + decodeLine(t, readLine, &readResp) + if len(readResp.Result.Contents) == 0 || readResp.Result.Contents[0].URI != "codeguard://rules" { + t.Fatalf("unexpected resources/read response: %#v", readResp) + } + var payload struct { + Rules []json.RawMessage `json:"rules"` + } + if err := json.Unmarshal([]byte(readResp.Result.Contents[0].Text), &payload); err != nil || len(payload.Rules) == 0 { + t.Fatalf("expected rule catalog in resource text, got err=%v payload=%#v", err, payload) + } +} + +func assertPromptsDiscovery(t *testing.T, lines []string) { + t.Helper() + if len(lines) != 3 { + t.Fatalf("expected 3 responses, got %d: %q", len(lines), lines) + } + listLine := findResponseLineByID(t, lines, `"prompts-list"`) + getLine := findResponseLineByID(t, lines, `"prompts-get"`) + + var listResp struct { + Result struct { + Prompts []struct { + Name string `json:"name"` + } `json:"prompts"` + } `json:"result"` + } + decodeLine(t, listLine, &listResp) + if !containsTool(toolNames(listResp.Result.Prompts), "review-diff") { + t.Fatalf("expected review-diff in prompts/list: %#v", listResp.Result.Prompts) + } + + var getResp struct { + Result struct { + Messages []struct { + Role string `json:"role"` + } `json:"messages"` + } `json:"result"` + } + decodeLine(t, getLine, &getResp) + if len(getResp.Result.Messages) == 0 { + t.Fatalf("expected prompt messages, got %#v", getResp) + } +} + +func assertScanStreaming(t *testing.T, lines []string) { + t.Helper() + sectionProgress := 0 + sawResult := false + for _, line := range lines { + var envelope struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + Params struct { + ProgressToken string `json:"progressToken"` + Message string `json:"message"` + } `json:"params"` + } + decodeLine(t, line, &envelope) + if envelope.Method == "notifications/progress" && envelope.Params.ProgressToken == "stream-tok" { + // Per-section messages look like "Code Quality: pass (0 findings)". + if strings.Contains(envelope.Params.Message, ":") && strings.Contains(envelope.Params.Message, "findings") { + sectionProgress++ + } + } + if strings.TrimSpace(string(envelope.ID)) == `"scan-1"` { + sawResult = true + } + } + if sectionProgress < 2 { + t.Fatalf("expected at least 2 per-section progress notifications, got %d: %q", sectionProgress, lines) + } + if !sawResult { + t.Fatalf("expected final scan result, got %q", lines) + } +} + +func assertVerifyFixFailsClosed(t *testing.T, lines []string) { + t.Helper() + line := findResponseLineByID(t, lines, `"vf-1"`) + var resp struct { + Result struct { + IsError bool `json:"isError"` + StructuredContent struct { + Verified bool `json:"verified"` + } `json:"structuredContent"` + } `json:"result"` + } + decodeLine(t, line, &resp) + if !resp.Result.IsError { + t.Fatalf("expected verify_fix to fail closed on a bogus diff: %s", line) + } + if resp.Result.StructuredContent.Verified { + t.Fatalf("expected structuredContent.verified=false on a failed fix: %s", line) + } +} + +func assertApplyFixFailsClosed(t *testing.T, lines []string) { + t.Helper() + line := findResponseLineByID(t, lines, `"af-1"`) + var resp struct { + Result struct { + IsError bool `json:"isError"` + StructuredContent struct { + Applied bool `json:"applied"` + Verified bool `json:"verified"` + } `json:"structuredContent"` + } `json:"result"` + } + decodeLine(t, line, &resp) + if !resp.Result.IsError { + t.Fatalf("expected apply_fix to fail closed on a bogus diff: %s", line) + } + if resp.Result.StructuredContent.Applied || resp.Result.StructuredContent.Verified { + t.Fatalf("expected apply_fix not applied / not verified on failure: %s", line) + } +} + func scanResponseMatchesID(raw *json.RawMessage, want int) bool { if raw == nil { return false diff --git a/tests/mcp/smoke_helpers_test.go b/tests/mcp/smoke_helpers_test.go index bf7f0b7..31b23ab 100644 --- a/tests/mcp/smoke_helpers_test.go +++ b/tests/mcp/smoke_helpers_test.go @@ -95,6 +95,20 @@ func setupCancelableScanConfig(t *testing.T, dir string) map[string]string { return map[string]string{"__CONFIG_PATH__": configPath} } +func setupStreamingConfig(t *testing.T, dir string) map[string]string { + t.Helper() + configPath := filepath.Join(dir, "codeguard.json") + if err := os.WriteFile(configPath, []byte(`{ + "name": "mcp-streaming-test", + "targets": [{"name": "repo", "path": "`+dir+`", "language": "go"}], + "checks": {"quality": true, "design": true, "security": false, "prompts": false, "ci": true}, + "output": {"format": "json"} +}`), 0o644); err != nil { + t.Fatalf("write config: %v", err) + } + return map[string]string{"__CONFIG_PATH__": configPath} +} + func loadTranscript(t *testing.T, rel string, replacements map[string]string) string { t.Helper() data, err := os.ReadFile(rel) @@ -173,13 +187,28 @@ func findResponseLineByID(t *testing.T, lines []string, want string) string { return "" } -func containsTool(tools []struct { - Name string `json:"name"` -}, name string) bool { - for _, tool := range tools { - if tool.Name == name { +func containsTool(names []string, name string) bool { + for _, candidate := range names { + if candidate == name { return true } } return false } + +func toolNames[T any](tools []T) []string { + names := make([]string, 0, len(tools)) + for _, tool := range tools { + data, err := json.Marshal(tool) + if err != nil { + continue + } + var named struct { + Name string `json:"name"` + } + if json.Unmarshal(data, &named) == nil { + names = append(names, named.Name) + } + } + return names +} diff --git a/tests/mcp/smoke_test.go b/tests/mcp/smoke_test.go index 2d238f3..79b06d6 100644 --- a/tests/mcp/smoke_test.go +++ b/tests/mcp/smoke_test.go @@ -38,6 +38,36 @@ func TestMCPHostSmokeProfiles(t *testing.T) { setup: setupCancelableScanConfig, assertion: assertScanProgressCancel, }, + { + name: "resource-agent", + transcript: "testdata/transcripts/resources_discovery.jsonl", + setup: setupPromptConfig, + assertion: assertResourcesDiscovery, + }, + { + name: "prompt-agent", + transcript: "testdata/transcripts/prompts_discovery.jsonl", + setup: setupPromptConfig, + assertion: assertPromptsDiscovery, + }, + { + name: "streaming-agent", + transcript: "testdata/transcripts/scan_streaming.jsonl", + setup: setupStreamingConfig, + assertion: assertScanStreaming, + }, + { + name: "verify-fix-agent", + transcript: "testdata/transcripts/verify_fix_failclosed.jsonl", + setup: setupPromptConfig, + assertion: assertVerifyFixFailsClosed, + }, + { + name: "apply-fix-agent", + transcript: "testdata/transcripts/apply_fix_failclosed.jsonl", + setup: setupPromptConfig, + assertion: assertApplyFixFailsClosed, + }, } for _, profile := range profiles { diff --git a/tests/mcp/testdata/transcripts/apply_fix_failclosed.jsonl b/tests/mcp/testdata/transcripts/apply_fix_failclosed.jsonl new file mode 100644 index 0000000..044ce8a --- /dev/null +++ b/tests/mcp/testdata/transcripts/apply_fix_failclosed.jsonl @@ -0,0 +1,3 @@ +{"jsonrpc":"2.0","id":"init-1","method":"initialize","params":{"protocolVersion":"2025-11-25"}} +{"jsonrpc":"2.0","method":"notifications/initialized"} +{"jsonrpc":"2.0","id":"af-1","method":"tools/call","params":{"name":"apply_fix","arguments":{"diff":"garbage not a diff","finding":{"rule_id":"demo.rule"}}}} diff --git a/tests/mcp/testdata/transcripts/prompts_discovery.jsonl b/tests/mcp/testdata/transcripts/prompts_discovery.jsonl new file mode 100644 index 0000000..e14c3c1 --- /dev/null +++ b/tests/mcp/testdata/transcripts/prompts_discovery.jsonl @@ -0,0 +1,4 @@ +{"jsonrpc":"2.0","id":"init-1","method":"initialize","params":{"protocolVersion":"2025-11-25","clientInfo":{"name":"Prompt Agent","version":"1.0.0"}}} +{"jsonrpc":"2.0","method":"notifications/initialized"} +{"jsonrpc":"2.0","id":"prompts-list","method":"prompts/list"} +{"jsonrpc":"2.0","id":"prompts-get","method":"prompts/get","params":{"name":"explain-rule","arguments":{"rule_id":"security.hardcoded-secret"}}} diff --git a/tests/mcp/testdata/transcripts/resources_discovery.jsonl b/tests/mcp/testdata/transcripts/resources_discovery.jsonl new file mode 100644 index 0000000..8ea9215 --- /dev/null +++ b/tests/mcp/testdata/transcripts/resources_discovery.jsonl @@ -0,0 +1,4 @@ +{"jsonrpc":"2.0","id":"init-1","method":"initialize","params":{"protocolVersion":"2025-11-25","clientInfo":{"name":"Resource Agent","version":"1.0.0"}}} +{"jsonrpc":"2.0","method":"notifications/initialized"} +{"jsonrpc":"2.0","id":"res-list","method":"resources/list"} +{"jsonrpc":"2.0","id":"res-read","method":"resources/read","params":{"uri":"codeguard://rules"}} diff --git a/tests/mcp/testdata/transcripts/scan_streaming.jsonl b/tests/mcp/testdata/transcripts/scan_streaming.jsonl new file mode 100644 index 0000000..2fa5607 --- /dev/null +++ b/tests/mcp/testdata/transcripts/scan_streaming.jsonl @@ -0,0 +1,3 @@ +{"jsonrpc":"2.0","id":"init-1","method":"initialize","params":{"protocolVersion":"2025-11-25"}} +{"jsonrpc":"2.0","method":"notifications/initialized"} +{"jsonrpc":"2.0","id":"scan-1","method":"tools/call","params":{"name":"scan","arguments":{"mode":"full"},"_meta":{"progressToken":"stream-tok"}}} diff --git a/tests/mcp/testdata/transcripts/verify_fix_failclosed.jsonl b/tests/mcp/testdata/transcripts/verify_fix_failclosed.jsonl new file mode 100644 index 0000000..f19aba0 --- /dev/null +++ b/tests/mcp/testdata/transcripts/verify_fix_failclosed.jsonl @@ -0,0 +1,3 @@ +{"jsonrpc":"2.0","id":"init-1","method":"initialize","params":{"protocolVersion":"2025-11-25"}} +{"jsonrpc":"2.0","method":"notifications/initialized"} +{"jsonrpc":"2.0","id":"vf-1","method":"tools/call","params":{"name":"verify_fix","arguments":{"diff":"garbage not a diff","finding":{"rule_id":"demo.rule"}}}} From 2efd47c9bd701f5508bb53f33ce3e5af37d0b03d Mon Sep 17 00:00:00 2001 From: Alex Wilkerson John Date: Thu, 18 Jun 2026 20:40:07 -0400 Subject: [PATCH 2/2] feat: mcp server features --- internal/cli/mcp_client.go | 353 ----------------------- internal/cli/mcp_client_bridge.go | 95 ++++++ internal/cli/mcp_client_context.go | 17 ++ internal/cli/mcp_client_progress.go | 58 ++++ internal/cli/mcp_client_requester.go | 119 ++++++++ internal/cli/mcp_client_roots.go | 36 +++ internal/cli/mcp_client_types.go | 29 ++ internal/cli/mcp_dispatch.go | 50 ++-- internal/cli/mcp_dispatch_content.go | 43 +++ internal/cli/mcp_fix.go | 317 ++++---------------- internal/cli/mcp_fix_common.go | 68 +++++ internal/cli/mcp_fix_sampling.go | 175 +++++++++++ internal/cli/mcp_http.go | 325 --------------------- internal/cli/mcp_http_handler.go | 65 +++++ internal/cli/mcp_http_io.go | 50 ++++ internal/cli/mcp_http_registry.go | 61 ++++ internal/cli/mcp_http_session.go | 57 ---- internal/cli/mcp_http_stream.go | 205 +++++++++++++ internal/cli/mcp_protocol.go | 269 ----------------- internal/cli/mcp_protocol_version.go | 28 ++ internal/cli/mcp_run.go | 30 +- internal/cli/mcp_tool_catalog.go | 169 +++++++++++ internal/cli/mcp_tool_config.go | 102 +++++++ internal/cli/mcp_tool_dispatch.go | 37 +++ internal/cli/mcp_tool_results.go | 41 +++ internal/cli/mcp_tool_rules.go | 29 ++ internal/cli/mcp_tool_schemas.go | 52 ++++ internal/cli/mcp_tools.go | 176 +---------- internal/codeguard/ai/fix/generator.go | 21 +- tests/mcp/http_helpers_test.go | 156 ++++++++++ tests/mcp/http_test.go | 130 +-------- tests/mcp/sampling_helpers_test.go | 119 ++++++++ tests/mcp/sampling_stdio_helpers_test.go | 92 ++++++ tests/mcp/sampling_test.go | 201 +------------ 34 files changed, 1985 insertions(+), 1790 deletions(-) create mode 100644 internal/cli/mcp_client_bridge.go create mode 100644 internal/cli/mcp_client_context.go create mode 100644 internal/cli/mcp_client_progress.go create mode 100644 internal/cli/mcp_client_requester.go create mode 100644 internal/cli/mcp_client_roots.go create mode 100644 internal/cli/mcp_client_types.go create mode 100644 internal/cli/mcp_dispatch_content.go create mode 100644 internal/cli/mcp_fix_common.go create mode 100644 internal/cli/mcp_fix_sampling.go create mode 100644 internal/cli/mcp_http_handler.go create mode 100644 internal/cli/mcp_http_io.go create mode 100644 internal/cli/mcp_http_registry.go create mode 100644 internal/cli/mcp_http_stream.go create mode 100644 internal/cli/mcp_protocol_version.go create mode 100644 internal/cli/mcp_tool_catalog.go create mode 100644 internal/cli/mcp_tool_config.go create mode 100644 internal/cli/mcp_tool_dispatch.go create mode 100644 internal/cli/mcp_tool_results.go create mode 100644 internal/cli/mcp_tool_rules.go create mode 100644 internal/cli/mcp_tool_schemas.go create mode 100644 tests/mcp/http_helpers_test.go create mode 100644 tests/mcp/sampling_helpers_test.go create mode 100644 tests/mcp/sampling_stdio_helpers_test.go diff --git a/internal/cli/mcp_client.go b/internal/cli/mcp_client.go index fbf9132..e1ff809 100644 --- a/internal/cli/mcp_client.go +++ b/internal/cli/mcp_client.go @@ -1,360 +1,7 @@ package cli -import ( - "context" - "encoding/json" - "fmt" - "strings" - "sync" - "time" - - service "github.com/devr-tools/codeguard/pkg/codeguard" -) - -// serverRequestTimeout bounds how long the server waits for a client to answer -// a server-initiated request (sampling/createMessage, roots/list). -const serverRequestTimeout = 120 * time.Second - -// clientCaller is the server→client capability surface used by tools: it reports -// the client's advertised capabilities and issues server-initiated requests -// (sampling, roots). Implemented per transport by clientBridge. -type clientCaller interface { - supports(capability string) bool - sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) - listRoots(ctx context.Context) ([]mcpRoot, error) - elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) -} - -// elicitResult is the client's answer to an elicitation/create request. -type elicitResult struct { - Action string `json:"action"` // "accept" | "decline" | "cancel" - Content json.RawMessage `json:"content"` -} - -func (e elicitResult) accepted() bool { return e.Action == "accept" } - -type mcpRoot struct { - URI string `json:"uri"` - Name string `json:"name,omitempty"` -} - -type clientCallerCtxKey struct{} - -func withClientCaller(ctx context.Context, c clientCaller) context.Context { - if c == nil { - return ctx - } - return context.WithValue(ctx, clientCallerCtxKey{}, c) -} - -func clientCallerFrom(ctx context.Context) clientCaller { - c, _ := ctx.Value(clientCallerCtxKey{}).(clientCaller) - return c -} - -// serverRequester correlates server-initiated requests with the client's -// responses. The transport writes the outbound request (via the send closure -// passed to call) and feeds inbound responses back through deliver. -type serverRequester struct { - mu sync.Mutex - pending map[string]chan json.RawMessage - counter int -} - -func newServerRequester() *serverRequester { - return &serverRequester{pending: map[string]chan json.RawMessage{}} -} - -// call issues a server-initiated request: it allocates an id, lets send write -// the request to the transport, and waits for the matching response (or a -// timeout / context cancellation). -func (r *serverRequester) call(ctx context.Context, send func(id string) error) (json.RawMessage, error) { - r.mu.Lock() - r.counter++ - id := fmt.Sprintf("srv-%d", r.counter) - ch := make(chan json.RawMessage, 1) - r.pending[id] = ch - r.mu.Unlock() - defer func() { - r.mu.Lock() - delete(r.pending, id) - r.mu.Unlock() - }() - - if err := send(id); err != nil { - return nil, err - } - - ctx, cancel := context.WithTimeout(ctx, serverRequestTimeout) - defer cancel() - select { - case resp := <-ch: - return parseServerResponse(resp) - case <-ctx.Done(): - return nil, ctx.Err() - } -} - -// deliver routes an inbound response to a waiting call. It returns false when no -// server-initiated request is awaiting that id (so the caller can dispatch the -// message normally). -func (r *serverRequester) deliver(id string, raw json.RawMessage) bool { - r.mu.Lock() - ch, ok := r.pending[id] - r.mu.Unlock() - if !ok { - return false - } - select { - case ch <- raw: - default: - } - return true -} - -func parseServerResponse(raw json.RawMessage) (json.RawMessage, error) { - var env struct { - Result json.RawMessage `json:"result"` - Error *struct { - Code int `json:"code"` - Message string `json:"message"` - } `json:"error"` - } - if err := json.Unmarshal(raw, &env); err != nil { - return nil, err - } - if env.Error != nil { - return nil, fmt.Errorf("client error %d: %s", env.Error.Code, env.Error.Message) - } - return env.Result, nil -} - -// isResponseMessage reports whether an inbound JSON-RPC message is a response to -// a server-initiated request: it carries an id and no method. -func isResponseMessage(req mcpRequest) bool { - return len(req.ID) > 0 && strings.TrimSpace(req.Method) == "" -} - -// decodeIDKey canonicalizes a JSON-RPC id into the string key used by -// serverRequester (which issues string ids like "srv-1"). -func decodeIDKey(raw json.RawMessage) string { - var s string - if json.Unmarshal(raw, &s) == nil { - return s - } - return strings.TrimSpace(string(raw)) -} - -// parseClientCapabilities extracts the capabilities object the client advertised -// during initialize, so the server knows whether sampling/roots are available. -func parseClientCapabilities(raw json.RawMessage) map[string]any { - var params struct { - Capabilities map[string]any `json:"capabilities"` - } - if err := json.Unmarshal(raw, ¶ms); err != nil { - return nil - } - return params.Capabilities -} - -// rootsCache memoizes the client's roots per connection so a config-path check -// does not issue a fresh roots/list round trip every time. It is invalidated on -// notifications/roots/list_changed. -type rootsCache struct { - mu sync.Mutex - loaded bool - roots []mcpRoot -} - -func (c *rootsCache) invalidate() { - c.mu.Lock() - c.loaded = false - c.roots = nil - c.mu.Unlock() -} - -// load returns the cached roots, fetching once via fetch on a miss. The lock is -// held across fetch so concurrent callers coalesce into a single round trip. -func (c *rootsCache) load(fetch func() ([]mcpRoot, error)) ([]mcpRoot, error) { - c.mu.Lock() - defer c.mu.Unlock() - if c.loaded { - return c.roots, nil - } - roots, err := fetch() - if err != nil { - return nil, err - } - c.roots = roots - c.loaded = true - return roots, nil -} - -// clientBridge is the transport-agnostic clientCaller. Each transport supplies -// the client capabilities, a serverRequester, a send closure that writes a -// server→client message over that transport, and a per-connection roots cache. -type clientBridge struct { - caps map[string]any - requester *serverRequester - send func(payload any) error - roots *rootsCache -} - -func (b *clientBridge) supports(capability string) bool { - if b == nil || b.caps == nil { - return false - } - _, ok := b.caps[capability] - return ok -} - -func (b *clientBridge) sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) { - if !b.supports("sampling") { - return nil, fmt.Errorf("client does not support sampling") - } - return b.requester.call(ctx, func(id string) error { - return b.send(map[string]any{ - "jsonrpc": "2.0", - "id": id, - "method": "sampling/createMessage", - "params": params, - }) - }) -} - -func (b *clientBridge) listRoots(ctx context.Context) ([]mcpRoot, error) { - if !b.supports("roots") { - return nil, fmt.Errorf("client does not support roots") - } - if b.roots != nil { - return b.roots.load(func() ([]mcpRoot, error) { return b.fetchRoots(ctx) }) - } - return b.fetchRoots(ctx) -} - -func (b *clientBridge) elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) { - if !b.supports("elicitation") { - return elicitResult{}, fmt.Errorf("client does not support elicitation") - } - raw, err := b.requester.call(ctx, func(id string) error { - return b.send(map[string]any{ - "jsonrpc": "2.0", - "id": id, - "method": "elicitation/create", - "params": map[string]any{ - "message": message, - "requestedSchema": schema, - }, - }) - }) - if err != nil { - return elicitResult{}, err - } - var result elicitResult - if err := json.Unmarshal(raw, &result); err != nil { - return elicitResult{}, err - } - return result, nil -} - -func (b *clientBridge) fetchRoots(ctx context.Context) ([]mcpRoot, error) { - raw, err := b.requester.call(ctx, func(id string) error { - return b.send(map[string]any{ - "jsonrpc": "2.0", - "id": id, - "method": "roots/list", - "params": map[string]any{}, - }) - }) - if err != nil { - return nil, err - } - var result struct { - Roots []mcpRoot `json:"roots"` - } - if err := json.Unmarshal(raw, &result); err != nil { - return nil, err - } - return result.Roots, nil -} - // mcp_client.go holds the per-call transport callbacks that the tool layer // reaches through context: a progress emitter (for streaming partial findings) // and, added in a later phase, a client caller (for sampling/roots). Threading // them via context keeps the many tool method signatures unchanged, and both // are nil-safe so tools work when a transport does not supply them. - -type progressFunc func(progress float64, total float64, message string) - -type progressCtxKey struct{} - -func withProgress(ctx context.Context, fn progressFunc) context.Context { - if fn == nil { - return ctx - } - return context.WithValue(ctx, progressCtxKey{}, fn) -} - -func progressFrom(ctx context.Context) progressFunc { - fn, _ := ctx.Value(progressCtxKey{}).(progressFunc) - return fn -} - -type clientRootsCtxKey struct{} - -// withClientRoots attaches the filesystem roots the connected client advertised -// (via the roots capability) so config-path confinement can permit them. -func withClientRoots(ctx context.Context, roots []string) context.Context { - if len(roots) == 0 { - return ctx - } - return context.WithValue(ctx, clientRootsCtxKey{}, roots) -} - -func clientRootsFrom(ctx context.Context) []string { - roots, _ := ctx.Value(clientRootsCtxKey{}).([]string) - return roots -} - -// countEnabledSections returns a best-effort count of the scan sections that -// will run for a config, used as the `total` for per-section progress. It is a -// hint, not a guarantee — if it drifts from the runner the progress bar simply -// may not land exactly on 100%. -func countEnabledSections(cfg service.Config, mode service.ScanMode) float64 { - count := 0 - if cfg.Checks.Quality { - count++ - } - if cfg.Checks.Design { - count++ - } - if cfg.Checks.Security { - count++ - } - if cfg.Checks.Prompts { - count++ - } - if cfg.Checks.CI { - count++ - } - if cfg.Checks.SupplyChain { - count++ - } - if cfg.Checks.Contracts != nil { - if *cfg.Checks.Contracts { - count++ - } - } else if mode == service.ScanModeDiff { - count++ - } - for _, pack := range cfg.RulePacks { - if len(pack.Rules) > 0 { - count++ - break - } - } - if count == 0 { - count = 1 - } - return float64(count) -} diff --git a/internal/cli/mcp_client_bridge.go b/internal/cli/mcp_client_bridge.go new file mode 100644 index 0000000..13ee48a --- /dev/null +++ b/internal/cli/mcp_client_bridge.go @@ -0,0 +1,95 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" +) + +// clientBridge is the transport-agnostic clientCaller. Each transport supplies +// the client capabilities, a serverRequester, a send closure that writes a +// server→client message over that transport, and a per-connection roots cache. +type clientBridge struct { + caps map[string]any + requester *serverRequester + send func(payload any) error + roots *rootsCache +} + +func (b *clientBridge) supports(capability string) bool { + if b == nil || b.caps == nil { + return false + } + _, ok := b.caps[capability] + return ok +} + +func (b *clientBridge) sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) { + if !b.supports("sampling") { + return nil, fmt.Errorf("client does not support sampling") + } + return b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "sampling/createMessage", + "params": params, + }) + }) +} + +func (b *clientBridge) listRoots(ctx context.Context) ([]mcpRoot, error) { + if !b.supports("roots") { + return nil, fmt.Errorf("client does not support roots") + } + if b.roots != nil { + return b.roots.load(func() ([]mcpRoot, error) { return b.fetchRoots(ctx) }) + } + return b.fetchRoots(ctx) +} + +func (b *clientBridge) elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) { + if !b.supports("elicitation") { + return elicitResult{}, fmt.Errorf("client does not support elicitation") + } + raw, err := b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "elicitation/create", + "params": map[string]any{ + "message": message, + "requestedSchema": schema, + }, + }) + }) + if err != nil { + return elicitResult{}, err + } + var result elicitResult + if err := json.Unmarshal(raw, &result); err != nil { + return elicitResult{}, err + } + return result, nil +} + +func (b *clientBridge) fetchRoots(ctx context.Context) ([]mcpRoot, error) { + raw, err := b.requester.call(ctx, func(id string) error { + return b.send(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": "roots/list", + "params": map[string]any{}, + }) + }) + if err != nil { + return nil, err + } + var result struct { + Roots []mcpRoot `json:"roots"` + } + if err := json.Unmarshal(raw, &result); err != nil { + return nil, err + } + return result.Roots, nil +} diff --git a/internal/cli/mcp_client_context.go b/internal/cli/mcp_client_context.go new file mode 100644 index 0000000..0a41b8d --- /dev/null +++ b/internal/cli/mcp_client_context.go @@ -0,0 +1,17 @@ +package cli + +import "context" + +type clientCallerCtxKey struct{} + +func withClientCaller(ctx context.Context, c clientCaller) context.Context { + if c == nil { + return ctx + } + return context.WithValue(ctx, clientCallerCtxKey{}, c) +} + +func clientCallerFrom(ctx context.Context) clientCaller { + c, _ := ctx.Value(clientCallerCtxKey{}).(clientCaller) + return c +} diff --git a/internal/cli/mcp_client_progress.go b/internal/cli/mcp_client_progress.go new file mode 100644 index 0000000..33b0eb1 --- /dev/null +++ b/internal/cli/mcp_client_progress.go @@ -0,0 +1,58 @@ +package cli + +import ( + "context" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +type progressFunc func(progress float64, total float64, message string) + +type progressCtxKey struct{} + +func withProgress(ctx context.Context, fn progressFunc) context.Context { + if fn == nil { + return ctx + } + return context.WithValue(ctx, progressCtxKey{}, fn) +} + +func progressFrom(ctx context.Context) progressFunc { + fn, _ := ctx.Value(progressCtxKey{}).(progressFunc) + return fn +} + +// countEnabledSections returns a best-effort count of the scan sections that +// will run for a config, used as the `total` for per-section progress. +func countEnabledSections(cfg service.Config, mode service.ScanMode) float64 { + enabled := []bool{ + cfg.Checks.Quality, + cfg.Checks.Design, + cfg.Checks.Security, + cfg.Checks.Prompts, + cfg.Checks.CI, + cfg.Checks.SupplyChain, + cfg.Checks.Contracts != nil && *cfg.Checks.Contracts, + cfg.Checks.Contracts == nil && mode == service.ScanModeDiff, + hasRulePackRules(cfg), + } + count := 0 + for _, ok := range enabled { + if ok { + count++ + } + } + if count == 0 { + count = 1 + } + return float64(count) +} + +func hasRulePackRules(cfg service.Config) bool { + for _, pack := range cfg.RulePacks { + if len(pack.Rules) > 0 { + return true + } + } + return false +} diff --git a/internal/cli/mcp_client_requester.go b/internal/cli/mcp_client_requester.go new file mode 100644 index 0000000..cd38b51 --- /dev/null +++ b/internal/cli/mcp_client_requester.go @@ -0,0 +1,119 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" + "sync" + "time" +) + +// serverRequestTimeout bounds how long the server waits for a client to answer +// a server-initiated request (sampling/createMessage, roots/list). +const serverRequestTimeout = 120 * time.Second + +// serverRequester correlates server-initiated requests with the client's +// responses. The transport writes the outbound request (via the send closure +// passed to call) and feeds inbound responses back through deliver. +type serverRequester struct { + mu sync.Mutex + pending map[string]chan json.RawMessage + counter int +} + +func newServerRequester() *serverRequester { + return &serverRequester{pending: map[string]chan json.RawMessage{}} +} + +// call issues a server-initiated request: it allocates an id, lets send write +// the request to the transport, and waits for the matching response (or a +// timeout / context cancellation). +func (r *serverRequester) call(ctx context.Context, send func(id string) error) (json.RawMessage, error) { + r.mu.Lock() + r.counter++ + id := fmt.Sprintf("srv-%d", r.counter) + ch := make(chan json.RawMessage, 1) + r.pending[id] = ch + r.mu.Unlock() + defer func() { + r.mu.Lock() + delete(r.pending, id) + r.mu.Unlock() + }() + + if err := send(id); err != nil { + return nil, err + } + + ctx, cancel := context.WithTimeout(ctx, serverRequestTimeout) + defer cancel() + select { + case resp := <-ch: + return parseServerResponse(resp) + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +// deliver routes an inbound response to a waiting call. It returns false when no +// server-initiated request is awaiting that id (so the caller can dispatch the +// message normally). +func (r *serverRequester) deliver(id string, raw json.RawMessage) bool { + r.mu.Lock() + ch, ok := r.pending[id] + r.mu.Unlock() + if !ok { + return false + } + select { + case ch <- raw: + default: + } + return true +} + +func parseServerResponse(raw json.RawMessage) (json.RawMessage, error) { + var env struct { + Result json.RawMessage `json:"result"` + Error *struct { + Code int `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if err := json.Unmarshal(raw, &env); err != nil { + return nil, err + } + if env.Error != nil { + return nil, fmt.Errorf("client error %d: %s", env.Error.Code, env.Error.Message) + } + return env.Result, nil +} + +// isResponseMessage reports whether an inbound JSON-RPC message is a response to +// a server-initiated request: it carries an id and no method. +func isResponseMessage(req mcpRequest) bool { + return len(req.ID) > 0 && strings.TrimSpace(req.Method) == "" +} + +// decodeIDKey canonicalizes a JSON-RPC id into the string key used by +// serverRequester (which issues string ids like "srv-1"). +func decodeIDKey(raw json.RawMessage) string { + var s string + if json.Unmarshal(raw, &s) == nil { + return s + } + return strings.TrimSpace(string(raw)) +} + +// parseClientCapabilities extracts the capabilities object the client advertised +// during initialize, so the server knows whether sampling/roots are available. +func parseClientCapabilities(raw json.RawMessage) map[string]any { + var params struct { + Capabilities map[string]any `json:"capabilities"` + } + if err := json.Unmarshal(raw, ¶ms); err != nil { + return nil + } + return params.Capabilities +} diff --git a/internal/cli/mcp_client_roots.go b/internal/cli/mcp_client_roots.go new file mode 100644 index 0000000..82618a1 --- /dev/null +++ b/internal/cli/mcp_client_roots.go @@ -0,0 +1,36 @@ +package cli + +import "sync" + +// rootsCache memoizes the client's roots per connection so a config-path check +// does not issue a fresh roots/list round trip every time. It is invalidated on +// notifications/roots/list_changed. +type rootsCache struct { + mu sync.Mutex + loaded bool + roots []mcpRoot +} + +func (c *rootsCache) invalidate() { + c.mu.Lock() + c.loaded = false + c.roots = nil + c.mu.Unlock() +} + +// load returns the cached roots, fetching once via fetch on a miss. The lock is +// held across fetch so concurrent callers coalesce into a single round trip. +func (c *rootsCache) load(fetch func() ([]mcpRoot, error)) ([]mcpRoot, error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.loaded { + return c.roots, nil + } + roots, err := fetch() + if err != nil { + return nil, err + } + c.roots = roots + c.loaded = true + return roots, nil +} diff --git a/internal/cli/mcp_client_types.go b/internal/cli/mcp_client_types.go new file mode 100644 index 0000000..3b800dc --- /dev/null +++ b/internal/cli/mcp_client_types.go @@ -0,0 +1,29 @@ +package cli + +import ( + "context" + "encoding/json" +) + +// clientCaller is the server→client capability surface used by tools: it reports +// the client's advertised capabilities and issues server-initiated requests +// (sampling, roots). Implemented per transport by clientBridge. +type clientCaller interface { + supports(capability string) bool + sampleMessage(ctx context.Context, params map[string]any) (json.RawMessage, error) + listRoots(ctx context.Context) ([]mcpRoot, error) + elicit(ctx context.Context, message string, schema map[string]any) (elicitResult, error) +} + +// elicitResult is the client's answer to an elicitation/create request. +type elicitResult struct { + Action string `json:"action"` // "accept" | "decline" | "cancel" + Content json.RawMessage `json:"content"` +} + +func (e elicitResult) accepted() bool { return e.Action == "accept" } + +type mcpRoot struct { + URI string `json:"uri"` + Name string `json:"name,omitempty"` +} diff --git a/internal/cli/mcp_dispatch.go b/internal/cli/mcp_dispatch.go index 9b16d59..5ab6da3 100644 --- a/internal/cli/mcp_dispatch.go +++ b/internal/cli/mcp_dispatch.go @@ -88,36 +88,32 @@ func buildInitializeResult(params json.RawMessage) map[string]any { // Notifications (notifications/initialized, notifications/cancelled) are not // handled here — they carry no id and are transport-specific. func (s *mcpToolService) dispatchSyncMethod(method string, id json.RawMessage, params json.RawMessage) (map[string]any, bool) { - switch method { - case "initialize": - return buildResultMessage(id, buildInitializeResult(params)), true - case "ping": - return buildResultMessage(id, map[string]any{}), true - case "tools/list": - return buildResultMessage(id, map[string]any{"tools": mcpTools()}), true - case "resources/list": - return buildResultMessage(id, map[string]any{"resources": s.resourcesList()}), true - case "resources/templates/list": - return buildResultMessage(id, map[string]any{"resourceTemplates": resourceTemplates()}), true - case "resources/read": - result, errMsg := s.readResource(params) - if errMsg != "" { - return buildErrorMessage(ptrID(id), -32602, errMsg), true - } - return buildResultMessage(id, result), true - case "prompts/list": - return buildResultMessage(id, map[string]any{"prompts": mcpPrompts()}), true - case "prompts/get": - result, errMsg := getPrompt(params) - if errMsg != "" { - return buildErrorMessage(ptrID(id), -32602, errMsg), true - } - return buildResultMessage(id, result), true - case "logging/setLevel": + if msg, ok := staticSyncMethod(method, id, params); ok { + return msg, true + } + if msg, ok := s.dispatchResourceMethod(method, id, params); ok { + return msg, true + } + if msg, ok := dispatchPromptMethod(method, id, params); ok { + return msg, true + } + if method == "logging/setLevel" { return buildResultMessage(id, map[string]any{}), true - default: + } + return nil, false +} + +func staticSyncMethod(method string, id json.RawMessage, params json.RawMessage) (map[string]any, bool) { + handlers := map[string]func() map[string]any{ + "initialize": func() map[string]any { return buildResultMessage(id, buildInitializeResult(params)) }, + "ping": func() map[string]any { return buildResultMessage(id, map[string]any{}) }, + "tools/list": func() map[string]any { return buildResultMessage(id, map[string]any{"tools": mcpTools()}) }, + } + handler, ok := handlers[method] + if !ok { return nil, false } + return handler(), true } // ptrID returns a pointer to a copy of id, or nil when id is empty, for use with diff --git a/internal/cli/mcp_dispatch_content.go b/internal/cli/mcp_dispatch_content.go new file mode 100644 index 0000000..946b6ab --- /dev/null +++ b/internal/cli/mcp_dispatch_content.go @@ -0,0 +1,43 @@ +package cli + +import "encoding/json" + +func (s *mcpToolService) dispatchResourceMethod(method string, id json.RawMessage, params json.RawMessage) (map[string]any, bool) { + switch method { + case "resources/list": + return buildResultMessage(id, map[string]any{"resources": s.resourcesList()}), true + case "resources/templates/list": + return buildResultMessage(id, map[string]any{"resourceTemplates": resourceTemplates()}), true + case "resources/read": + return resourceReadMessage(s, id, params), true + default: + return nil, false + } +} + +func dispatchPromptMethod(method string, id json.RawMessage, params json.RawMessage) (map[string]any, bool) { + switch method { + case "prompts/list": + return buildResultMessage(id, map[string]any{"prompts": mcpPrompts()}), true + case "prompts/get": + return promptGetMessage(id, params), true + default: + return nil, false + } +} + +func resourceReadMessage(s *mcpToolService, id json.RawMessage, params json.RawMessage) map[string]any { + result, errMsg := s.readResource(params) + if errMsg != "" { + return buildErrorMessage(ptrID(id), -32602, errMsg) + } + return buildResultMessage(id, result) +} + +func promptGetMessage(id json.RawMessage, params json.RawMessage) map[string]any { + result, errMsg := getPrompt(params) + if errMsg != "" { + return buildErrorMessage(ptrID(id), -32602, errMsg) + } + return buildResultMessage(id, result) +} diff --git a/internal/cli/mcp_fix.go b/internal/cli/mcp_fix.go index 8c1e8f7..7fdc9f5 100644 --- a/internal/cli/mcp_fix.go +++ b/internal/cli/mcp_fix.go @@ -3,10 +3,10 @@ package cli import ( "context" "encoding/json" + "errors" "fmt" "strings" - internalfix "github.com/devr-tools/codeguard/internal/codeguard/ai/fix" runnersupport "github.com/devr-tools/codeguard/internal/codeguard/runner/support" service "github.com/devr-tools/codeguard/pkg/codeguard" ) @@ -42,59 +42,37 @@ func (a fixToolArgs) options() service.FixOptions { } func (s *mcpToolService) callVerifyFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { - var args fixToolArgs - if err := json.Unmarshal(raw, &args); err != nil { - return nil, fmt.Errorf("invalid verify_fix arguments") - } - if strings.TrimSpace(args.Diff) == "" { - return toolErrorResult("verify_fix requires a unified diff"), nil - } - - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + fixCtx, err := loadFixContext(ctx, s, raw) if err != nil { - return toolErrorResult(err.Error()), nil + if errors.Is(err, errInvalidFixArgs) { + return nil, fmt.Errorf("invalid verify_fix arguments") + } + return toolResultFromError(err), nil } - cfg, err := s.loadConfig(confinedPath, args.Profile) - if err != nil { - return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + if result, stop := requireFixDiff("verify_fix", fixCtx.args.Diff); stop { + return result, nil } - result, err := service.VerifyFix(ctx, cfg, args.Finding, service.FixCandidate{Diff: args.Diff}, args.options()) - if err != nil { - data := map[string]any{ - "verified": false, - "error": err.Error(), - "attempted_diff": args.Diff, - } - // Surface what still fails so an agent can iterate, by re-evaluating the - // diff against policy (does not mutate the working tree). - if report, perr := service.RunPatch(ctx, cfg, args.Diff); perr == nil { - data["remaining_findings"] = report - } - return toolErrorResultData(fmt.Sprintf("fix did not verify: %v", err), data), nil + result, failure, verifyErr, ok := verifyFixCandidate(ctx, fixCtx.cfg, fixCtx.args) + if ok { + return toolSuccessResult(result), nil } - return toolSuccessResult(result), nil + return toolErrorResultData(fmt.Sprintf("fix did not verify: %v", verifyErr), failure), nil } func (s *mcpToolService) callProposeFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { - var args fixToolArgs - if err := json.Unmarshal(raw, &args); err != nil { - return nil, fmt.Errorf("invalid propose_fix arguments") - } - if strings.TrimSpace(args.Finding.RuleID) == "" && strings.TrimSpace(args.Finding.Message) == "" { - return toolErrorResult("propose_fix requires a finding to fix"), nil - } - - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + fixCtx, err := loadFixContext(ctx, s, raw) if err != nil { - return toolErrorResult(err.Error()), nil + if errors.Is(err, errInvalidFixArgs) { + return nil, fmt.Errorf("invalid propose_fix arguments") + } + return toolResultFromError(err), nil } - cfg, err := s.loadConfig(confinedPath, args.Profile) - if err != nil { - return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + if result, stop := requireFixFinding(fixCtx.args.Finding.RuleID, fixCtx.args.Finding.Message); stop { + return result, nil } - generator, err := s.resolveFixGenerator(ctx, cfg) + generator, err := s.resolveFixGenerator(ctx, fixCtx.cfg) if err != nil { return toolErrorResult(fmt.Sprintf("initialize fix generator: %v", err)), nil } @@ -103,11 +81,11 @@ func (s *mcpToolService) callProposeFix(ctx context.Context, raw json.RawMessage } result, err := service.GenerateVerifiedFix(ctx, service.FixGenerateRequest{ - Config: cfg, - Finding: args.Finding, - Analysis: firstNonEmpty(args.Finding.Why, args.Finding.Message), + Config: fixCtx.cfg, + Finding: fixCtx.args.Finding, + Analysis: firstNonEmpty(fixCtx.args.Finding.Why, fixCtx.args.Finding.Message), Generator: generator, - Options: args.options(), + Options: fixCtx.args.options(), }) if err != nil { return toolErrorResultData(fmt.Sprintf("generate verified fix: %v", err), map[string]any{ @@ -118,238 +96,53 @@ func (s *mcpToolService) callProposeFix(ctx context.Context, raw json.RawMessage return toolSuccessResult(result), nil } -// callApplyFix verifies a candidate diff and, only if it passes, writes it to -// the working tree. When the client supports elicitation it first asks the user -// to confirm. This is the one tool that mutates the repo (destructiveHint). func (s *mcpToolService) callApplyFix(ctx context.Context, raw json.RawMessage) (map[string]any, error) { - var args fixToolArgs - if err := json.Unmarshal(raw, &args); err != nil { - return nil, fmt.Errorf("invalid apply_fix arguments") - } - if strings.TrimSpace(args.Diff) == "" { - return toolErrorResult("apply_fix requires a unified diff"), nil - } - - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + fixCtx, err := loadFixContext(ctx, s, raw) if err != nil { - return toolErrorResult(err.Error()), nil + if errors.Is(err, errInvalidFixArgs) { + return nil, fmt.Errorf("invalid apply_fix arguments") + } + return toolResultFromError(err), nil } - cfg, err := s.loadConfig(confinedPath, args.Profile) - if err != nil { - return toolErrorResult(fmt.Sprintf("load config: %v", err)), nil + if result, stop := requireFixDiff("apply_fix", fixCtx.args.Diff); stop { + return result, nil } - result, err := service.VerifyFix(ctx, cfg, args.Finding, service.FixCandidate{Diff: args.Diff}, args.options()) - if err != nil { - data := map[string]any{"verified": false, "applied": false, "error": err.Error(), "attempted_diff": args.Diff} - if report, perr := service.RunPatch(ctx, cfg, args.Diff); perr == nil { - data["remaining_findings"] = report - } - return toolErrorResultData(fmt.Sprintf("fix did not verify, not applied: %v", err), data), nil + verified, failure, verifyErr, ok := verifyFixCandidate(ctx, fixCtx.cfg, fixCtx.args) + if !ok { + return toolErrorResultData(fmt.Sprintf("fix did not verify, not applied: %v", verifyErr), failure), nil } - - // Confirm with the user before writing, when the client supports it. - if caller := clientCallerFrom(ctx); caller != nil && caller.supports("elicitation") { - accepted, err := confirmApply(ctx, caller, result.ChangedFiles) - if err != nil { - return toolErrorResult(fmt.Sprintf("confirmation failed: %v", err)), nil - } - if !accepted { - return toolSuccessResult(map[string]any{ - "applied": false, - "declined": true, - "diff": result.Diff, - "changed_files": result.ChangedFiles, - }), nil - } + if reply := confirmApplyResult(ctx, clientCallerFrom(ctx), verified.ChangedFiles, verified.Diff); reply != nil { + return reply, nil } - - if err := runnersupport.ApplyUnifiedDiff(cfg, result.Diff); err != nil { + if err := runnersupport.ApplyUnifiedDiff(fixCtx.cfg, verified.Diff); err != nil { return toolErrorResult(fmt.Sprintf("verified fix failed to apply to the working tree: %v", err)), nil } return toolSuccessResult(map[string]any{ "applied": true, - "diff": result.Diff, - "summary": result.Summary, - "changed_files": result.ChangedFiles, - "report": result.Report, - "test_results": result.TestResults, + "diff": verified.Diff, + "summary": verified.Summary, + "changed_files": verified.ChangedFiles, + "report": verified.Report, + "test_results": verified.TestResults, }), nil } -func confirmApply(ctx context.Context, caller clientCaller, changedFiles []string) (bool, error) { - message := fmt.Sprintf("Apply the verified fix to %d file(s): %s?", len(changedFiles), strings.Join(changedFiles, ", ")) - schema := map[string]any{ - "type": "object", - "properties": map[string]any{ - "apply": map[string]any{"type": "boolean", "description": "Apply the verified fix to the working tree"}, - }, - "required": []string{"apply"}, - } - res, err := caller.elicit(ctx, message, schema) - if err != nil { - return false, err - } - if !res.accepted() { - return false, nil - } - // If the host returned content, honor an explicit apply:false. - var content struct { - Apply *bool `json:"apply"` - } - if json.Unmarshal(res.Content, &content) == nil && content.Apply != nil { - return *content.Apply, nil - } - return true, nil -} - -// resolveFixGenerator picks a fix generator: the connected client's LLM via MCP -// sampling when the client supports it (no API key needed), otherwise a -// configured AI provider. Returns (nil, nil) when neither is available. -func (s *mcpToolService) resolveFixGenerator(ctx context.Context, cfg service.Config) (internalfix.Generator, error) { - if caller := clientCallerFrom(ctx); caller != nil && caller.supports("sampling") { - return samplingGenerator{caller: caller}, nil +func confirmApplyResult(ctx context.Context, caller clientCaller, changedFiles []string, diff string) map[string]any { + if caller == nil || !caller.supports("elicitation") { + return nil } - generator, available, err := internalfix.NewAIGenerator(cfg.AI) + accepted, err := confirmApply(ctx, caller, changedFiles) if err != nil { - return nil, err + return toolErrorResult(fmt.Sprintf("confirmation failed: %v", err)) } - if !available { - return nil, nil + if accepted { + return nil } - return generator, nil -} - -// samplingGenerator generates a fix candidate by asking the connected client's -// LLM through MCP sampling, so no server-side AI provider or API key is needed. -type samplingGenerator struct { - caller clientCaller -} - -func (g samplingGenerator) GenerateFix(ctx context.Context, input internalfix.GenerateInput) (internalfix.Candidate, error) { - params := map[string]any{ - "systemPrompt": "You are a precise code-fixing assistant. Return ONLY a valid unified diff (git format) that fixes the reported finding and changes nothing unrelated. Do not include any prose, explanation, or markdown fences.", - "messages": []map[string]any{{ - "role": "user", - "content": map[string]any{ - "type": "text", - "text": buildSamplingFixPrompt(input), - }, - }}, - "maxTokens": 2048, - // Ask the host to attach this server's context (the repo it is scanning) - // so the model can read the affected files, and bias model selection - // toward capability over speed for a correct patch. - "includeContext": "thisServer", - "modelPreferences": map[string]any{ - "intelligencePriority": 0.8, - "speedPriority": 0.3, - }, - } - raw, err := g.caller.sampleMessage(ctx, params) - if err != nil { - return internalfix.Candidate{}, err - } - text, err := samplingResultText(raw) - if err != nil { - return internalfix.Candidate{}, err - } - candidate := candidateFromText(text) - if strings.TrimSpace(candidate.Diff) == "" { - return internalfix.Candidate{}, fmt.Errorf("sampling response did not contain a diff") - } - return candidate, nil -} - -func buildSamplingFixPrompt(input internalfix.GenerateInput) string { - var b strings.Builder - b.WriteString("Fix this codeguard finding by producing a minimal unified diff.\n\n") - if input.Finding.RuleID != "" { - fmt.Fprintf(&b, "Rule: %s\n", input.Finding.RuleID) - } - if input.Finding.Path != "" { - fmt.Fprintf(&b, "File: %s", input.Finding.Path) - if input.Finding.Line > 0 { - fmt.Fprintf(&b, ":%d", input.Finding.Line) - } - b.WriteString("\n") - } - if input.Finding.Message != "" { - fmt.Fprintf(&b, "Issue: %s\n", input.Finding.Message) - } - if analysis := firstNonEmpty(input.Analysis, input.Finding.Why); analysis != "" { - fmt.Fprintf(&b, "Why: %s\n", analysis) - } - if input.Finding.HowToFix != "" { - fmt.Fprintf(&b, "How to fix: %s\n", input.Finding.HowToFix) - } - if input.Instructions != "" { - fmt.Fprintf(&b, "\n%s\n", input.Instructions) - } - b.WriteString("\nReturn only the unified diff.") - return b.String() -} - -// samplingResultText extracts the assistant text from a sampling/createMessage -// result, tolerating both a single content block and a content array. -func samplingResultText(raw json.RawMessage) (string, error) { - var single struct { - Content struct { - Text string `json:"text"` - } `json:"content"` - } - if err := json.Unmarshal(raw, &single); err == nil && strings.TrimSpace(single.Content.Text) != "" { - return single.Content.Text, nil - } - var multi struct { - Content []struct { - Type string `json:"type"` - Text string `json:"text"` - } `json:"content"` - } - if err := json.Unmarshal(raw, &multi); err == nil { - for _, block := range multi.Content { - if strings.TrimSpace(block.Text) != "" { - return block.Text, nil - } - } - } - return "", fmt.Errorf("sampling response had no text content") -} - -// candidateFromText turns a model reply into a fix candidate, accepting a raw -// diff, a {"summary","diff"} JSON object, or a fenced code block. -func candidateFromText(text string) internalfix.Candidate { - trimmed := strings.TrimSpace(text) - var asJSON internalfix.Candidate - if json.Unmarshal([]byte(trimmed), &asJSON) == nil && strings.TrimSpace(asJSON.Diff) != "" { - return asJSON - } - if fenced := extractFencedBlock(trimmed); fenced != "" { - return internalfix.Candidate{Diff: fenced} - } - if looksLikeDiff(trimmed) { - return internalfix.Candidate{Diff: trimmed} - } - return internalfix.Candidate{} -} - -func extractFencedBlock(text string) string { - start := strings.Index(text, "```") - if start < 0 { - return "" - } - rest := text[start+3:] - if nl := strings.IndexByte(rest, '\n'); nl >= 0 { - rest = rest[nl+1:] // drop the optional language tag line - } - end := strings.Index(rest, "```") - if end < 0 { - return strings.TrimSpace(rest) - } - return strings.TrimSpace(rest[:end]) -} - -func looksLikeDiff(text string) bool { - return strings.Contains(text, "@@") || strings.HasPrefix(text, "diff ") || strings.HasPrefix(text, "--- ") + return toolSuccessResult(map[string]any{ + "applied": false, + "declined": true, + "diff": diff, + "changed_files": changedFiles, + }) } diff --git a/internal/cli/mcp_fix_common.go b/internal/cli/mcp_fix_common.go new file mode 100644 index 0000000..1cfd992 --- /dev/null +++ b/internal/cli/mcp_fix_common.go @@ -0,0 +1,68 @@ +package cli + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strings" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +type fixContext struct { + args fixToolArgs + cfg service.Config +} + +var errInvalidFixArgs = errors.New("invalid fix arguments") + +func loadFixContext(ctx context.Context, s *mcpToolService, raw json.RawMessage) (fixContext, error) { + var args fixToolArgs + if err := json.Unmarshal(raw, &args); err != nil { + return fixContext{}, errInvalidFixArgs + } + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) + if err != nil { + return fixContext{}, err + } + cfg, err := s.loadConfig(confinedPath, args.Profile) + if err != nil { + return fixContext{}, fmt.Errorf("load config: %w", err) + } + return fixContext{args: args, cfg: cfg}, nil +} + +func requireFixDiff(name string, diff string) (map[string]any, bool) { + if strings.TrimSpace(diff) == "" { + return toolErrorResult(name + " requires a unified diff"), true + } + return nil, false +} + +func requireFixFinding(ruleID string, message string) (map[string]any, bool) { + if strings.TrimSpace(ruleID) == "" && strings.TrimSpace(message) == "" { + return toolErrorResult("propose_fix requires a finding to fix"), true + } + return nil, false +} + +func verifyFixCandidate(ctx context.Context, cfg service.Config, args fixToolArgs) (service.VerifiedFix, map[string]any, error, bool) { + result, err := service.VerifyFix(ctx, cfg, args.Finding, service.FixCandidate{Diff: args.Diff}, args.options()) + if err == nil { + return result, nil, nil, true + } + data := map[string]any{ + "verified": false, + "error": err.Error(), + "attempted_diff": args.Diff, + } + if report, perr := service.RunPatch(ctx, cfg, args.Diff); perr == nil { + data["remaining_findings"] = report + } + return service.VerifiedFix{}, data, err, false +} + +func toolResultFromError(err error) map[string]any { + return toolErrorResult(err.Error()) +} diff --git a/internal/cli/mcp_fix_sampling.go b/internal/cli/mcp_fix_sampling.go new file mode 100644 index 0000000..8733119 --- /dev/null +++ b/internal/cli/mcp_fix_sampling.go @@ -0,0 +1,175 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" + + internalfix "github.com/devr-tools/codeguard/internal/codeguard/ai/fix" + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +func confirmApply(ctx context.Context, caller clientCaller, changedFiles []string) (bool, error) { + message := fmt.Sprintf("Apply the verified fix to %d file(s): %s?", len(changedFiles), strings.Join(changedFiles, ", ")) + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "apply": map[string]any{"type": "boolean", "description": "Apply the verified fix to the working tree"}, + }, + "required": []string{"apply"}, + } + res, err := caller.elicit(ctx, message, schema) + if err != nil { + return false, err + } + if !res.accepted() { + return false, nil + } + var content struct { + Apply *bool `json:"apply"` + } + if json.Unmarshal(res.Content, &content) == nil && content.Apply != nil { + return *content.Apply, nil + } + return true, nil +} + +func (s *mcpToolService) resolveFixGenerator(ctx context.Context, cfg service.Config) (internalfix.Generator, error) { + if caller := clientCallerFrom(ctx); caller != nil && caller.supports("sampling") { + return samplingGenerator{caller: caller}, nil + } + generator, available, err := internalfix.NewAIGenerator(cfg.AI) + if err != nil { + return nil, err + } + if !available { + return nil, nil + } + return generator, nil +} + +type samplingGenerator struct { + caller clientCaller +} + +func (g samplingGenerator) GenerateFix(ctx context.Context, input internalfix.GenerateInput) (internalfix.Candidate, error) { + params := map[string]any{ + "systemPrompt": "You are a precise code-fixing assistant. Return ONLY a valid unified diff (git format) that fixes the reported finding and changes nothing unrelated. Do not include any prose, explanation, or markdown fences.", + "messages": []map[string]any{{ + "role": "user", + "content": map[string]any{ + "type": "text", + "text": buildSamplingFixPrompt(input), + }, + }}, + "maxTokens": 2048, + "includeContext": "thisServer", + "modelPreferences": map[string]any{ + "intelligencePriority": 0.8, + "speedPriority": 0.3, + }, + } + raw, err := g.caller.sampleMessage(ctx, params) + if err != nil { + return internalfix.Candidate{}, err + } + text, err := samplingResultText(raw) + if err != nil { + return internalfix.Candidate{}, err + } + candidate := candidateFromText(text) + if strings.TrimSpace(candidate.Diff) == "" { + return internalfix.Candidate{}, fmt.Errorf("sampling response did not contain a diff") + } + return candidate, nil +} + +func buildSamplingFixPrompt(input internalfix.GenerateInput) string { + var b strings.Builder + b.WriteString("Fix this codeguard finding by producing a minimal unified diff.\n\n") + if input.Finding.RuleID != "" { + fmt.Fprintf(&b, "Rule: %s\n", input.Finding.RuleID) + } + if input.Finding.Path != "" { + fmt.Fprintf(&b, "File: %s", input.Finding.Path) + if input.Finding.Line > 0 { + fmt.Fprintf(&b, ":%d", input.Finding.Line) + } + b.WriteString("\n") + } + if input.Finding.Message != "" { + fmt.Fprintf(&b, "Issue: %s\n", input.Finding.Message) + } + if analysis := firstNonEmpty(input.Analysis, input.Finding.Why); analysis != "" { + fmt.Fprintf(&b, "Why: %s\n", analysis) + } + if input.Finding.HowToFix != "" { + fmt.Fprintf(&b, "How to fix: %s\n", input.Finding.HowToFix) + } + if input.Instructions != "" { + fmt.Fprintf(&b, "\n%s\n", input.Instructions) + } + b.WriteString("\nReturn only the unified diff.") + return b.String() +} + +func samplingResultText(raw json.RawMessage) (string, error) { + var single struct { + Content struct { + Text string `json:"text"` + } `json:"content"` + } + if err := json.Unmarshal(raw, &single); err == nil && strings.TrimSpace(single.Content.Text) != "" { + return single.Content.Text, nil + } + var multi struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + } + if err := json.Unmarshal(raw, &multi); err == nil { + for _, block := range multi.Content { + if strings.TrimSpace(block.Text) != "" { + return block.Text, nil + } + } + } + return "", fmt.Errorf("sampling response had no text content") +} + +func candidateFromText(text string) internalfix.Candidate { + trimmed := strings.TrimSpace(text) + var asJSON internalfix.Candidate + if json.Unmarshal([]byte(trimmed), &asJSON) == nil && strings.TrimSpace(asJSON.Diff) != "" { + return asJSON + } + if fenced := extractFencedBlock(trimmed); fenced != "" { + return internalfix.Candidate{Diff: fenced} + } + if looksLikeDiff(trimmed) { + return internalfix.Candidate{Diff: trimmed} + } + return internalfix.Candidate{} +} + +func extractFencedBlock(text string) string { + start := strings.Index(text, "```") + if start < 0 { + return "" + } + rest := text[start+3:] + if nl := strings.IndexByte(rest, '\n'); nl >= 0 { + rest = rest[nl+1:] + } + end := strings.Index(rest, "```") + if end < 0 { + return strings.TrimSpace(rest) + } + return strings.TrimSpace(rest[:end]) +} + +func looksLikeDiff(text string) bool { + return strings.Contains(text, "@@") || strings.HasPrefix(text, "diff ") || strings.HasPrefix(text, "--- ") +} diff --git a/internal/cli/mcp_http.go b/internal/cli/mcp_http.go index 8d0822a..0dc9de2 100644 --- a/internal/cli/mcp_http.go +++ b/internal/cli/mcp_http.go @@ -1,16 +1,9 @@ package cli import ( - "context" - "crypto/rand" "crypto/subtle" - "encoding/hex" - "encoding/json" - "fmt" - "io" "net/http" "strings" - "sync" ) // mcp_http.go implements a Streamable-HTTP transport for the MCP server so @@ -65,321 +58,3 @@ func (a mcpAuthConfig) authorize(r *http.Request) bool { } return subtle.ConstantTimeCompare([]byte(got), []byte(want)) == 1 } - -type mcpHTTPHandler struct { - tools *mcpToolService - auth mcpAuthConfig - path string - sem chan struct{} - sessions *sessionRegistry -} - -// newMCPHTTPHandler builds the HTTP handler for the MCP server. It is split out -// from the listener so tests can mount it on httptest.NewServer. -func newMCPHTTPHandler(tools *mcpToolService, auth mcpAuthConfig, path string) http.Handler { - if strings.TrimSpace(path) == "" { - path = mcpDefaultHTTPPath - } - h := &mcpHTTPHandler{ - tools: tools, - auth: auth, - path: path, - sem: make(chan struct{}, mcpHTTPMaxConcurrentTools), - sessions: newSessionRegistry(), - } - mux := http.NewServeMux() - mux.HandleFunc(mcpHealthPath, h.handleHealth) - mux.HandleFunc(path, h.handleMCP) - return mux -} - -// callerForRequest returns the client caller bound to the request's session, or -// nil when no session/header is present. -func (h *mcpHTTPHandler) callerForRequest(r *http.Request) clientCaller { - sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) - if !ok { - return nil - } - return sess.caller() -} - -func (h *mcpHTTPHandler) handleHealth(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "text/plain; charset=utf-8") - w.WriteHeader(http.StatusOK) - _, _ = io.WriteString(w, "ok\n") -} - -func (h *mcpHTTPHandler) handleMCP(w http.ResponseWriter, r *http.Request) { - switch r.Method { - case http.MethodPost: - h.handlePost(w, r) - case http.MethodGet: - h.handleGetStream(w, r) - case http.MethodDelete: - if !h.auth.authorize(r) { - w.Header().Set("WWW-Authenticate", "Bearer") - w.WriteHeader(http.StatusUnauthorized) - return - } - h.sessions.remove(r.Header.Get(mcpSessionHeader)) - w.WriteHeader(http.StatusOK) - default: - w.Header().Set("Allow", "GET, POST, DELETE") - w.WriteHeader(http.StatusMethodNotAllowed) - } -} - -// handleGetStream opens the server→client SSE stream for a session. The server -// writes sampling/createMessage and roots/list requests over this stream; the -// client answers them on subsequent POSTs. -func (h *mcpHTTPHandler) handleGetStream(w http.ResponseWriter, r *http.Request) { - if !h.auth.authorize(r) { - w.Header().Set("WWW-Authenticate", "Bearer") - w.WriteHeader(http.StatusUnauthorized) - return - } - sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) - if !ok { - w.WriteHeader(http.StatusNotFound) - return - } - flusher, ok := w.(http.Flusher) - if !ok { - w.WriteHeader(http.StatusInternalServerError) - return - } - w.Header().Set("Content-Type", contentTypeEventStream) - w.Header().Set("Cache-Control", "no-cache") - w.WriteHeader(http.StatusOK) - flusher.Flush() - - sess.attachStream(w, flusher) - defer sess.detachStream() - // Emit an SSE comment so the client knows the stream is attached and the - // server can now deliver server-initiated requests over it. - _, _ = io.WriteString(w, ": ready\n\n") - flusher.Flush() - <-r.Context().Done() - // The client disconnected; drop the session so the map does not retain dead - // sessions (the common client keeps one stream open for the session's life). - h.sessions.remove(sess.id) -} - -func (h *mcpHTTPHandler) handlePost(w http.ResponseWriter, r *http.Request) { - if !h.auth.authorize(r) { - w.Header().Set("WWW-Authenticate", "Bearer") - writeJSONStatus(w, http.StatusUnauthorized, buildErrorMessage(nil, -32001, "unauthorized")) - return - } - - body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, mcpHTTPMaxBodyBytes)) - if err != nil { - writeJSONStatus(w, http.StatusRequestEntityTooLarge, buildErrorMessage(nil, -32600, "request too large")) - return - } - trimmed := strings.TrimSpace(string(body)) - if trimmed == "" { - writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) - return - } - - // JSON-RPC batch (array) — process each message and return an array of the - // non-notification responses. Batched requests do not stream progress. - if strings.HasPrefix(trimmed, "[") { - h.handleBatch(w, r, []byte(trimmed)) - return - } - - var req mcpRequest - if err := json.Unmarshal([]byte(trimmed), &req); err != nil { - writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) - return - } - if req.JSONRPC != "2.0" { - writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(req.idPtr(), -32600, "invalid request")) - return - } - - // A response (id, no method) is the client answering a server-initiated - // request; route it to the session's pending caller. - if isResponseMessage(req) { - if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { - sess.requester.deliver(decodeIDKey(req.ID), json.RawMessage(trimmed)) - } - w.WriteHeader(http.StatusAccepted) - return - } - - if req.Method == "notifications/roots/list_changed" { - if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { - sess.roots.invalidate() - } - w.WriteHeader(http.StatusAccepted) - return - } - - // Notifications carry no id and expect no response body. - if len(req.ID) == 0 { - w.WriteHeader(http.StatusAccepted) - return - } - - if req.Method == "initialize" { - sess := h.sessions.create(parseClientCapabilities(req.Params)) - w.Header().Set(mcpSessionHeader, sess.id) - } - - if req.Method == "tools/call" { - h.streamToolCall(w, r, req) - return - } - - msg, handled := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) - if !handled { - writeJSONStatus(w, http.StatusOK, buildErrorMessage(req.idPtr(), -32601, "method not found")) - return - } - writeJSONStatus(w, http.StatusOK, msg) -} - -func (h *mcpHTTPHandler) handleBatch(w http.ResponseWriter, r *http.Request, body []byte) { - var reqs []mcpRequest - if err := json.Unmarshal(body, &reqs); err != nil { - writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) - return - } - ctx := withClientCaller(r.Context(), h.callerForRequest(r)) - responses := make([]map[string]any, 0, len(reqs)) - for _, req := range reqs { - if req.JSONRPC != "2.0" { - responses = append(responses, buildErrorMessage(req.idPtr(), -32600, "invalid request")) - continue - } - if len(req.ID) == 0 { - continue // notification — no response - } - responses = append(responses, h.produceMessage(ctx, req)) - } - if len(responses) == 0 { - w.WriteHeader(http.StatusAccepted) - return - } - writeJSONStatus(w, http.StatusOK, responses) -} - -// produceMessage resolves a single request to its response envelope, running -// tools/call synchronously (no progress streaming). Used for batched requests. -func (h *mcpHTTPHandler) produceMessage(ctx context.Context, req mcpRequest) map[string]any { - if req.Method == "tools/call" { - if err := h.acquire(ctx); err != nil { - return buildErrorMessage(req.idPtr(), -32603, "server busy") - } - defer h.release() - result, err := h.tools.callToolWithContext(ctx, req.Params) - if err != nil { - return buildErrorMessage(req.idPtr(), -32602, err.Error()) - } - return buildResultMessage(req.ID, result) - } - msg, handled := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) - if !handled { - return buildErrorMessage(req.idPtr(), -32601, "method not found") - } - return msg -} - -// streamToolCall runs a tools/call and streams progress + result as SSE so the -// HTTP transport matches the stdio transport's progress behavior. Cancellation -// is driven by the request context (client disconnect). -func (h *mcpHTTPHandler) streamToolCall(w http.ResponseWriter, r *http.Request, req mcpRequest) { - flusher, ok := w.(http.Flusher) - if !ok { - // No streaming support — fall back to a single JSON response. - writeJSONStatus(w, http.StatusOK, h.produceMessage(r.Context(), req)) - return - } - - if err := h.acquire(r.Context()); err != nil { - writeJSONStatus(w, http.StatusServiceUnavailable, buildErrorMessage(req.idPtr(), -32603, "server busy")) - return - } - defer h.release() - - w.Header().Set("Content-Type", contentTypeEventStream) - w.Header().Set("Cache-Control", "no-cache") - w.WriteHeader(http.StatusOK) - - ctx := withClientCaller(r.Context(), h.callerForRequest(r)) - progressToken := progressTokenFromParams(req.Params) - if progressToken != nil { - _ = writeSSE(w, buildProgressMessage(*progressToken, 0, 1, "Started")) - flusher.Flush() - var progressMu sync.Mutex - ctx = withProgress(ctx, func(progress float64, total float64, message string) { - progressMu.Lock() - defer progressMu.Unlock() - _ = writeSSE(w, buildProgressMessage(*progressToken, progress, total, message)) - flusher.Flush() - }) - } - - result, err := h.tools.callToolWithContext(ctx, req.Params) - - if progressToken != nil { - message := "Completed" - if err != nil { - message = "Stopped" - } - _ = writeSSE(w, buildProgressMessage(*progressToken, 1, 1, message)) - flusher.Flush() - } - - var msg map[string]any - if err != nil { - msg = buildErrorMessage(req.idPtr(), -32602, err.Error()) - } else { - msg = buildResultMessage(req.ID, result) - } - _ = writeSSE(w, msg) - flusher.Flush() -} - -func (h *mcpHTTPHandler) acquire(ctx context.Context) error { - select { - case h.sem <- struct{}{}: - return nil - case <-ctx.Done(): - return ctx.Err() - } -} - -func (h *mcpHTTPHandler) release() { <-h.sem } - -func writeJSONStatus(w http.ResponseWriter, status int, payload any) { - w.Header().Set("Content-Type", contentTypeJSON) - w.WriteHeader(status) - data, err := json.Marshal(payload) - if err != nil { - return - } - _, _ = w.Write(data) - _, _ = w.Write([]byte("\n")) -} - -func writeSSE(w io.Writer, payload any) error { - data, err := json.Marshal(payload) - if err != nil { - return err - } - _, err = fmt.Fprintf(w, "data: %s\n\n", data) - return err -} - -// newSessionID returns a random hex session identifier. -func newSessionID() string { - buf := make([]byte, 16) - if _, err := rand.Read(buf); err != nil { - return "codeguard-session" - } - return hex.EncodeToString(buf) -} diff --git a/internal/cli/mcp_http_handler.go b/internal/cli/mcp_http_handler.go new file mode 100644 index 0000000..c046b53 --- /dev/null +++ b/internal/cli/mcp_http_handler.go @@ -0,0 +1,65 @@ +package cli + +import ( + "net/http" + "strings" +) + +type mcpHTTPHandler struct { + tools *mcpToolService + auth mcpAuthConfig + path string + sem chan struct{} + sessions *sessionRegistry +} + +func newMCPHTTPHandler(tools *mcpToolService, auth mcpAuthConfig, path string) http.Handler { + if strings.TrimSpace(path) == "" { + path = mcpDefaultHTTPPath + } + h := &mcpHTTPHandler{ + tools: tools, + auth: auth, + path: path, + sem: make(chan struct{}, mcpHTTPMaxConcurrentTools), + sessions: newSessionRegistry(), + } + mux := http.NewServeMux() + mux.HandleFunc(mcpHealthPath, h.handleHealth) + mux.HandleFunc(path, h.handleMCP) + return mux +} + +func (h *mcpHTTPHandler) handleHealth(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/plain; charset=utf-8") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok\n")) +} + +func (h *mcpHTTPHandler) handleMCP(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + h.handlePost(w, r) + case http.MethodGet: + h.handleGetStream(w, r) + case http.MethodDelete: + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + return + } + h.sessions.remove(r.Header.Get(mcpSessionHeader)) + w.WriteHeader(http.StatusOK) + default: + w.Header().Set("Allow", "GET, POST, DELETE") + w.WriteHeader(http.StatusMethodNotAllowed) + } +} + +func callerForRequest(h *mcpHTTPHandler, r *http.Request) clientCaller { + sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) + if !ok { + return nil + } + return sess.caller() +} diff --git a/internal/cli/mcp_http_io.go b/internal/cli/mcp_http_io.go new file mode 100644 index 0000000..0b5e2e3 --- /dev/null +++ b/internal/cli/mcp_http_io.go @@ -0,0 +1,50 @@ +package cli + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" +) + +func acquireRequestSlot(ctx context.Context, sem chan struct{}) error { + select { + case sem <- struct{}{}: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func releaseRequestSlot(sem chan struct{}) { <-sem } + +func writeJSONStatus(w http.ResponseWriter, status int, payload any) { + w.Header().Set("Content-Type", contentTypeJSON) + w.WriteHeader(status) + data, err := json.Marshal(payload) + if err != nil { + return + } + _, _ = w.Write(data) + _, _ = w.Write([]byte("\n")) +} + +func writeSSE(w io.Writer, payload any) error { + data, err := json.Marshal(payload) + if err != nil { + return err + } + _, err = fmt.Fprintf(w, "data: %s\n\n", data) + return err +} + +func newSessionID() string { + buf := make([]byte, 16) + if _, err := rand.Read(buf); err != nil { + return "codeguard-session" + } + return hex.EncodeToString(buf) +} diff --git a/internal/cli/mcp_http_registry.go b/internal/cli/mcp_http_registry.go new file mode 100644 index 0000000..29a184f --- /dev/null +++ b/internal/cli/mcp_http_registry.go @@ -0,0 +1,61 @@ +package cli + +import ( + "strings" + "sync" +) + +const maxHTTPSessions = 512 + +type sessionRegistry struct { + mu sync.Mutex + sessions map[string]*httpSession + order []string + max int +} + +func newSessionRegistry() *sessionRegistry { + return &sessionRegistry{sessions: map[string]*httpSession{}, max: maxHTTPSessions} +} + +func (r *sessionRegistry) create(caps map[string]any) *httpSession { + sess := &httpSession{id: newSessionID(), caps: caps, requester: newServerRequester(), roots: &rootsCache{}} + r.mu.Lock() + defer r.mu.Unlock() + if r.max > 0 && len(r.order) >= r.max { + oldest := r.order[0] + r.order = r.order[1:] + delete(r.sessions, oldest) + } + r.sessions[sess.id] = sess + r.order = append(r.order, sess.id) + return sess +} + +func (r *sessionRegistry) get(id string) (*httpSession, bool) { + if strings.TrimSpace(id) == "" { + return nil, false + } + r.mu.Lock() + sess, ok := r.sessions[id] + r.mu.Unlock() + return sess, ok +} + +func (r *sessionRegistry) remove(id string) { + if strings.TrimSpace(id) == "" { + return + } + r.mu.Lock() + defer r.mu.Unlock() + if _, ok := r.sessions[id]; !ok { + return + } + delete(r.sessions, id) + for i, existing := range r.order { + if existing == id { + r.order = append(r.order[:i], r.order[i+1:]...) + return + } + } +} diff --git a/internal/cli/mcp_http_session.go b/internal/cli/mcp_http_session.go index f43b428..2190346 100644 --- a/internal/cli/mcp_http_session.go +++ b/internal/cli/mcp_http_session.go @@ -3,7 +3,6 @@ package cli import ( "errors" "net/http" - "strings" "sync" ) @@ -15,10 +14,6 @@ import ( var errNoClientStream = errors.New("client SSE stream is not connected") -// maxHTTPSessions bounds the session map so streamless clients that never send -// DELETE cannot grow it without limit; the oldest session is evicted on create. -const maxHTTPSessions = 512 - type httpSession struct { id string caps map[string]any @@ -64,55 +59,3 @@ func (sess *httpSession) sendToClient(payload any) error { func (sess *httpSession) caller() clientCaller { return &clientBridge{caps: sess.caps, requester: sess.requester, send: sess.sendToClient, roots: sess.roots} } - -type sessionRegistry struct { - mu sync.Mutex - sessions map[string]*httpSession - order []string - max int -} - -func newSessionRegistry() *sessionRegistry { - return &sessionRegistry{sessions: map[string]*httpSession{}, max: maxHTTPSessions} -} - -func (r *sessionRegistry) create(caps map[string]any) *httpSession { - sess := &httpSession{id: newSessionID(), caps: caps, requester: newServerRequester(), roots: &rootsCache{}} - r.mu.Lock() - if r.max > 0 && len(r.order) >= r.max { - oldest := r.order[0] - r.order = r.order[1:] - delete(r.sessions, oldest) - } - r.sessions[sess.id] = sess - r.order = append(r.order, sess.id) - r.mu.Unlock() - return sess -} - -func (r *sessionRegistry) get(id string) (*httpSession, bool) { - if strings.TrimSpace(id) == "" { - return nil, false - } - r.mu.Lock() - sess, ok := r.sessions[id] - r.mu.Unlock() - return sess, ok -} - -func (r *sessionRegistry) remove(id string) { - if strings.TrimSpace(id) == "" { - return - } - r.mu.Lock() - if _, ok := r.sessions[id]; ok { - delete(r.sessions, id) - for i, existing := range r.order { - if existing == id { - r.order = append(r.order[:i], r.order[i+1:]...) - break - } - } - } - r.mu.Unlock() -} diff --git a/internal/cli/mcp_http_stream.go b/internal/cli/mcp_http_stream.go new file mode 100644 index 0000000..237ed4b --- /dev/null +++ b/internal/cli/mcp_http_stream.go @@ -0,0 +1,205 @@ +package cli + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "sync" +) + +func (h *mcpHTTPHandler) handleGetStream(w http.ResponseWriter, r *http.Request) { + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + return + } + sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)) + if !ok { + w.WriteHeader(http.StatusNotFound) + return + } + flusher, ok := w.(http.Flusher) + if !ok { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", contentTypeEventStream) + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + flusher.Flush() + + sess.attachStream(w, flusher) + defer sess.detachStream() + _, _ = io.WriteString(w, ": ready\n\n") + flusher.Flush() + <-r.Context().Done() + h.sessions.remove(sess.id) +} + +func (h *mcpHTTPHandler) handlePost(w http.ResponseWriter, r *http.Request) { + if !h.auth.authorize(r) { + w.Header().Set("WWW-Authenticate", "Bearer") + writeJSONStatus(w, http.StatusUnauthorized, buildErrorMessage(nil, -32001, "unauthorized")) + return + } + + body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, mcpHTTPMaxBodyBytes)) + if err != nil { + writeJSONStatus(w, http.StatusRequestEntityTooLarge, buildErrorMessage(nil, -32600, "request too large")) + return + } + trimmed := strings.TrimSpace(string(body)) + if trimmed == "" { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + if strings.HasPrefix(trimmed, "[") { + h.handleBatch(w, r, []byte(trimmed)) + return + } + + var req mcpRequest + if err := json.Unmarshal([]byte(trimmed), &req); err != nil { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + status, handled := h.handleSpecialPost(w, r, req, trimmed) + if handled { + if status > 0 { + w.WriteHeader(status) + } + return + } + + msg, ok := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) + if !ok { + writeJSONStatus(w, http.StatusOK, buildErrorMessage(req.idPtr(), -32601, "method not found")) + return + } + writeJSONStatus(w, http.StatusOK, msg) +} + +func (h *mcpHTTPHandler) handleSpecialPost(w http.ResponseWriter, r *http.Request, req mcpRequest, trimmed string) (int, bool) { + if req.JSONRPC != "2.0" { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(req.idPtr(), -32600, "invalid request")) + return 0, true + } + if isResponseMessage(req) { + if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { + sess.requester.deliver(decodeIDKey(req.ID), json.RawMessage(trimmed)) + } + return http.StatusAccepted, true + } + if req.Method == "notifications/roots/list_changed" { + if sess, ok := h.sessions.get(r.Header.Get(mcpSessionHeader)); ok { + sess.roots.invalidate() + } + return http.StatusAccepted, true + } + if len(req.ID) == 0 { + return http.StatusAccepted, true + } + if req.Method == "initialize" { + sess := h.sessions.create(parseClientCapabilities(req.Params)) + w.Header().Set(mcpSessionHeader, sess.id) + } + if req.Method == "tools/call" { + h.streamToolCall(w, r, req) + return 0, true + } + return 0, false +} + +func (h *mcpHTTPHandler) handleBatch(w http.ResponseWriter, r *http.Request, body []byte) { + var reqs []mcpRequest + if err := json.Unmarshal(body, &reqs); err != nil { + writeJSONStatus(w, http.StatusBadRequest, buildErrorMessage(nil, -32700, "parse error")) + return + } + ctx := withClientCaller(r.Context(), callerForRequest(h, r)) + responses := make([]map[string]any, 0, len(reqs)) + for _, req := range reqs { + if req.JSONRPC != "2.0" { + responses = append(responses, buildErrorMessage(req.idPtr(), -32600, "invalid request")) + continue + } + if len(req.ID) == 0 { + continue + } + responses = append(responses, h.produceMessage(ctx, req)) + } + if len(responses) == 0 { + w.WriteHeader(http.StatusAccepted) + return + } + writeJSONStatus(w, http.StatusOK, responses) +} + +func (h *mcpHTTPHandler) produceMessage(ctx context.Context, req mcpRequest) map[string]any { + if req.Method == "tools/call" { + if err := acquireRequestSlot(ctx, h.sem); err != nil { + return buildErrorMessage(req.idPtr(), -32603, "server busy") + } + defer releaseRequestSlot(h.sem) + result, err := h.tools.callToolWithContext(ctx, req.Params) + if err != nil { + return buildErrorMessage(req.idPtr(), -32602, err.Error()) + } + return buildResultMessage(req.ID, result) + } + msg, handled := h.tools.dispatchSyncMethod(req.Method, req.ID, req.Params) + if !handled { + return buildErrorMessage(req.idPtr(), -32601, "method not found") + } + return msg +} + +func (h *mcpHTTPHandler) streamToolCall(w http.ResponseWriter, r *http.Request, req mcpRequest) { + flusher, ok := w.(http.Flusher) + if !ok { + writeJSONStatus(w, http.StatusOK, h.produceMessage(r.Context(), req)) + return + } + if err := acquireRequestSlot(r.Context(), h.sem); err != nil { + writeJSONStatus(w, http.StatusServiceUnavailable, buildErrorMessage(req.idPtr(), -32603, "server busy")) + return + } + defer releaseRequestSlot(h.sem) + + w.Header().Set("Content-Type", contentTypeEventStream) + w.Header().Set("Cache-Control", "no-cache") + w.WriteHeader(http.StatusOK) + + ctx := withClientCaller(r.Context(), callerForRequest(h, r)) + progressToken := progressTokenFromParams(req.Params) + if progressToken != nil { + _ = writeSSE(w, buildProgressMessage(*progressToken, 0, 1, "Started")) + flusher.Flush() + var progressMu sync.Mutex + ctx = withProgress(ctx, func(progress float64, total float64, message string) { + progressMu.Lock() + defer progressMu.Unlock() + _ = writeSSE(w, buildProgressMessage(*progressToken, progress, total, message)) + flusher.Flush() + }) + } + + result, err := h.tools.callToolWithContext(ctx, req.Params) + if progressToken != nil { + message := "Completed" + if err != nil { + message = "Stopped" + } + _ = writeSSE(w, buildProgressMessage(*progressToken, 1, 1, message)) + flusher.Flush() + } + if err != nil { + _ = writeSSE(w, buildErrorMessage(req.idPtr(), -32602, err.Error())) + flusher.Flush() + return + } + _ = writeSSE(w, buildResultMessage(req.ID, result)) + flusher.Flush() +} diff --git a/internal/cli/mcp_protocol.go b/internal/cli/mcp_protocol.go index 17a51b3..7f1e458 100644 --- a/internal/cli/mcp_protocol.go +++ b/internal/cli/mcp_protocol.go @@ -1,270 +1 @@ package cli - -import ( - "encoding/json" - "strings" -) - -func toolSuccessResult(payload any) map[string]any { - text := mustJSON(payload) - return map[string]any{ - "content": []map[string]any{{ - "type": "text", - "text": text, - }}, - "structuredContent": payload, - "isError": false, - } -} - -func toolErrorResult(message string) map[string]any { - return map[string]any{ - "content": []map[string]any{{ - "type": "text", - "text": message, - }}, - "isError": true, - } -} - -// toolErrorResultData is toolErrorResult with machine-readable structuredContent -// so callers can act on the failure (e.g. a fix that did not verify carries the -// attempted diff and remaining findings). -func toolErrorResultData(message string, data any) map[string]any { - result := toolErrorResult(message) - if data != nil { - result["structuredContent"] = data - } - return result -} - -func mustJSON(value any) string { - data, err := json.Marshal(value) - if err != nil { - return "{}" - } - return string(data) -} - -func mcpTools() []mcpTool { - tools := []mcpTool{ - { - Name: "scan", - Title: "Scan Repository", - Description: "Run codeguard against the configured repository targets and return a structured report.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "mode": map[string]any{"type": "string", "enum": []string{"full", "diff"}}, - "base_ref": map[string]any{"type": "string"}, - }, - }, - OutputSchema: reportOutputSchema(), - }, - { - Name: "validate_config", - Title: "Validate Config", - Description: "Validate the configured codeguard policy file and return a machine-readable result.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - }, - }, - OutputSchema: objectOutputSchema(), - }, - { - Name: "validate_patch", - Title: "Validate Patch", - Description: "Evaluate a unified diff against policy without mutating the working tree and return a structured report.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "diff": map[string]any{"type": "string"}, - }, - "required": []string{"diff"}, - }, - OutputSchema: reportOutputSchema(), - }, - { - Name: "explain", - Title: "Explain Rule", - Description: "Return machine-first explanation metadata for a codeguard rule.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "rule_id": map[string]any{"type": "string"}, - }, - "required": []string{"rule_id"}, - }, - OutputSchema: objectOutputSchema(), - }, - { - Name: "list_rules", - Title: "List Rules", - Description: "Return the rule catalog that applies to the current or requested configuration.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - }, - }, - OutputSchema: objectOutputSchema(), - }, - { - Name: "verify_fix", - Title: "Verify Fix", - Description: "Verify a candidate unified diff against a finding: apply it in an isolated workspace, re-scan the changed lines, run the nearest inferred tests, and return the result only if it passes. Does not modify the working tree.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "finding": fixFindingSchema(), - "diff": map[string]any{"type": "string", "description": "Candidate unified diff to verify."}, - "base_ref": map[string]any{"type": "string"}, - "max_nearest_tests": map[string]any{"type": "integer"}, - "test_commands": map[string]any{"type": "array"}, - }, - "required": []string{"diff"}, - }, - OutputSchema: objectOutputSchema(), - }, - { - Name: "propose_fix", - Title: "Propose Fix", - Description: "Generate a candidate fix for a finding (via the client's LLM when MCP sampling is supported, else a configured AI provider), verify it in an isolated workspace with re-scan and nearest tests, and return it only if it passes. Does not modify the working tree.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "finding": fixFindingSchema(), - "base_ref": map[string]any{"type": "string"}, - "max_nearest_tests": map[string]any{"type": "integer"}, - "test_commands": map[string]any{"type": "array"}, - }, - "required": []string{"finding"}, - }, - OutputSchema: objectOutputSchema(), - }, - { - Name: "apply_fix", - Title: "Apply Fix", - Description: "Verify a candidate unified diff and, only if it passes, write it to the working tree. Asks the user to confirm first when the client supports elicitation. This tool modifies files on disk.", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{ - "config_path": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "finding": fixFindingSchema(), - "diff": map[string]any{"type": "string", "description": "Candidate unified diff to verify and apply."}, - "base_ref": map[string]any{"type": "string"}, - "max_nearest_tests": map[string]any{"type": "integer"}, - "test_commands": map[string]any{"type": "array"}, - }, - "required": []string{"diff"}, - }, - OutputSchema: objectOutputSchema(), - Annotations: writeToolAnnotations("Apply Fix"), - }, - } - - // Most codeguard tools only read the repository; tools that already declared - // annotations (e.g. the destructive apply_fix) keep them. - for i := range tools { - if tools[i].Annotations == nil { - tools[i].Annotations = readOnlyToolAnnotations(tools[i].Title) - } - } - return tools -} - -// writeToolAnnotations marks a tool that mutates the working tree. -func writeToolAnnotations(title string) map[string]any { - return map[string]any{ - "title": title, - "readOnlyHint": false, - "destructiveHint": true, - "idempotentHint": false, - "openWorldHint": false, - } -} - -// fixFindingSchema describes the finding object accepted by the fix tools. -func fixFindingSchema() map[string]any { - return map[string]any{ - "type": "object", - "description": "The finding to fix, as returned in a scan/validate_patch report.", - "properties": map[string]any{ - "rule_id": map[string]any{"type": "string"}, - "path": map[string]any{"type": "string"}, - "line": map[string]any{"type": "integer"}, - "message": map[string]any{"type": "string"}, - "why": map[string]any{"type": "string"}, - "how_to_fix": map[string]any{"type": "string"}, - }, - } -} - -// readOnlyToolAnnotations returns MCP tool annotations marking a tool as -// read-only and non-destructive. -func readOnlyToolAnnotations(title string) map[string]any { - return map[string]any{ - "title": title, - "readOnlyHint": true, - "destructiveHint": false, - "idempotentHint": true, - "openWorldHint": false, - } -} - -// objectOutputSchema is a permissive output schema for tools that return an -// arbitrary JSON object in structuredContent. -func objectOutputSchema() map[string]any { - return map[string]any{"type": "object"} -} - -// reportOutputSchema describes the codeguard report object returned by scan and -// validate_patch in structuredContent. -func reportOutputSchema() map[string]any { - return map[string]any{ - "type": "object", - "properties": map[string]any{ - "name": map[string]any{"type": "string"}, - "profile": map[string]any{"type": "string"}, - "sections": map[string]any{"type": "array"}, - "summary": map[string]any{"type": "object"}, - }, - } -} - -func negotiateMCPProtocolVersion(raw json.RawMessage) string { - var params struct { - ProtocolVersion string `json:"protocolVersion"` - } - if err := json.Unmarshal(raw, ¶ms); err != nil { - return mcpProtocolVersionCompat - } - switch strings.TrimSpace(params.ProtocolVersion) { - case mcpProtocolVersionCurrent, mcpProtocolVersionCompat: - return params.ProtocolVersion - default: - return mcpProtocolVersionCompat - } -} - -func normalizeMCPArguments(raw json.RawMessage) json.RawMessage { - if len(raw) == 0 || strings.TrimSpace(string(raw)) == "null" { - return json.RawMessage([]byte("{}")) - } - return raw -} diff --git a/internal/cli/mcp_protocol_version.go b/internal/cli/mcp_protocol_version.go new file mode 100644 index 0000000..047f21d --- /dev/null +++ b/internal/cli/mcp_protocol_version.go @@ -0,0 +1,28 @@ +package cli + +import ( + "encoding/json" + "strings" +) + +func negotiateMCPProtocolVersion(raw json.RawMessage) string { + var params struct { + ProtocolVersion string `json:"protocolVersion"` + } + if err := json.Unmarshal(raw, ¶ms); err != nil { + return mcpProtocolVersionCompat + } + switch strings.TrimSpace(params.ProtocolVersion) { + case mcpProtocolVersionCurrent, mcpProtocolVersionCompat: + return params.ProtocolVersion + default: + return mcpProtocolVersionCompat + } +} + +func normalizeMCPArguments(raw json.RawMessage) json.RawMessage { + if len(raw) == 0 || strings.TrimSpace(string(raw)) == "null" { + return json.RawMessage([]byte("{}")) + } + return raw +} diff --git a/internal/cli/mcp_run.go b/internal/cli/mcp_run.go index 4321b69..c18855e 100644 --- a/internal/cli/mcp_run.go +++ b/internal/cli/mcp_run.go @@ -138,17 +138,10 @@ func (s *mcpServer) handleLine(line string, stdout io.Writer) error { } func (s *mcpServer) handleRequestMethod(req mcpRequest, stdout io.Writer) error { + if handled, err := s.handleAsyncRequest(req, stdout); handled { + return err + } switch req.Method { - case "initialize": - return s.handleInitializeResponse(req, stdout) - case "notifications/initialized": - return nil - case "notifications/cancelled": - s.handleCancelledNotification(req.Params) - return nil - case "notifications/roots/list_changed": - s.rootsCache.invalidate() - return nil case "ping": return s.responder.writeResult(stdout, req.ID, map[string]any{}) case "tools/list": @@ -166,6 +159,23 @@ func (s *mcpServer) handleRequestMethod(req mcpRequest, stdout io.Writer) error } } +func (s *mcpServer) handleAsyncRequest(req mcpRequest, stdout io.Writer) (bool, error) { + switch req.Method { + case "initialize": + return true, s.handleInitializeResponse(req, stdout) + case "notifications/initialized": + return true, nil + case "notifications/cancelled": + s.handleCancelledNotification(req.Params) + return true, nil + case "notifications/roots/list_changed": + s.rootsCache.invalidate() + return true, nil + default: + return false, nil + } +} + // handleSyncMethod serves the request methods that produce a single synchronous // response (resources/*, prompts/*, logging/*) by delegating to the shared // transport-neutral router and writing its envelope to stdout. diff --git a/internal/cli/mcp_tool_catalog.go b/internal/cli/mcp_tool_catalog.go new file mode 100644 index 0000000..08468da --- /dev/null +++ b/internal/cli/mcp_tool_catalog.go @@ -0,0 +1,169 @@ +package cli + +func mcpTools() []mcpTool { + tools := []mcpTool{ + scanTool(), + validateConfigTool(), + validatePatchTool(), + explainTool(), + listRulesTool(), + verifyFixTool(), + proposeFixTool(), + applyFixTool(), + } + for i := range tools { + if tools[i].Annotations == nil { + tools[i].Annotations = readOnlyToolAnnotations(tools[i].Title) + } + } + return tools +} + +func scanTool() mcpTool { + return mcpTool{ + Name: "scan", + Title: "Scan Repository", + Description: "Run codeguard against the configured repository targets and return a structured report.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "mode": map[string]any{"type": "string", "enum": []string{"full", "diff"}}, + "base_ref": map[string]any{"type": "string"}, + }, + }, + OutputSchema: reportOutputSchema(), + } +} + +func validateConfigTool() mcpTool { + return mcpTool{ + Name: "validate_config", + Title: "Validate Config", + Description: "Validate the configured codeguard policy file and return a machine-readable result.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + }, + }, + OutputSchema: objectOutputSchema(), + } +} + +func validatePatchTool() mcpTool { + return mcpTool{ + Name: "validate_patch", + Title: "Validate Patch", + Description: "Evaluate a unified diff against policy without mutating the working tree and return a structured report.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "diff": map[string]any{"type": "string"}, + }, + "required": []string{"diff"}, + }, + OutputSchema: reportOutputSchema(), + } +} + +func explainTool() mcpTool { + return mcpTool{ + Name: "explain", + Title: "Explain Rule", + Description: "Return machine-first explanation metadata for a codeguard rule.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "rule_id": map[string]any{"type": "string"}, + }, + "required": []string{"rule_id"}, + }, + OutputSchema: objectOutputSchema(), + } +} + +func listRulesTool() mcpTool { + return mcpTool{ + Name: "list_rules", + Title: "List Rules", + Description: "Return the rule catalog that applies to the current or requested configuration.", + InputSchema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + }, + }, + OutputSchema: objectOutputSchema(), + } +} + +func verifyFixTool() mcpTool { + return buildFixTool( + "verify_fix", + "Verify Fix", + "Verify a candidate unified diff against a finding: apply it in an isolated workspace, re-scan the changed lines, run the nearest inferred tests, and return the result only if it passes. Does not modify the working tree.", + true, + false, + ) +} + +func proposeFixTool() mcpTool { + return buildFixTool( + "propose_fix", + "Propose Fix", + "Generate a candidate fix for a finding (via the client's LLM when MCP sampling is supported, else a configured AI provider), verify it in an isolated workspace with re-scan and nearest tests, and return it only if it passes. Does not modify the working tree.", + false, + false, + ) +} + +func applyFixTool() mcpTool { + tool := buildFixTool( + "apply_fix", + "Apply Fix", + "Verify a candidate unified diff and, only if it passes, write it to the working tree. Asks the user to confirm first when the client supports elicitation. This tool modifies files on disk.", + true, + true, + ) + tool.Annotations = writeToolAnnotations("Apply Fix") + return tool +} + +func buildFixTool(name string, title string, description string, requireDiff bool, destructive bool) mcpTool { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "config_path": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "finding": fixFindingSchema(), + "diff": map[string]any{"type": "string", "description": "Candidate unified diff to verify and apply."}, + "base_ref": map[string]any{"type": "string"}, + "max_nearest_tests": map[string]any{"type": "integer"}, + "test_commands": map[string]any{"type": "array"}, + }, + } + if requireDiff { + schema["required"] = []string{"diff"} + } else { + schema["required"] = []string{"finding"} + } + tool := mcpTool{ + Name: name, + Title: title, + Description: description, + InputSchema: schema, + OutputSchema: objectOutputSchema(), + } + if destructive { + tool.Annotations = writeToolAnnotations(title) + } + return tool +} diff --git a/internal/cli/mcp_tool_config.go b/internal/cli/mcp_tool_config.go new file mode 100644 index 0000000..b7b5970 --- /dev/null +++ b/internal/cli/mcp_tool_config.go @@ -0,0 +1,102 @@ +package cli + +import ( + "context" + "os" + "path/filepath" + "strings" + "time" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +func (s *mcpToolService) loadConfig(configPath string, profile string) (service.Config, error) { + path := strings.TrimSpace(configPath) + if path == "" { + path = s.defaultConfigPath + } + overrideProfile := strings.TrimSpace(profile) + if overrideProfile == "" { + overrideProfile = strings.TrimSpace(s.defaultProfile) + } + return loadConfigWithProfile(path, overrideProfile) +} + +func confineConfigArg(ctx context.Context, s *mcpToolService, configPath string) (string, error) { + candidate := strings.TrimSpace(configPath) + if candidate == "" { + return "", nil + } + return confinePath(allowedRoots(ctx, s), candidate) +} + +const rootsFetchTimeout = 10 * time.Second + +func allowedRoots(ctx context.Context, s *mcpToolService) []string { + roots := make([]string, 0, 3) + if def := strings.TrimSpace(s.defaultConfigPath); def != "" { + roots = append(roots, configDirOf(def)) + } + if wd, err := os.Getwd(); err == nil { + roots = append(roots, wd) + } + if caller := clientCallerFrom(ctx); caller != nil && caller.supports("roots") { + rctx, cancel := context.WithTimeout(ctx, rootsFetchTimeout) + if clientRoots, err := caller.listRoots(rctx); err == nil { + for _, root := range clientRoots { + if p := rootURIToPath(root.URI); p != "" { + roots = append(roots, p) + } + } + } + cancel() + } + return roots +} + +func rootURIToPath(uri string) string { + uri = strings.TrimSpace(uri) + if uri == "" { + return "" + } + if strings.HasPrefix(uri, "file://") { + return strings.TrimPrefix(uri, "file://") + } + if strings.Contains(uri, "://") { + return "" + } + return uri +} + +func configDirOf(path string) string { + if info, err := os.Stat(path); err == nil && info.IsDir() { + return path + } + return filepath.Dir(path) +} + +func confinePath(allowedRoots []string, candidate string) (string, error) { + abs := resolvePath(candidate) + for _, root := range allowedRoots { + if strings.TrimSpace(root) == "" { + continue + } + rootAbs := resolvePath(root) + if abs == rootAbs || strings.HasPrefix(abs, rootAbs+string(os.PathSeparator)) { + return abs, nil + } + } + return "", errConfigPathNotPermitted +} + +func resolvePath(path string) string { + abs, err := filepath.Abs(path) + if err != nil { + return filepath.Clean(path) + } + abs = filepath.Clean(abs) + if resolved, err := filepath.EvalSymlinks(abs); err == nil { + return resolved + } + return abs +} diff --git a/internal/cli/mcp_tool_dispatch.go b/internal/cli/mcp_tool_dispatch.go new file mode 100644 index 0000000..e100d67 --- /dev/null +++ b/internal/cli/mcp_tool_dispatch.go @@ -0,0 +1,37 @@ +package cli + +import ( + "context" + "encoding/json" + "fmt" + "strings" +) + +func (s *mcpToolService) callToolWithContext(ctx context.Context, raw json.RawMessage) (map[string]any, error) { + var call struct { + Name string `json:"name"` + Arguments json.RawMessage `json:"arguments"` + } + if err := json.Unmarshal(raw, &call); err != nil { + return nil, fmt.Errorf("invalid tool call") + } + + handler, ok := toolHandlers(s)[strings.TrimSpace(call.Name)] + if !ok { + return nil, fmt.Errorf("unknown tool: %s", call.Name) + } + return handler(ctx, normalizeMCPArguments(call.Arguments)) +} + +func toolHandlers(s *mcpToolService) map[string]func(context.Context, json.RawMessage) (map[string]any, error) { + return map[string]func(context.Context, json.RawMessage) (map[string]any, error){ + "scan": s.callScan, + "validate_config": s.callValidateConfig, + "validate_patch": s.callValidatePatch, + "explain": s.callExplain, + "list_rules": s.callListRules, + "verify_fix": s.callVerifyFix, + "propose_fix": s.callProposeFix, + "apply_fix": s.callApplyFix, + } +} diff --git a/internal/cli/mcp_tool_results.go b/internal/cli/mcp_tool_results.go new file mode 100644 index 0000000..62191d3 --- /dev/null +++ b/internal/cli/mcp_tool_results.go @@ -0,0 +1,41 @@ +package cli + +import "encoding/json" + +func toolSuccessResult(payload any) map[string]any { + text := mustJSON(payload) + return map[string]any{ + "content": []map[string]any{{ + "type": "text", + "text": text, + }}, + "structuredContent": payload, + "isError": false, + } +} + +func toolErrorResult(message string) map[string]any { + return map[string]any{ + "content": []map[string]any{{ + "type": "text", + "text": message, + }}, + "isError": true, + } +} + +func toolErrorResultData(message string, data any) map[string]any { + result := toolErrorResult(message) + if data != nil { + result["structuredContent"] = data + } + return result +} + +func mustJSON(value any) string { + data, err := json.Marshal(value) + if err != nil { + return "{}" + } + return string(data) +} diff --git a/internal/cli/mcp_tool_rules.go b/internal/cli/mcp_tool_rules.go new file mode 100644 index 0000000..5b468e7 --- /dev/null +++ b/internal/cli/mcp_tool_rules.go @@ -0,0 +1,29 @@ +package cli + +import ( + "strings" + + service "github.com/devr-tools/codeguard/pkg/codeguard" +) + +func (s *mcpToolService) resolveExplainRule(configPath string, profile string, ruleID string) (service.RuleMetadata, bool, error) { + if strings.TrimSpace(configPath) != "" { + cfg, err := s.loadConfig(configPath, profile) + if err != nil { + return service.RuleMetadata{}, false, err + } + rule, ok := service.ExplainRuleForConfig(cfg, ruleID) + return rule, ok, nil + } + + rule, ok := service.ExplainRule(ruleID) + if strings.TrimSpace(profile) == "" { + return rule, ok, nil + } + cfg, err := s.loadConfig("", profile) + if err != nil { + return service.RuleMetadata{}, false, err + } + rule, ok = service.ExplainRuleForConfig(cfg, ruleID) + return rule, ok, nil +} diff --git a/internal/cli/mcp_tool_schemas.go b/internal/cli/mcp_tool_schemas.go new file mode 100644 index 0000000..28c241e --- /dev/null +++ b/internal/cli/mcp_tool_schemas.go @@ -0,0 +1,52 @@ +package cli + +func writeToolAnnotations(title string) map[string]any { + return map[string]any{ + "title": title, + "readOnlyHint": false, + "destructiveHint": true, + "idempotentHint": false, + "openWorldHint": false, + } +} + +func fixFindingSchema() map[string]any { + return map[string]any{ + "type": "object", + "description": "The finding to fix, as returned in a scan/validate_patch report.", + "properties": map[string]any{ + "rule_id": map[string]any{"type": "string"}, + "path": map[string]any{"type": "string"}, + "line": map[string]any{"type": "integer"}, + "message": map[string]any{"type": "string"}, + "why": map[string]any{"type": "string"}, + "how_to_fix": map[string]any{"type": "string"}, + }, + } +} + +func readOnlyToolAnnotations(title string) map[string]any { + return map[string]any{ + "title": title, + "readOnlyHint": true, + "destructiveHint": false, + "idempotentHint": true, + "openWorldHint": false, + } +} + +func objectOutputSchema() map[string]any { + return map[string]any{"type": "object"} +} + +func reportOutputSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "name": map[string]any{"type": "string"}, + "profile": map[string]any{"type": "string"}, + "sections": map[string]any{"type": "array"}, + "summary": map[string]any{"type": "object"}, + }, + } +} diff --git a/internal/cli/mcp_tools.go b/internal/cli/mcp_tools.go index 380b9e6..964f0f7 100644 --- a/internal/cli/mcp_tools.go +++ b/internal/cli/mcp_tools.go @@ -5,11 +5,8 @@ import ( "encoding/json" "errors" "fmt" - "os" - "path/filepath" "strings" "sync" - "time" service "github.com/devr-tools/codeguard/pkg/codeguard" ) @@ -19,37 +16,6 @@ import ( // the HTTP transport does not become a filesystem oracle. var errConfigPathNotPermitted = errors.New("config_path is not within an allowed root") -func (s *mcpToolService) callToolWithContext(ctx context.Context, raw json.RawMessage) (map[string]any, error) { - var call struct { - Name string `json:"name"` - Arguments json.RawMessage `json:"arguments"` - } - if err := json.Unmarshal(raw, &call); err != nil { - return nil, fmt.Errorf("invalid tool call") - } - - switch strings.TrimSpace(call.Name) { - case "scan": - return s.callScan(ctx, normalizeMCPArguments(call.Arguments)) - case "validate_config": - return s.callValidateConfig(ctx, normalizeMCPArguments(call.Arguments)) - case "validate_patch": - return s.callValidatePatch(ctx, normalizeMCPArguments(call.Arguments)) - case "explain": - return s.callExplain(ctx, normalizeMCPArguments(call.Arguments)) - case "list_rules": - return s.callListRules(ctx, normalizeMCPArguments(call.Arguments)) - case "verify_fix": - return s.callVerifyFix(ctx, normalizeMCPArguments(call.Arguments)) - case "propose_fix": - return s.callProposeFix(ctx, normalizeMCPArguments(call.Arguments)) - case "apply_fix": - return s.callApplyFix(ctx, normalizeMCPArguments(call.Arguments)) - default: - return nil, fmt.Errorf("unknown tool: %s", call.Name) - } -} - func (s *mcpToolService) callScan(ctx context.Context, raw json.RawMessage) (map[string]any, error) { var args struct { ConfigPath string `json:"config_path"` @@ -61,7 +27,7 @@ func (s *mcpToolService) callScan(ctx context.Context, raw json.RawMessage) (map return nil, fmt.Errorf("invalid scan arguments") } - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) if err != nil { return toolErrorResult(err.Error()), nil } @@ -111,7 +77,7 @@ func (s *mcpToolService) callValidateConfig(ctx context.Context, raw json.RawMes return nil, fmt.Errorf("invalid validate_config arguments") } - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) if err != nil { return toolErrorResult(err.Error()), nil } @@ -142,7 +108,7 @@ func (s *mcpToolService) callValidatePatch(ctx context.Context, raw json.RawMess return toolErrorResult("validate_patch requires a unified diff"), nil } - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) if err != nil { return toolErrorResult(err.Error()), nil } @@ -170,7 +136,7 @@ func (s *mcpToolService) callExplain(ctx context.Context, raw json.RawMessage) ( return toolErrorResult("explain requires rule_id"), nil } - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) if err != nil { return toolErrorResult(err.Error()), nil } @@ -196,7 +162,7 @@ func (s *mcpToolService) callListRules(ctx context.Context, raw json.RawMessage) if strings.TrimSpace(args.ConfigPath) == "" && strings.TrimSpace(args.Profile) == "" { return toolSuccessResult(map[string]any{"rules": service.Rules()}), nil } - confinedPath, err := s.confineConfigArg(ctx, args.ConfigPath) + confinedPath, err := confineConfigArg(ctx, s, args.ConfigPath) if err != nil { return toolErrorResult(err.Error()), nil } @@ -206,135 +172,3 @@ func (s *mcpToolService) callListRules(ctx context.Context, raw json.RawMessage) } return toolSuccessResult(map[string]any{"rules": service.RulesForConfig(cfg)}), nil } - -func (s *mcpToolService) loadConfig(configPath string, profile string) (service.Config, error) { - path := strings.TrimSpace(configPath) - if path == "" { - path = s.defaultConfigPath - } - overrideProfile := strings.TrimSpace(profile) - if overrideProfile == "" { - overrideProfile = strings.TrimSpace(s.defaultProfile) - } - return loadConfigWithProfile(path, overrideProfile) -} - -// confineConfigArg validates a caller-supplied config_path against the allowed -// roots and returns it resolved to an absolute path. An empty argument returns -// "" so the server's trusted default config is used. The server's own default -// config path is never passed through here, so a trusted out-of-tree --config -// keeps working. -func (s *mcpToolService) confineConfigArg(ctx context.Context, configPath string) (string, error) { - candidate := strings.TrimSpace(configPath) - if candidate == "" { - return "", nil - } - return confinePath(s.allowedRoots(ctx), candidate) -} - -// rootsFetchTimeout bounds the server→client roots/list round trip done while -// resolving a config path, so a config load never blocks on a slow client. -const rootsFetchTimeout = 10 * time.Second - -// allowedRoots is the set of directories a caller-supplied config_path may live -// under: the directory of the server's default config, the working directory, -// any roots advertised by the connected client (via the roots capability), and -// any roots injected explicitly (tests). -func (s *mcpToolService) allowedRoots(ctx context.Context) []string { - roots := make([]string, 0, 4) - if def := strings.TrimSpace(s.defaultConfigPath); def != "" { - roots = append(roots, configDirOf(def)) - } - if wd, err := os.Getwd(); err == nil { - roots = append(roots, wd) - } - if caller := clientCallerFrom(ctx); caller != nil && caller.supports("roots") { - rctx, cancel := context.WithTimeout(ctx, rootsFetchTimeout) - if clientRoots, err := caller.listRoots(rctx); err == nil { - for _, root := range clientRoots { - if p := rootURIToPath(root.URI); p != "" { - roots = append(roots, p) - } - } - } - cancel() - } - roots = append(roots, clientRootsFrom(ctx)...) - return roots -} - -// rootURIToPath converts a roots entry to a filesystem path. MCP roots are -// file:// URIs; a bare path is accepted as-is. -func rootURIToPath(uri string) string { - uri = strings.TrimSpace(uri) - if uri == "" { - return "" - } - if strings.HasPrefix(uri, "file://") { - return strings.TrimPrefix(uri, "file://") - } - if strings.Contains(uri, "://") { - return "" // non-file scheme is not a local path - } - return uri -} - -// configDirOf returns the directory that should anchor confinement for a config -// path: the path itself if it is a directory, otherwise its parent. -func configDirOf(path string) string { - if info, err := os.Stat(path); err == nil && info.IsDir() { - return path - } - return filepath.Dir(path) -} - -// confinePath resolves candidate to an absolute, cleaned path and requires it to -// sit within one of allowedRoots. Symlinks are resolved best-effort to defeat -// link escapes; on resolution failure it falls back to the cleaned path. -func confinePath(allowedRoots []string, candidate string) (string, error) { - abs := resolvePath(candidate) - for _, root := range allowedRoots { - if strings.TrimSpace(root) == "" { - continue - } - rootAbs := resolvePath(root) - if abs == rootAbs || strings.HasPrefix(abs, rootAbs+string(os.PathSeparator)) { - return abs, nil - } - } - return "", errConfigPathNotPermitted -} - -func resolvePath(path string) string { - abs, err := filepath.Abs(path) - if err != nil { - return filepath.Clean(path) - } - abs = filepath.Clean(abs) - if resolved, err := filepath.EvalSymlinks(abs); err == nil { - return resolved - } - return abs -} - -func (s *mcpToolService) resolveExplainRule(configPath string, profile string, ruleID string) (service.RuleMetadata, bool, error) { - if strings.TrimSpace(configPath) != "" { - cfg, err := s.loadConfig(configPath, profile) - if err != nil { - return service.RuleMetadata{}, false, err - } - rule, ok := service.ExplainRuleForConfig(cfg, ruleID) - return rule, ok, nil - } - - rule, ok := service.ExplainRule(ruleID) - if strings.TrimSpace(profile) == "" { - return rule, ok, nil - } - cfg, err := s.loadConfig("", profile) - if err != nil { - return service.RuleMetadata{}, false, err - } - rule, ok = service.ExplainRuleForConfig(cfg, ruleID) - return rule, ok, nil -} diff --git a/internal/codeguard/ai/fix/generator.go b/internal/codeguard/ai/fix/generator.go index b458798..67da6f0 100644 --- a/internal/codeguard/ai/fix/generator.go +++ b/internal/codeguard/ai/fix/generator.go @@ -70,7 +70,10 @@ func sourceExcerpt(cfg core.Config, finding core.Finding) string { return "" } for _, target := range cfg.Targets { - fullPath := filepath.Join(target.Path, filepath.FromSlash(finding.Path)) + fullPath, err := containedFindingPath(target.Path, finding.Path) + if err != nil { + continue + } data, err := os.ReadFile(fullPath) if err != nil { continue @@ -86,6 +89,22 @@ func sourceExcerpt(cfg core.Config, finding core.Finding) string { return "" } +func containedFindingPath(targetPath string, findingPath string) (string, error) { + baseDir, err := filepath.Abs(targetPath) + if err != nil { + return "", err + } + fullPath := filepath.Clean(filepath.Join(baseDir, filepath.FromSlash(findingPath))) + rel, err := filepath.Rel(baseDir, fullPath) + if err != nil { + return "", err + } + if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("finding path %q escapes target path %q", findingPath, targetPath) + } + return fullPath, nil +} + func maxInt(a int, b int) int { if a > b { return a diff --git a/tests/mcp/http_helpers_test.go b/tests/mcp/http_helpers_test.go new file mode 100644 index 0000000..1f7cd60 --- /dev/null +++ b/tests/mcp/http_helpers_test.go @@ -0,0 +1,156 @@ +package mcp_test + +import ( + "encoding/json" + "net/http" + "strings" + "testing" +) + +func assertInitializeCapabilities(t *testing.T, base string) { + t.Helper() + resp, body := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}`) + if resp.StatusCode != http.StatusOK { + t.Fatalf("expected 200, got %d body=%s", resp.StatusCode, body) + } + if resp.Header.Get("Mcp-Session-Id") == "" { + t.Fatalf("expected Mcp-Session-Id header on initialize") + } + var out struct { + Result struct { + ProtocolVersion string `json:"protocolVersion"` + Capabilities map[string]any `json:"capabilities"` + } `json:"result"` + } + decodeLine(t, body, &out) + if out.Result.ProtocolVersion != "2025-06-18" { + t.Fatalf("unexpected protocol version: %s", out.Result.ProtocolVersion) + } + for _, capability := range []string{"tools", "resources", "prompts", "logging"} { + if _, ok := out.Result.Capabilities[capability]; !ok { + t.Fatalf("missing capability %q: %#v", capability, out.Result.Capabilities) + } + } +} + +func assertToolsListAnnotations(t *testing.T, base string) { + t.Helper() + _, body := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`) + var out struct { + Result struct { + Tools []struct { + Name string `json:"name"` + Annotations struct { + ReadOnlyHint bool `json:"readOnlyHint"` + } `json:"annotations"` + } `json:"tools"` + } `json:"result"` + } + decodeLine(t, body, &out) + if len(out.Result.Tools) != 8 { + t.Fatalf("expected 8 tools, got %d", len(out.Result.Tools)) + } + for _, tool := range out.Result.Tools { + if tool.Name == "apply_fix" { + continue + } + if !tool.Annotations.ReadOnlyHint { + t.Fatalf("tool %q missing readOnlyHint", tool.Name) + } + } +} + +func assertToolCallStreamsSSE(t *testing.T, base string) { + t.Helper() + resp, body := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"list_rules","arguments":{},"_meta":{"progressToken":"p1"}}}`) + if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") { + t.Fatalf("expected SSE content type, got %q", ct) + } + var sawProgress, sawResult bool + for _, frame := range strings.Split(body, "\n\n") { + frame = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(frame), "data:")) + if frame == "" { + continue + } + var msg struct { + Method string `json:"method"` + Result struct { + IsError bool `json:"isError"` + } `json:"result"` + } + decodeLine(t, frame, &msg) + if msg.Method == "notifications/progress" { + sawProgress = true + } + if msg.Method == "" && !msg.Result.IsError { + sawResult = true + } + } + if !sawProgress || !sawResult { + t.Fatalf("expected progress + result over SSE (progress=%v result=%v): %s", sawProgress, sawResult, body) + } +} + +func assertResourceRead(t *testing.T, base string) { + t.Helper() + _, body := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"codeguard://rules"}}`) + var out struct { + Result struct { + Contents []struct { + URI string `json:"uri"` + Text string `json:"text"` + } `json:"contents"` + } `json:"result"` + } + decodeLine(t, body, &out) + if len(out.Result.Contents) == 0 || out.Result.Contents[0].URI != "codeguard://rules" { + t.Fatalf("unexpected resources/read response: %s", body) + } + var payload struct { + Rules []json.RawMessage `json:"rules"` + } + if err := json.Unmarshal([]byte(out.Result.Contents[0].Text), &payload); err != nil || len(payload.Rules) == 0 { + t.Fatalf("expected rule catalog in resource text, got err=%v len=%d", err, len(payload.Rules)) + } +} + +func assertPromptGet(t *testing.T, base string) { + t.Helper() + _, body := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":5,"method":"prompts/get","params":{"name":"review-diff","arguments":{"diff":"--- a\n+++ b\n"}}}`) + if !strings.Contains(body, "validate_patch") { + t.Fatalf("expected review-diff prompt to mention validate_patch: %s", body) + } +} + +func assertHealthAndMissingSession(t *testing.T, base string) { + t.Helper() + resp, body := httpGet(t, base+"/healthz") + if resp.StatusCode != http.StatusOK || !strings.Contains(body, "ok") { + t.Fatalf("unexpected health response: %d %s", resp.StatusCode, body) + } + resp, _ = httpGet(t, base+"/mcp") + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("expected 404 for GET /mcp without session, got %d", resp.StatusCode) + } +} + +func assertAuthEnforcement(t *testing.T, base string) { + t.Helper() + initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}` + cases := []struct { + name string + header map[string]string + want int + }{ + {"missing", nil, http.StatusUnauthorized}, + {"wrong", map[string]string{"Authorization": "Bearer nope"}, http.StatusUnauthorized}, + {"correct", map[string]string{"Authorization": "Bearer " + httpTestToken}, http.StatusOK}, + {"correct-no-scheme", map[string]string{"Authorization": httpTestToken}, http.StatusOK}, + } + for _, tc := range cases { + resp, body := mcpPost(t, base, tc.header, initBody) + if resp.StatusCode != tc.want { + t.Fatalf("%s: expected %d, got %d body=%s", tc.name, tc.want, resp.StatusCode, body) + } + } +} diff --git a/tests/mcp/http_test.go b/tests/mcp/http_test.go index 53f7d15..a56bcb7 100644 --- a/tests/mcp/http_test.go +++ b/tests/mcp/http_test.go @@ -1,7 +1,6 @@ package mcp_test import ( - "encoding/json" "io" "net" "net/http" @@ -39,146 +38,31 @@ func TestMCPServeHTTP(t *testing.T) { authBase := startHTTPServer(t, httpTestToken) t.Run("initialize-capabilities", func(t *testing.T) { - resp, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}`) - if resp.StatusCode != http.StatusOK { - t.Fatalf("expected 200, got %d body=%s", resp.StatusCode, body) - } - if resp.Header.Get("Mcp-Session-Id") == "" { - t.Fatalf("expected Mcp-Session-Id header on initialize") - } - var out struct { - Result struct { - ProtocolVersion string `json:"protocolVersion"` - Capabilities map[string]any `json:"capabilities"` - } `json:"result"` - } - decodeLine(t, body, &out) - if out.Result.ProtocolVersion != "2025-06-18" { - t.Fatalf("unexpected protocol version: %s", out.Result.ProtocolVersion) - } - for _, capability := range []string{"tools", "resources", "prompts", "logging"} { - if _, ok := out.Result.Capabilities[capability]; !ok { - t.Fatalf("missing capability %q: %#v", capability, out.Result.Capabilities) - } - } + assertInitializeCapabilities(t, openBase) }) t.Run("tools-list-annotations", func(t *testing.T) { - _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`) - var out struct { - Result struct { - Tools []struct { - Name string `json:"name"` - Annotations struct { - ReadOnlyHint bool `json:"readOnlyHint"` - } `json:"annotations"` - } `json:"tools"` - } `json:"result"` - } - decodeLine(t, body, &out) - if len(out.Result.Tools) != 8 { - t.Fatalf("expected 8 tools, got %d", len(out.Result.Tools)) - } - for _, tool := range out.Result.Tools { - if tool.Name == "apply_fix" { - continue // apply_fix is the one destructive tool - } - if !tool.Annotations.ReadOnlyHint { - t.Fatalf("tool %q missing readOnlyHint", tool.Name) - } - } + assertToolsListAnnotations(t, openBase) }) t.Run("tool-call-streams-sse", func(t *testing.T) { - resp, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"list_rules","arguments":{},"_meta":{"progressToken":"p1"}}}`) - if ct := resp.Header.Get("Content-Type"); !strings.HasPrefix(ct, "text/event-stream") { - t.Fatalf("expected SSE content type, got %q", ct) - } - var sawProgress, sawResult bool - for _, frame := range strings.Split(body, "\n\n") { - frame = strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(frame), "data:")) - if frame == "" { - continue - } - var msg struct { - Method string `json:"method"` - Result struct { - IsError bool `json:"isError"` - } `json:"result"` - } - decodeLine(t, frame, &msg) - if msg.Method == "notifications/progress" { - sawProgress = true - } - if msg.Method == "" && !msg.Result.IsError { - sawResult = true - } - } - if !sawProgress || !sawResult { - t.Fatalf("expected progress + result over SSE (progress=%v result=%v): %s", sawProgress, sawResult, body) - } + assertToolCallStreamsSSE(t, openBase) }) t.Run("resource-read", func(t *testing.T) { - _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"codeguard://rules"}}`) - var out struct { - Result struct { - Contents []struct { - URI string `json:"uri"` - Text string `json:"text"` - } `json:"contents"` - } `json:"result"` - } - decodeLine(t, body, &out) - if len(out.Result.Contents) == 0 || out.Result.Contents[0].URI != "codeguard://rules" { - t.Fatalf("unexpected resources/read response: %s", body) - } - var payload struct { - Rules []json.RawMessage `json:"rules"` - } - if err := json.Unmarshal([]byte(out.Result.Contents[0].Text), &payload); err != nil || len(payload.Rules) == 0 { - t.Fatalf("expected rule catalog in resource text, got err=%v len=%d", err, len(payload.Rules)) - } + assertResourceRead(t, openBase) }) t.Run("prompt-get", func(t *testing.T) { - _, body := mcpPost(t, openBase, nil, `{"jsonrpc":"2.0","id":5,"method":"prompts/get","params":{"name":"review-diff","arguments":{"diff":"--- a\n+++ b\n"}}}`) - if !strings.Contains(body, "validate_patch") { - t.Fatalf("expected review-diff prompt to mention validate_patch: %s", body) - } + assertPromptGet(t, openBase) }) t.Run("health-and-get-stream-needs-session", func(t *testing.T) { - resp, body := httpGet(t, openBase+"/healthz") - if resp.StatusCode != http.StatusOK || !strings.Contains(body, "ok") { - t.Fatalf("unexpected health response: %d %s", resp.StatusCode, body) - } - // GET /mcp is now the server→client SSE stream; without a session id it - // has nothing to attach to and returns 404. - resp, _ = httpGet(t, openBase+"/mcp") - if resp.StatusCode != http.StatusNotFound { - t.Fatalf("expected 404 for GET /mcp without session, got %d", resp.StatusCode) - } + assertHealthAndMissingSession(t, openBase) }) t.Run("auth-enforcement", func(t *testing.T) { - initBody := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}}` - cases := []struct { - name string - header map[string]string - want int - }{ - {"missing", nil, http.StatusUnauthorized}, - {"wrong", map[string]string{"Authorization": "Bearer nope"}, http.StatusUnauthorized}, - {"correct", map[string]string{"Authorization": "Bearer " + httpTestToken}, http.StatusOK}, - {"correct-no-scheme", map[string]string{"Authorization": httpTestToken}, http.StatusOK}, - } - for _, tc := range cases { - resp, body := mcpPost(t, authBase, tc.header, initBody) - if resp.StatusCode != tc.want { - t.Fatalf("%s: expected %d, got %d body=%s", tc.name, tc.want, resp.StatusCode, body) - } - } + assertAuthEnforcement(t, authBase) }) } diff --git a/tests/mcp/sampling_helpers_test.go b/tests/mcp/sampling_helpers_test.go new file mode 100644 index 0000000..d8f3fc4 --- /dev/null +++ b/tests/mcp/sampling_helpers_test.go @@ -0,0 +1,119 @@ +package mcp_test + +import ( + "bufio" + "encoding/json" + "net/http" + "strings" + "testing" + "time" +) + +func initializeHTTPSession(t *testing.T, base string, capabilities string) string { + t.Helper() + body := `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":` + capabilities + `}}` + resp, _ := mcpPost(t, base, nil, body) + session := resp.Header.Get("Mcp-Session-Id") + if session == "" { + t.Fatalf("no session id returned") + } + return session +} + +func openSessionStream(t *testing.T, base string, session string) *bufio.Reader { + t.Helper() + streamReq, _ := http.NewRequest(http.MethodGet, base+"/mcp", nil) + streamReq.Header.Set("Mcp-Session-Id", session) + streamResp, err := http.DefaultClient.Do(streamReq) + if err != nil { + t.Fatalf("open stream: %v", err) + } + t.Cleanup(func() { _ = streamResp.Body.Close() }) + streamReader := bufio.NewReader(streamResp.Body) + if _, err := streamReader.ReadString('\n'); err != nil { + t.Fatalf("read stream readiness: %v", err) + } + return streamReader +} + +func waitForSamplingRequest(t *testing.T, streamReader *bufio.Reader, base string, session string) { + t.Helper() + deadline := time.Now().Add(30 * time.Second) + for time.Now().Before(deadline) { + line, err := streamReader.ReadString('\n') + if err != nil { + t.Fatalf("read stream: %v", err) + } + data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") + if !ok { + continue + } + var msg struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) == nil && msg.Method == "sampling/createMessage" { + mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, samplingResponse(t, msg.ID)) + return + } + } + t.Fatalf("server did not issue sampling/createMessage over the stream") +} + +func waitForRootsRequest(t *testing.T, streamReader *bufio.Reader, base string, session string, dir string) { + t.Helper() + deadline := time.Now().Add(20 * time.Second) + for time.Now().Before(deadline) { + line, err := streamReader.ReadString('\n') + if err != nil { + t.Fatalf("read stream: %v", err) + } + data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") + if !ok { + continue + } + var msg struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) == nil && msg.Method == "roots/list" { + rootsResp, _ := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": msg.ID, + "result": map[string]any{"roots": []map[string]any{{"uri": "file://" + dir, "name": "temp"}}}, + }) + mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, string(rootsResp)) + return + } + } + t.Fatalf("server did not issue roots/list") +} + +func assertNoSecondRootsRequest(t *testing.T, streamReader *bufio.Reader) { + t.Helper() + secondRoots := make(chan struct{}, 1) + go func() { + for { + line, err := streamReader.ReadString('\n') + if err != nil { + return + } + data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") + if !ok { + continue + } + var msg struct { + Method string `json:"method"` + } + if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) == nil && msg.Method == "roots/list" { + secondRoots <- struct{}{} + return + } + } + }() + select { + case <-secondRoots: + t.Fatalf("server issued a second roots/list instead of using the cache") + case <-time.After(500 * time.Millisecond): + } +} diff --git a/tests/mcp/sampling_stdio_helpers_test.go b/tests/mcp/sampling_stdio_helpers_test.go new file mode 100644 index 0000000..de11207 --- /dev/null +++ b/tests/mcp/sampling_stdio_helpers_test.go @@ -0,0 +1,92 @@ +package mcp_test + +import ( + "bufio" + "encoding/json" + "io" + "os" + "os/exec" + "strings" + "testing" + "time" +) + +type stdioSamplingOutcome struct { + sawSampling bool + sawFix bool +} + +type samplingMessage struct { + ID json.RawMessage `json:"id"` + Method string `json:"method"` +} + +func startStdioSamplingServer(t *testing.T, cfg string) (*exec.Cmd, io.WriteCloser, io.ReadCloser) { + t.Helper() + cmd := exec.Command(os.Args[0], "-test.run=TestMCPServeHelperProcess", "--", cfg) + cmd.Env = append(os.Environ(), "GO_WANT_MCP_HELPER_PROCESS=1") + stdin, err := cmd.StdinPipe() + if err != nil { + t.Fatalf("stdin pipe: %v", err) + } + stdout, err := cmd.StdoutPipe() + if err != nil { + t.Fatalf("stdout pipe: %v", err) + } + if err := cmd.Start(); err != nil { + t.Fatalf("start: %v", err) + } + return cmd, stdin, stdout +} + +func assertStdioSamplingRoundTrip(t *testing.T, scanner *bufio.Scanner, write func(string)) { + t.Helper() + result := make(chan stdioSamplingOutcome, 1) + go func() { + result <- collectStdioSamplingOutcome(t, scanner, write) + }() + select { + case out := <-result: + if !out.sawSampling { + t.Fatalf("server did not issue sampling/createMessage") + } + if !out.sawFix { + t.Fatalf("propose_fix never returned a result") + } + case <-time.After(30 * time.Second): + t.Fatalf("timed out waiting for sampling round trip") + } +} + +func collectStdioSamplingOutcome(t *testing.T, scanner *bufio.Scanner, write func(string)) stdioSamplingOutcome { + t.Helper() + var out stdioSamplingOutcome + for scanner.Scan() { + line := strings.TrimSpace(scanner.Text()) + if line == "" { + continue + } + msg, ok := decodeSamplingMessage(line) + if !ok { + continue + } + if msg.Method == "sampling/createMessage" { + out.sawSampling = true + write(samplingResponse(t, msg.ID)) + continue + } + if strings.TrimSpace(string(msg.ID)) == `"fix"` { + out.sawFix = true + return out + } + } + return out +} + +func decodeSamplingMessage(line string) (samplingMessage, bool) { + var msg samplingMessage + if json.Unmarshal([]byte(line), &msg) != nil { + return msg, false + } + return msg, true +} diff --git a/tests/mcp/sampling_test.go b/tests/mcp/sampling_test.go index 006e176..3b8aa23 100644 --- a/tests/mcp/sampling_test.go +++ b/tests/mcp/sampling_test.go @@ -4,9 +4,7 @@ import ( "bufio" "encoding/json" "io" - "net/http" "os" - "os/exec" "path/filepath" "strings" "testing" @@ -59,19 +57,7 @@ func TestMCPStdioSampling(t *testing.T) { dir := t.TempDir() cfg := writeFixtureConfig(t, dir) - cmd := exec.Command(os.Args[0], "-test.run=TestMCPServeHelperProcess", "--", cfg) - cmd.Env = append(os.Environ(), "GO_WANT_MCP_HELPER_PROCESS=1") - stdin, err := cmd.StdinPipe() - if err != nil { - t.Fatalf("stdin pipe: %v", err) - } - stdout, err := cmd.StdoutPipe() - if err != nil { - t.Fatalf("stdout pipe: %v", err) - } - if err := cmd.Start(); err != nil { - t.Fatalf("start: %v", err) - } + cmd, stdin, stdout := startStdioSamplingServer(t, cfg) t.Cleanup(func() { _ = stdin.Close() _ = cmd.Process.Kill() @@ -90,50 +76,7 @@ func TestMCPStdioSampling(t *testing.T) { scanner := bufio.NewScanner(stdout) scanner.Buffer(make([]byte, 0, 64*1024), 10*1024*1024) - type outcome struct { - sawSampling bool - sawFix bool - } - result := make(chan outcome, 1) - go func() { - var out outcome - for scanner.Scan() { - line := strings.TrimSpace(scanner.Text()) - if line == "" { - continue - } - var msg struct { - ID json.RawMessage `json:"id"` - Method string `json:"method"` - } - if json.Unmarshal([]byte(line), &msg) != nil { - continue - } - if msg.Method == "sampling/createMessage" { - out.sawSampling = true - write(samplingResponse(t, msg.ID)) - continue - } - if strings.TrimSpace(string(msg.ID)) == `"fix"` { - out.sawFix = true - result <- out - return - } - } - result <- out - }() - - select { - case out := <-result: - if !out.sawSampling { - t.Fatalf("server did not issue sampling/createMessage") - } - if !out.sawFix { - t.Fatalf("propose_fix never returned a result") - } - case <-time.After(30 * time.Second): - t.Fatalf("timed out waiting for sampling round trip") - } + assertStdioSamplingRoundTrip(t, scanner, write) } // TestMCPHTTPSampling drives propose_fix over HTTP and answers the server's @@ -143,62 +86,15 @@ func TestMCPHTTPSampling(t *testing.T) { cfg := writeFixtureConfig(t, dir) base := startHTTPServerWithConfig(t, "", cfg) - // initialize, advertising sampling, to obtain a session id. - resp, _ := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{"sampling":{}}}}`) - session := resp.Header.Get("Mcp-Session-Id") - if session == "" { - t.Fatalf("no session id returned") - } - - // Open the server→client SSE stream and wait for the readiness comment. - streamReq, _ := http.NewRequest(http.MethodGet, base+"/mcp", nil) - streamReq.Header.Set("Mcp-Session-Id", session) - streamResp, err := http.DefaultClient.Do(streamReq) - if err != nil { - t.Fatalf("open stream: %v", err) - } - t.Cleanup(func() { _ = streamResp.Body.Close() }) - streamReader := bufio.NewReader(streamResp.Body) - if _, err := streamReader.ReadString('\n'); err != nil { // ": ready" - t.Fatalf("read stream readiness: %v", err) - } - - // Fire propose_fix; it blocks server-side until we answer the sampling request. + session := initializeHTTPSession(t, base, `{"sampling":{}}`) + streamReader := openSessionStream(t, base, session) fixDone := make(chan string, 1) go func() { _, body := mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"propose_fix","arguments":{"finding":{"rule_id":"demo.rule","message":"bad value","path":"main.go","line":1}}}}`) fixDone <- body }() - - // Read the sampling request off the stream and answer it on a POST. - sawSampling := false - deadline := time.Now().Add(30 * time.Second) - for time.Now().Before(deadline) && !sawSampling { - line, err := streamReader.ReadString('\n') - if err != nil { - t.Fatalf("read stream: %v", err) - } - line = strings.TrimSpace(line) - data, ok := strings.CutPrefix(line, "data:") - if !ok { - continue - } - var msg struct { - ID json.RawMessage `json:"id"` - Method string `json:"method"` - } - if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) != nil { - continue - } - if msg.Method == "sampling/createMessage" { - sawSampling = true - mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, samplingResponse(t, msg.ID)) - } - } - if !sawSampling { - t.Fatalf("server did not issue sampling/createMessage over the stream") - } + waitForSamplingRequest(t, streamReader, base, session) select { case <-fixDone: @@ -215,70 +111,18 @@ func TestMCPHTTPRootsConfinement(t *testing.T) { cfg := writeFixtureConfig(t, dir) base := startHTTPServer(t, "") // default config; cfg lives in an out-of-tree temp dir - resp, _ := mcpPost(t, base, nil, `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{"roots":{}}}}`) - session := resp.Header.Get("Mcp-Session-Id") - if session == "" { - t.Fatalf("no session id returned") - } - - streamReq, _ := http.NewRequest(http.MethodGet, base+"/mcp", nil) - streamReq.Header.Set("Mcp-Session-Id", session) - streamResp, err := http.DefaultClient.Do(streamReq) - if err != nil { - t.Fatalf("open stream: %v", err) - } - t.Cleanup(func() { _ = streamResp.Body.Close() }) - streamReader := bufio.NewReader(streamResp.Body) - if _, err := streamReader.ReadString('\n'); err != nil { - t.Fatalf("read readiness: %v", err) - } - - // validate_config with an out-of-tree config_path; the server must fetch - // roots to decide whether it is permitted. + session := initializeHTTPSession(t, base, `{"roots":{}}`) + streamReader := openSessionStream(t, base, session) validateDone := make(chan string, 1) go func() { req := `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"validate_config","arguments":{"config_path":"` + cfg + `"}}}` _, body := mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, req) validateDone <- body }() - - // Answer the server's roots/list request, advertising the temp dir as a root. - sawRoots := false - deadline := time.Now().Add(20 * time.Second) - for time.Now().Before(deadline) && !sawRoots { - line, err := streamReader.ReadString('\n') - if err != nil { - t.Fatalf("read stream: %v", err) - } - data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") - if !ok { - continue - } - var msg struct { - ID json.RawMessage `json:"id"` - Method string `json:"method"` - } - if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) != nil { - continue - } - if msg.Method == "roots/list" { - sawRoots = true - rootsResp, _ := json.Marshal(map[string]any{ - "jsonrpc": "2.0", - "id": msg.ID, - "result": map[string]any{"roots": []map[string]any{{"uri": "file://" + dir, "name": "temp"}}}, - }) - mcpPost(t, base, map[string]string{"Mcp-Session-Id": session}, string(rootsResp)) - } - } - if !sawRoots { - t.Fatalf("server did not issue roots/list") - } + waitForRootsRequest(t, streamReader, base, session, dir) select { case body := <-validateDone: - // With the temp dir advertised as a root, confinement must permit the - // path — so the response must not be the "not within an allowed root" error. if strings.Contains(body, "not within an allowed root") { t.Fatalf("config_path was rejected despite being an advertised root: %s", body) } @@ -286,7 +130,6 @@ func TestMCPHTTPRootsConfinement(t *testing.T) { t.Fatalf("validate_config did not complete after roots answer") } - // A second config_path call must reuse the cached roots — no new roots/list. secondDone := make(chan string, 1) go func() { req := `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"validate_config","arguments":{"config_path":"` + cfg + `"}}}` @@ -294,27 +137,6 @@ func TestMCPHTTPRootsConfinement(t *testing.T) { secondDone <- body }() - secondRoots := make(chan struct{}, 1) - go func() { - for { - line, err := streamReader.ReadString('\n') - if err != nil { - return - } - data, ok := strings.CutPrefix(strings.TrimSpace(line), "data:") - if !ok { - continue - } - var msg struct { - Method string `json:"method"` - } - if json.Unmarshal([]byte(strings.TrimSpace(data)), &msg) == nil && msg.Method == "roots/list" { - secondRoots <- struct{}{} - return - } - } - }() - select { case body := <-secondDone: if strings.Contains(body, "not within an allowed root") { @@ -323,10 +145,5 @@ func TestMCPHTTPRootsConfinement(t *testing.T) { case <-time.After(20 * time.Second): t.Fatalf("second validate_config did not complete") } - select { - case <-secondRoots: - t.Fatalf("server issued a second roots/list instead of using the cache") - case <-time.After(500 * time.Millisecond): - // no second roots/list — cache hit, as expected - } + assertNoSecondRootsRequest(t, streamReader) }