Merge 5b4f3d575a into 412705783d

pkg/agent/agent_utils.go

@@ -15,6 +15,7 @@ import (
 	"github.com/sipeed/picoclaw/pkg/config"
 	"github.com/sipeed/picoclaw/pkg/providers"
 	"github.com/sipeed/picoclaw/pkg/session"
+	"github.com/sipeed/picoclaw/pkg/tools"
 	"github.com/sipeed/picoclaw/pkg/utils"
 )
 
@@ -39,6 +40,31 @@ func outboundContextFromInbound(
 	return outboundCtx
 }
 
+func withMCPHeadersFromRaw(ctx context.Context, raw map[string]string) context.Context {
+	if len(raw) == 0 {
+		return ctx
+	}
+	var headers map[string]string
+	for k, v := range raw {
+		after, ok := strings.CutPrefix(k, "mcp:")
+		if !ok {
+			continue
+		}
+		after = strings.TrimSpace(after)
+		if after == "" {
+			continue
+		}
+		if headers == nil {
+			headers = make(map[string]string)
+		}
+		headers[after] = v
+	}
+	if len(headers) == 0 {
+		return ctx
+	}
+	return tools.WithMCPHeaders(ctx, headers)
+}
+
 func outboundScopeFromSessionScope(scope *session.SessionScope) *bus.OutboundScope {
 	if scope == nil {
 		return nil

pkg/agent/pipeline_execute.go

@@ -528,6 +528,9 @@ toolLoop:
 			ts.sessionKey,
 			ts.opts.Dispatch.SessionScope,
 		)
+		if inbound := ts.opts.Dispatch.InboundContext; inbound != nil {
+			execCtx = withMCPHeadersFromRaw(execCtx, inbound.Raw)
+		}
 		toolResult := ts.agent.Tools.ExecuteWithContext(
 			execCtx,
 			toolName,

pkg/channels/pico/pico.go

@@ -44,6 +44,23 @@ var allowedInlineImageMIMETypes = map[string]struct{}{
 	"image/bmp":  {},
 }
 
+// reservedRawKeys are internal metadata keys that cannot be overwritten by client-provided metadata.
+// This prevents clients from spoofing internal identity or session information.
+var reservedRawKeys = map[string]struct{}{
+	"platform":     {},
+	"session_id":   {},
+	"conn_id":      {},
+	"message_id":   {},
+	"sender_id":    {},
+	"chat_id":      {},
+	"message_kind": {},
+}
+
+func isReservedRawKey(key string) bool {
+	_, reserved := reservedRawKeys[key]
+	return reserved
+}
+
 func outboundMessageIsThought(msg bus.OutboundMessage) bool {
 	if len(msg.Context.Raw) == 0 {
 		return false
@@ -948,6 +965,17 @@ func (c *PicoChannel) handleMessageSend(pc *picoConn, msg PicoMessage) {
 		"conn_id":    pc.id,
 	}
 
+	if payloadMeta, ok := msg.Payload["metadata"].(map[string]any); ok {
+		for k, v := range payloadMeta {
+			if s, ok := v.(string); ok && s != "" {
+				if isReservedRawKey(k) {
+					continue
+				}
+				metadata[k] = s
+			}
+		}
+	}
+
 	logger.DebugCF("pico", "Received message", map[string]any{
 		"session_id": sessionID,
 		"preview":    truncate(content, 50),

pkg/channels/pico/pico_test.go

@@ -547,3 +547,32 @@ func newTestPicoWebSocket(t *testing.T) (*websocket.Conn, <-chan PicoMessage, fu
 	defer resp.Body.Close()
 	return clientConn, received, cleanup
 }
+
+func TestIsReservedRawKey(t *testing.T) {
+	reserved := []string{
+		"platform",
+		"session_id",
+		"conn_id",
+		"message_id",
+		"sender_id",
+		"chat_id",
+		"message_kind",
+	}
+	for _, key := range reserved {
+		if !isReservedRawKey(key) {
+			t.Errorf("isReservedRawKey(%q) = false, want true", key)
+		}
+	}
+
+	allowed := []string{
+		"mcp:Authorization",
+		"mcp:X-Custom-Header",
+		"custom_field",
+		"user_data",
+	}
+	for _, key := range allowed {
+		if isReservedRawKey(key) {
+			t.Errorf("isReservedRawKey(%q) = true, want false", key)
+		}
+	}
+}

pkg/config/config.go

@@ -1110,6 +1110,35 @@ func (c *SkillRegistryConfig) DecodeParam(target any) error {
 	return json.Unmarshal(data, target)
 }
 
+// DynamicHeadersConfig controls which headers can be forwarded from channel context to MCP servers.
+// By default (when nil or empty), no dynamic headers are allowed - this is secure by default.
+type DynamicHeadersConfig struct {
+	// Allowed is a list of header names that may be forwarded from channel context.
+	// Only headers explicitly listed here will be passed to this MCP server.
+	// Header names are matched case-insensitively.
+	Allowed []string `json:"allowed,omitempty"`
+	// MaxCount limits the number of dynamic headers per request (default: 10).
+	MaxCount int `json:"max_count,omitempty"`
+	// MaxValueLen limits the maximum length of each header value (default: 4096).
+	MaxValueLen int `json:"max_value_len,omitempty"`
+}
+
+// GetMaxCount returns the max header count, defaulting to 10.
+func (c *DynamicHeadersConfig) GetMaxCount() int {
+	if c == nil || c.MaxCount <= 0 {
+		return 10
+	}
+	return c.MaxCount
+}
+
+// GetMaxValueLen returns the max value length, defaulting to 4096.
+func (c *DynamicHeadersConfig) GetMaxValueLen() int {
+	if c == nil || c.MaxValueLen <= 0 {
+		return 4096
+	}
+	return c.MaxValueLen
+}
+
 // MCPServerConfig defines configuration for a single MCP server
 type MCPServerConfig struct {
 	// Enabled indicates whether this MCP server is active
@@ -1132,6 +1161,9 @@ type MCPServerConfig struct {
 	URL string `json:"url,omitempty"`
 	// Headers are HTTP headers to send with requests (sse/http only)
 	Headers map[string]string `json:"headers,omitempty"`
+	// DynamicHeaders controls which headers from channel context can be forwarded to this server.
+	// When nil or empty, no dynamic headers are allowed (secure by default).
+	DynamicHeaders *DynamicHeadersConfig `json:"dynamic_headers,omitempty"`
 }
 
 // MCPConfig defines configuration for all MCP servers

pkg/mcp/manager.go

@@ -18,12 +18,14 @@ import (
 	"github.com/sipeed/picoclaw/pkg/config"
 	runtimeevents "github.com/sipeed/picoclaw/pkg/events"
 	"github.com/sipeed/picoclaw/pkg/logger"
+	toolshared "github.com/sipeed/picoclaw/pkg/tools/shared"
 )
 
 // headerTransport is an http.RoundTripper that adds custom headers to requests
 type headerTransport struct {
-	base    http.RoundTripper
-	headers map[string]string
+	base           http.RoundTripper
+	headers        map[string]string
+	dynamicHeaders *config.DynamicHeadersConfig
 }
 
 func expandHomeCommandPath(command string) string {
@@ -45,15 +47,16 @@ func expandHomeCommandPath(command string) string {
 }
 
 func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	// Clone the request to avoid modifying the original
 	req = req.Clone(req.Context())
-
-	// Add custom headers
 	for key, value := range t.headers {
 		req.Header.Set(key, value)
 	}
-
-	// Use the base transport
+	if dynamic := toolshared.MCPHeaders(req.Context()); len(dynamic) > 0 {
+		filtered := t.filterDynamicHeaders(dynamic)
+		for key, value := range filtered {
+			req.Header.Set(key, value)
+		}
+	}
 	base := t.base
 	if base == nil {
 		base = http.DefaultTransport
@@ -61,6 +64,39 @@ func (t *headerTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	return base.RoundTrip(req)
 }
 
+// filterDynamicHeaders applies the allowlist and limits from DynamicHeadersConfig.
+// Returns nil if no dynamic headers config is set (secure by default).
+func (t *headerTransport) filterDynamicHeaders(headers map[string]string) map[string]string {
+	if t.dynamicHeaders == nil || len(t.dynamicHeaders.Allowed) == 0 {
+		return nil
+	}
+
+	allowedSet := make(map[string]struct{}, len(t.dynamicHeaders.Allowed))
+	for _, h := range t.dynamicHeaders.Allowed {
+		allowedSet[strings.ToLower(h)] = struct{}{}
+	}
+
+	maxCount := t.dynamicHeaders.GetMaxCount()
+	maxValueLen := t.dynamicHeaders.GetMaxValueLen()
+
+	result := make(map[string]string)
+	count := 0
+	for key, value := range headers {
+		if count >= maxCount {
+			break
+		}
+		if _, ok := allowedSet[strings.ToLower(key)]; !ok {
+			continue
+		}
+		if len(value) > maxValueLen {
+			value = value[:maxValueLen]
+		}
+		result[key] = value
+		count++
+	}
+	return result
+}
+
 // loadEnvFile loads environment variables from a file in .env format
 // Each line should be in the format: KEY=value
 // Lines starting with # are comments
@@ -381,15 +417,14 @@ func connectServer(
 			DisableStandaloneSSE: disableStandaloneSSE,
 		}
 
-		// Add custom headers if provided
+		sseTransport.HTTPClient = &http.Client{
+			Transport: &headerTransport{
+				base:           http.DefaultTransport,
+				headers:        cfg.Headers,
+				dynamicHeaders: cfg.DynamicHeaders,
+			},
+		}
 		if len(cfg.Headers) > 0 {
-			// Create a custom HTTP client with header-injecting transport
-			sseTransport.HTTPClient = &http.Client{
-				Transport: &headerTransport{
-					base:    http.DefaultTransport,
-					headers: cfg.Headers,
-				},
-			}
 			logger.DebugCF("mcp", "Added custom HTTP headers",
 				map[string]any{
 					"server":       name,

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,300 @@ 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"},
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed: []string{"Authorization", "X-Custom"},
+		},
+	}
+
+	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)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	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"},
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed: []string{"Authorization"},
+		},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"Authorization": "Bearer dynamic",
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	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)
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	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)
+	}
+}
+
+func TestHeaderTransport_NoAllowlist_BlocksDynamic(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
+	})
+
+	// No dynamicHeaders config = secure by default, blocks all dynamic headers
+	transport := &headerTransport{
+		base:    base,
+		headers: map[string]string{"X-Static": "val"},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"Authorization": "Bearer should-be-blocked",
+		"X-Custom":      "also-blocked",
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	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 blocked without allowlist, got %q", got)
+	}
+	if got := captured.Get("X-Custom"); got != "" {
+		t.Errorf("X-Custom should be blocked without allowlist, got %q", got)
+	}
+}
+
+func TestHeaderTransport_AllowlistFilters(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,
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed: []string{"X-Allowed-Header"},
+		},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"X-Allowed-Header": "this-passes",
+		"X-Blocked-Header": "this-is-blocked",
+		"Authorization":    "also-blocked",
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	if got := captured.Get("X-Allowed-Header"); got != "this-passes" {
+		t.Errorf("X-Allowed-Header = %q, want %q", got, "this-passes")
+	}
+	if got := captured.Get("X-Blocked-Header"); got != "" {
+		t.Errorf("X-Blocked-Header should be blocked, got %q", got)
+	}
+	if got := captured.Get("Authorization"); got != "" {
+		t.Errorf("Authorization should be blocked, got %q", got)
+	}
+}
+
+func TestHeaderTransport_AllowlistCaseInsensitive(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,
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed: []string{"X-Grafana-Token"}, // lowercase in allowlist
+		},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"x-grafana-token": "token-value", // different case in request
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	if got := captured.Get("X-Grafana-Token"); got != "token-value" {
+		t.Errorf("X-Grafana-Token = %q, want %q (case-insensitive match)", got, "token-value")
+	}
+}
+
+func TestHeaderTransport_MaxCountLimit(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,
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed:  []string{"X-Header-1", "X-Header-2", "X-Header-3"},
+			MaxCount: 2,
+		},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"X-Header-1": "val1",
+		"X-Header-2": "val2",
+		"X-Header-3": "val3",
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	// Only 2 headers should be set (MaxCount=2), but we can't predict which due to map iteration order
+	count := 0
+	for _, h := range []string{"X-Header-1", "X-Header-2", "X-Header-3"} {
+		if captured.Get(h) != "" {
+			count++
+		}
+	}
+	if count != 2 {
+		t.Errorf("Expected exactly 2 headers due to MaxCount=2, got %d", count)
+	}
+}
+
+func TestHeaderTransport_MaxValueLenLimit(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,
+		dynamicHeaders: &config.DynamicHeadersConfig{
+			Allowed:     []string{"X-Long-Header"},
+			MaxValueLen: 10,
+		},
+	}
+
+	ctx := toolshared.WithMCPHeaders(context.Background(), map[string]string{
+		"X-Long-Header": "this-value-is-way-too-long-and-should-be-truncated",
+	})
+	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, "http://example.com", nil)
+
+	resp, err := transport.RoundTrip(req)
+	if err != nil {
+		t.Fatalf("RoundTrip() error = %v", err)
+	}
+	resp.Body.Close()
+
+	got := captured.Get("X-Long-Header")
+	if len(got) != 10 {
+		t.Errorf("X-Long-Header length = %d, want 10 (truncated)", len(got))
+	}
+	if got != "this-value" {
+		t.Errorf("X-Long-Header = %q, want %q", got, "this-value")
+	}
+}
+
+type roundTripFunc func(*http.Request) (*http.Response, error)
+
+func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
+	return f(req)
+}

