diff --git a/go.work.sum b/go.work.sum index cd3aab5..e4468d5 100644 --- a/go.work.sum +++ b/go.work.sum @@ -1,5 +1,6 @@ al.essio.dev/pkg/shellescape v1.5.1/go.mod h1:6sIqp7X2P6mThCQ7twERpZTuigpr6KbZWtls1U8I890= buf.build/gen/go/bufbuild/protovalidate/grpc/go v1.6.1-20240508200655-46a4cf4ba109.1/go.mod h1:12iIaR0LjReZQXXxBXMzWTMUIN6n4Y3HuSTwPpUQYSg= +buf.build/gen/go/bufbuild/protovalidate/grpc/go v1.6.2-20240508200655-46a4cf4ba109.1/go.mod h1:cz5A7G0AEk7z2kciJ1jK72u/8kRVlcfAhi2B9jYoMZw= cel.dev/expr v0.25.1/go.mod h1:hrXvqGP6G6gyx8UAHSHJ5RGk//1Oj5nXQ2NI02Nrsg4= cloud.google.com/go v0.121.2/go.mod h1:nRFlrHq39MNVWu+zESP2PosMWA0ryJw8KUBZ2iZpxbw= cloud.google.com/go/auth v0.16.1/go.mod h1:1howDHJ5IETh/LwYs3ZxvlkXF48aSqqJUM+5o02dNOI= @@ -56,6 +57,7 @@ github.com/cloudwego/eino v0.7.13/go.mod h1:nA8Vacmuqv3pqKBQbTWENBLQ8MmGmPt/Wqiy github.com/cloudwego/eino-ext/components/model/claude v0.1.17/go.mod h1:2sGGgwpR60LW+RdG/hcjdGBVEVEJ6EkKxlva8mju9BI= github.com/cloudwego/eino-ext/components/model/gemini v0.1.5/go.mod h1:yWVzN9Y5TU9MERRStVnfa28uXQbbAk+z/2SCi+G32Vc= github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5/go.mod h1:KdCmV+x/BuvyMxRnYBlmVaq4OLiKW6iRQfvC62cvdkI= +github.com/cncf/xds/go v0.0.0-20260202195803-dba9d589def2/go.mod h1:qwXFYgsP6T7XnJtbKlf1HP8AjxZZyzxMmc+Lq5GjlU4= github.com/distribution/reference v0.6.0/go.mod h1:BbU0aIcezP1/5jX/8MP0YiH4SdvB5Y4f/wlDRiLyi3E= github.com/docker/docker v28.1.1+incompatible/go.mod h1:eEKB0N0r5NX/I1kEveEz05bcu8tLC/8azJZsviup8Sk= github.com/docker/go-connections v0.5.0/go.mod h1:ov60Kzw0kKElRwhNs9UlUHAE/F9Fe6GLaXnqyDdmEXc= @@ -63,8 +65,10 @@ github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDD github.com/eino-contrib/jsonschema v1.0.3/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= github.com/envoyproxy/go-control-plane v0.14.0/go.mod h1:NcS5X47pLl/hfqxU70yPwL9ZMkUlwlKxtAohpi2wBEU= github.com/envoyproxy/go-control-plane/envoy v1.36.0/go.mod h1:ty89S1YCCVruQAm9OtKeEkQLTb+Lkz0k8v9W0Oxsv98= +github.com/envoyproxy/go-control-plane/envoy v1.37.0/go.mod h1:DReE9MMrmecPy+YvQOAOHNYMALuowAnbjjEMkkWOi6A= github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= github.com/envoyproxy/protoc-gen-validate v1.3.0/go.mod h1:HvYl7zwPa5mffgyeTUHA9zHIH36nmrm7oCbo4YKoSWA= +github.com/envoyproxy/protoc-gen-validate v1.3.3/go.mod h1:TsndJ/ngyIdQRhMcVVGDDHINPLWB7C82oDArY51KfB0= github.com/felixge/fgprof v0.9.5/go.mod h1:yKl+ERSa++RYOs32d8K6WEXCB4uXdLls4ZaZPpayhMM= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/gabriel-vasile/mimetype v1.4.10 h1:zyueNbySn/z8mJZHLt6IPw0KoZsiQNszIpU+bX4+ZK0= @@ -146,6 +150,7 @@ github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+x github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= go.opentelemetry.io/contrib/detectors/gcp v1.39.0/go.mod h1:t/OGqzHBa5v6RHZwrDBJ2OirWc+4q/w2fTbLZwAKjTk= +go.opentelemetry.io/contrib/detectors/gcp v1.42.0/go.mod h1:W9zQ439utxymRrXsUOzZbFX4JhLxXU4+ZnCt8GG7yA8= go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho v0.61.0/go.mod h1:/V0rmKWoHzXI2ROCfKE2PKPoo6hdlU1GRtzwzuO/3jc= go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0/go.mod h1:snMWehoOh2wsEwnvvwtDyFCxVeDAODenXHtn5vzrKjo= go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0/go.mod h1:UHB22Z8QsdRDrnAtX4PntOl36ajSxcdUMt1sF7Y6E7Q= @@ -166,6 +171,7 @@ google.golang.org/genai v1.21.0/go.mod h1:QPj5NGJw+3wEOHg+PrsWwJKvG6UC84ex5FR7qA google.golang.org/genproto v0.0.0-20250528174236-200df99c418a h1:KXuwdBmgjb4T3l4ZzXhP6HxxFKXD9FcK5/8qfJI4WwU= google.golang.org/genproto v0.0.0-20250528174236-200df99c418a/go.mod h1:Nlk93rrS2X7rV8hiC2gh2A/AJspZhElz9Oh2KGsjLEY= google.golang.org/genproto/googleapis/api v0.0.0-20260120221211-b8f7ae30c516/go.mod h1:p3MLuOwURrGBRoEyFHBT3GjUwaCQVKeNqqWxlcISGdw= +google.golang.org/genproto/googleapis/api v0.0.0-20260226221140-a57be14db171/go.mod h1:M5krXqk4GhBKvB596udGL3UyjL4I1+cTbK0orROM9ng= gorm.io/driver/clickhouse v0.7.0/go.mod h1:TmNo0wcVTsD4BBObiRnCahUgHJHjBIwuRejHwYt3JRs= gorm.io/driver/mysql v1.5.7/go.mod h1:sEtPWMiqiN1N1cMXoXmBbd8C6/l+TESwriotuRRpkDM= gorm.io/driver/postgres v1.6.0/go.mod h1:vUw0mrGgrTK+uPHEhAdV4sfFELrByKVGnaVRkXDhtWo= diff --git a/internal/agent/inject.go b/internal/agent/inject.go index e010c9c..8d5fcdf 100644 --- a/internal/agent/inject.go +++ b/internal/agent/inject.go @@ -1,6 +1,13 @@ package agent -import "errors" +import ( + "errors" + "maps" +) + +// AgentNameHeader carries the agent slug (e.g. "claude-code") to the MCP +// server, which uses it to attribute requests to a specific agent. +const AgentNameHeader = "X-Agent-Name" // InjectAll injects the SafeDep MCP config into every detected agent. // workspaceDir="" skips workspace injection. Best-effort: all agents @@ -13,15 +20,19 @@ func InjectAll(agents []Agent, cfg MCPConfig, workspaceDir string) error { continue } + // Clone the headers so each agent gets its own name without mutating + // the caller's shared map. + ac := withAgentName(cfg, a.Name()) + if inj, ok := a.AsGlobalInjector(); ok { - if err := inj.InjectGlobal(cfg); err != nil { + if err := inj.InjectGlobal(ac); err != nil { errs = append(errs, err) } } if workspaceDir != "" { if inj, ok := a.AsWorkspaceInjector(); ok { - if err := inj.InjectWorkspace(workspaceDir, cfg); err != nil { + if err := inj.InjectWorkspace(workspaceDir, ac); err != nil { errs = append(errs, err) } } @@ -31,6 +42,16 @@ func InjectAll(agents []Agent, cfg MCPConfig, workspaceDir string) error { return errors.Join(errs...) } +// withAgentName returns a copy of cfg with AgentNameHeader set to name, +// cloning Headers so the caller's map is never mutated. +func withAgentName(cfg MCPConfig, name string) MCPConfig { + h := make(map[string]string, len(cfg.Headers)+1) + maps.Copy(h, cfg.Headers) + h[AgentNameHeader] = name + cfg.Headers = h + return cfg +} + // RemoveAll removes the SafeDep MCP config from every detected agent. // Same error semantics as InjectAll. func RemoveAll(agents []Agent, workspaceDir string) error { diff --git a/internal/agent/inject_test.go b/internal/agent/inject_test.go index d888fba..d5c1e9e 100644 --- a/internal/agent/inject_test.go +++ b/internal/agent/inject_test.go @@ -136,6 +136,46 @@ func TestInjectAll(t *testing.T) { a := &fakeAgent{name: "x", detected: true, global: nil} require.NoError(t, InjectAll([]Agent{a}, cfg, "")) }) + + t.Run("stamps X-Agent-Name from the agent's Name into global config", func(t *testing.T) { + gi := &fakeGlobalInjector{} + a := &fakeAgent{name: "claude-code", detected: true, global: gi} + + require.NoError(t, InjectAll([]Agent{a}, cfg, "")) + require.NotNil(t, gi.injected) + assert.Equal(t, "claude-code", gi.injected.Headers[AgentNameHeader]) + }) + + t.Run("stamps X-Agent-Name into workspace config", func(t *testing.T) { + wi := &fakeWorkspaceInjector{} + a := &fakeAgent{name: "cursor", detected: true, workspace: wi} + + require.NoError(t, InjectAll([]Agent{a}, cfg, "/project")) + require.NotNil(t, wi.injected) + assert.Equal(t, "cursor", wi.injected.Headers[AgentNameHeader]) + }) + + t.Run("each agent gets its own name, not a shared one", func(t *testing.T) { + gi1, gi2 := &fakeGlobalInjector{}, &fakeGlobalInjector{} + a1 := &fakeAgent{name: "claude-code", detected: true, global: gi1} + a2 := &fakeAgent{name: "cursor", detected: true, global: gi2} + + require.NoError(t, InjectAll([]Agent{a1, a2}, cfg, "")) + assert.Equal(t, "claude-code", gi1.injected.Headers[AgentNameHeader]) + assert.Equal(t, "cursor", gi2.injected.Headers[AgentNameHeader]) + }) + + t.Run("does not mutate the caller's shared headers map", func(t *testing.T) { + shared := MCPConfig{URL: "https://mcp.safedep.io", Headers: map[string]string{"X-Tenant-ID": "t1"}} + gi := &fakeGlobalInjector{} + a := &fakeAgent{name: "claude-code", detected: true, global: gi} + + require.NoError(t, InjectAll([]Agent{a}, shared, "")) + assert.Equal(t, "t1", gi.injected.Headers["X-Tenant-ID"], "original headers preserved") + assert.Equal(t, "claude-code", gi.injected.Headers[AgentNameHeader]) + _, leaked := shared.Headers[AgentNameHeader] + assert.False(t, leaked, "per-agent header must not leak into the shared config") + }) } func TestRemoveAll(t *testing.T) { diff --git a/internal/cmd/protect/mcp/service_test.go b/internal/cmd/protect/mcp/service_test.go index c29e518..30e1c23 100644 --- a/internal/cmd/protect/mcp/service_test.go +++ b/internal/cmd/protect/mcp/service_test.go @@ -106,6 +106,7 @@ func TestMCPServiceInstall(t *testing.T) { assert.Equal(t, "Bearer key123", gi.injected.Headers["Authorization"]) assert.Equal(t, "tenant-1", gi.injected.Headers["X-Tenant-ID"]) assert.NotEmpty(t, gi.injected.Headers["X-Endpoint-ID"]) + assert.Equal(t, "claude-code", gi.injected.Headers[agent.AgentNameHeader]) }) t.Run("skips undetected agents", func(t *testing.T) {