diff --git a/pkg/agent/agent_utils.go b/pkg/agent/agent_utils.go index 6d1b5e70f..cb38fbefe 100644 --- a/pkg/agent/agent_utils.go +++ b/pkg/agent/agent_utils.go @@ -47,7 +47,11 @@ func withMCPHeadersFromRaw(ctx context.Context, raw map[string]string) context.C var headers map[string]string for k, v := range raw { after, ok := strings.CutPrefix(k, "mcp:") - if !ok || after == "" { + if !ok { + continue + } + after = strings.TrimSpace(after) + if after == "" { continue } if headers == nil { diff --git a/pkg/mcp/manager_test.go b/pkg/mcp/manager_test.go index 5789a37a9..e623f4f0c 100644 --- a/pkg/mcp/manager_test.go +++ b/pkg/mcp/manager_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "io" + "net/http" "os" "path/filepath" "strings" @@ -17,6 +18,7 @@ import ( "github.com/sipeed/picoclaw/pkg/config" runtimeevents "github.com/sipeed/picoclaw/pkg/events" + toolshared "github.com/sipeed/picoclaw/pkg/tools/shared" ) func TestLoadEnvFile(t *testing.T) { @@ -628,3 +630,102 @@ func (t *scriptedTransport) Close() error { func (t *scriptedTransport) SessionID() string { return t.sessionID } + +func TestHeaderTransport_DynamicHeaders(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"X-Static": "from-config"}, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "Authorization": "Bearer tok123", + "X-Custom": "dynamic-val", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + _, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + + if got := captured.Get("X-Static"); got != "from-config" { + t.Errorf("X-Static = %q, want %q", got, "from-config") + } + if got := captured.Get("Authorization"); got != "Bearer tok123" { + t.Errorf("Authorization = %q, want %q", got, "Bearer tok123") + } + if got := captured.Get("X-Custom"); got != "dynamic-val" { + t.Errorf("X-Custom = %q, want %q", got, "dynamic-val") + } +} + +func TestHeaderTransport_DynamicOverridesStatic(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"Authorization": "Bearer static"}, + } + + ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{ + "Authorization": "Bearer dynamic", + }) + req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil) + + _, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + + if got := captured.Get("Authorization"); got != "Bearer dynamic" { + t.Errorf("Authorization = %q, want dynamic to override static %q", got, "Bearer dynamic") + } +} + +func TestHeaderTransport_NoDynamicHeaders(t *testing.T) { + captured := make(http.Header) + base := roundTripFunc(func(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + captured[k] = v + } + return &http.Response{StatusCode: 200, Body: http.NoBody}, nil + }) + + transport := &headerTransport{ + base: base, + headers: map[string]string{"X-Static": "val"}, + } + + req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, "http://example.com", nil) + _, err := transport.RoundTrip(req) + if err != nil { + t.Fatalf("RoundTrip() error = %v", err) + } + + if got := captured.Get("X-Static"); got != "val" { + t.Errorf("X-Static = %q, want %q", got, "val") + } + if got := captured.Get("Authorization"); got != "" { + t.Errorf("Authorization should be empty, got %q", got) + } +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/pkg/tools/shared/base.go b/pkg/tools/shared/base.go index 0b23c298a..ab8cb7ffa 100644 --- a/pkg/tools/shared/base.go +++ b/pkg/tools/shared/base.go @@ -132,14 +132,30 @@ func ToolSessionScope(ctx context.Context) *session.SessionScope { } // WithMCPHeaders returns a child context carrying per-request headers for MCP HTTP transports. +// The map is cloned to prevent callers from mutating it after storage. func WithMCPHeaders(ctx context.Context, headers map[string]string) context.Context { - return context.WithValue(ctx, ctxKeyMCPHeaders, headers) + if len(headers) == 0 { + return ctx + } + clone := make(map[string]string, len(headers)) + for k, v := range headers { + clone[k] = v + } + return context.WithValue(ctx, ctxKeyMCPHeaders, clone) } // MCPHeaders extracts per-request MCP headers from ctx, or nil if unset. +// Returns a clone so callers cannot mutate the stored map. func MCPHeaders(ctx context.Context) map[string]string { v, _ := ctx.Value(ctxKeyMCPHeaders).(map[string]string) - return v + if len(v) == 0 { + return nil + } + clone := make(map[string]string, len(v)) + for k, val := range v { + clone[k] = val + } + return clone } // AsyncCallback is a function type that async tools use to notify completion.