pkg/tools/shared/base.go

@@ -50,6 +50,7 @@ var (
 	ctxKeyAgentID          = &toolCtxKey{"agentID"}
 	ctxKeySessionKey       = &toolCtxKey{"sessionKey"}
 	ctxKeySessionScope     = &toolCtxKey{"sessionScope"}
+	ctxKeyMCPHeaders       = &toolCtxKey{"mcpHeaders"}
 )
 
 // WithToolContext returns a child context carrying channel and chatID.
@@ -130,6 +131,33 @@ func ToolSessionScope(ctx context.Context) *session.SessionScope {
 	return session.CloneScope(scope)
 }
 
+// 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 {
+	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)
+	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.
 // When an async tool finishes its work, it calls this callback with the result.
 //

pkg/tools/shared_facade.go

@@ -61,6 +61,14 @@ func WithToolSessionContext(
 	return toolshared.WithToolSessionContext(ctx, agentID, sessionKey, scope)
 }
 
+func WithMCPHeaders(ctx context.Context, headers map[string]string) context.Context {
+	return toolshared.WithMCPHeaders(ctx, headers)
+}
+
+func MCPHeaders(ctx context.Context) map[string]string {
+	return toolshared.MCPHeaders(ctx)
+}
+
 func ToolChannel(ctx context.Context) string {
 	return toolshared.ToolChannel(ctx)
 }