Skip to content

feat: stamp agent name in header #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
arunanshub merged 1 commit into from
Jun 17, 2026
Merged

feat: stamp agent name in header #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions go.work.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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=
Expand Down Expand Up @@ -56,15 +57,18 @@ 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=
github.com/docker/go-units v0.5.0/go.mod h1:fgPhTUdO+D/Jk86RDLlptpiXQzgHJF7gydDDbaIK4Dk=
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=
Expand Down Expand Up @@ -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=
Expand All @@ -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=
Expand Down
27 changes: 24 additions & 3 deletions internal/agent/inject.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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)
}
}
Expand All @@ -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 {
Expand Down
40 changes: 40 additions & 0 deletions internal/agent/inject_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
1 change: 1 addition & 0 deletions internal/cmd/protect/mcp/service_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Loading