diff --git a/CHANGELOG.md b/CHANGELOG.md index 976620c00d..62dfef8d26 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -421,6 +421,12 @@ This release adds sandbox authentication improvements, compaction-model context ## [Unreleased] +- Harden served-agent safety and network controls: + - A2A, MCP HTTP, and chat now default to restricted tool safety; autonomous execution requires `--safety autonomous`, and `safety: autonomous` in YAML fails startup with guidance to use that flag. + - Non-loopback listeners require authentication or `--insecure-no-auth`; Unix sockets remain exempt. A2A and MCP HTTP use `--auth-token`, while chat continues to use `--api-key`; A2A and chat can explicitly allow browser origins with `--cors-origin`. + - A2A context IDs cannot access sessions created by another serve surface. The session-schema migration records session origins; older binaries reject upgraded databases with a newer-database error. + - `mcp.CreateToolHandler` now requires an explicit safety policy. + ## What's New - Splits the sidebar's Token Usage click target in two: clicking the token/context part (glyph, token count, context `%`, the "compacting…" marker, or the `⚠ capped` marker) opens the `/context` dialog, while clicking the cost part (`$` figure, sub-session count) keeps opening `/cost`; the "Token Usage" section title is no longer clickable diff --git a/cmd/root/a2a.go b/cmd/root/a2a.go index ed1c9b55d1..6ba84e3f2f 100644 --- a/cmd/root/a2a.go +++ b/cmd/root/a2a.go @@ -1,19 +1,33 @@ package root import ( + "errors" + "fmt" + "io" + "net" + "strings" + "github.com/spf13/cobra" "github.com/docker/docker-agent/pkg/a2a" "github.com/docker/docker-agent/pkg/cli" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/httpsec" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/telemetry" ) type a2aFlags struct { - agentName string - listenAddr string - sessionDB string - runConfig config.RuntimeConfig + agentName string + listenAddr string + sessionDB string + safety string + authToken string + corsOrigin string + insecureNoAuth bool + stdout io.Writer + runConfig config.RuntimeConfig } func newA2ACmd() *cobra.Command { @@ -32,6 +46,10 @@ func newA2ACmd() *cobra.Command { cmd.PersistentFlags().StringVarP(&flags.agentName, "agent", "a", "", "Name of the agent to run (defaults to the team's first agent)") cmd.PersistentFlags().StringVarP(&flags.listenAddr, "listen", "l", "127.0.0.1:8082", "Address to listen on") cmd.PersistentFlags().StringVarP(&flags.sessionDB, "session-db", "s", "", "Path to the session database (default: /session.db)") + cmd.PersistentFlags().StringVar(&flags.safety, "safety", "", "Tool safety policy (strict, balanced, restricted, autonomous)") + cmd.PersistentFlags().StringVar(&flags.authToken, "auth-token", "", "Bearer token required for all A2A requests") + cmd.PersistentFlags().StringVar(&flags.corsOrigin, "cors-origin", "", "Allowed browser origin(s), comma-separated; empty disables CORS") + cmd.PersistentFlags().BoolVar(&flags.insecureNoAuth, "insecure-no-auth", false, "Allow unauthenticated non-loopback binding (insecure)") addRuntimeConfigFlags(cmd, &flags.runConfig) return cmd @@ -44,7 +62,22 @@ func (f *a2aFlags) runA2ACommand(cmd *cobra.Command, args []string) (commandErr telemetry.TrackCommandError(ctx, "serve", append([]string{"a2a"}, args...), commandErr) }() - out := cli.NewPrinter(cmd.OutOrStdout()) + if err := validateSafetyFlag(f.safety); err != nil { + return err + } + if f.corsOrigin != "" { + if _, err := httpsec.ParseOrigins(f.corsOrigin); err != nil { + return fmt.Errorf("invalid --cors-origin: %w", err) + } + } + if !isLoopbackListenAddr(f.listenAddr) && f.authToken == "" && !f.insecureNoAuth { + return errors.New("non-loopback A2A listeners require --auth-token or --insecure-no-auth") + } + + out := cli.NewPrinter(f.stdout) + if f.stdout == nil { + out = cli.NewPrinter(cmd.OutOrStdout()) + } agentFilename := args[0] ln, cleanup, err := newListener(ctx, f.listenAddr) @@ -54,5 +87,28 @@ func (f *a2aFlags) runA2ACommand(cmd *cobra.Command, args []string) (commandErr defer cleanup() out.Println("Listening on", ln.Addr().String()) - return a2a.Run(ctx, agentFilename, f.agentName, sessionDBPath(f.sessionDB), &f.runConfig, ln) + return a2a.Run(ctx, agentFilename, f.agentName, sessionDBPath(f.sessionDB), &f.runConfig, ln, a2a.RunOptions{ + CLISafety: session.SafetyPolicy(f.safety), + AuthToken: f.authToken, + CORSOrigin: f.corsOrigin, + OnSafetyPolicy: func(resolved servesafety.Resolved) { + out.Printf("Tool safety policy: %s (source: %s)\n", resolved.Policy, resolved.Source) + }, + }) +} + +func isLoopbackListenAddr(addr string) bool { + if strings.HasPrefix(addr, "unix://") { + return true + } + host, _, err := net.SplitHostPort(addr) + if err != nil { + return false + } + host = strings.Trim(host, "[]") + if strings.EqualFold(host, "localhost") { + return true + } + ip := net.ParseIP(host) + return ip != nil && ip.IsLoopback() } diff --git a/cmd/root/a2a_test.go b/cmd/root/a2a_test.go new file mode 100644 index 0000000000..99498a43b1 --- /dev/null +++ b/cmd/root/a2a_test.go @@ -0,0 +1,28 @@ +package root + +import "testing" + +func TestIsLoopbackListenAddr(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + addr string + want bool + }{ + {"127.0.0.1:8082", true}, + {"[::1]:8082", true}, + {"localhost:8082", true}, + {"unix:///tmp/agent.sock", true}, + {"unix://", true}, + {":8082", false}, + {"0.0.0.0:8082", false}, + {"[::]:8082", false}, + {"192.168.1.1:8082", false}, + } { + t.Run(tc.addr, func(t *testing.T) { + if got := isLoopbackListenAddr(tc.addr); got != tc.want { + t.Errorf("isLoopbackListenAddr(%q) = %v, want %v", tc.addr, got, tc.want) + } + }) + } +} diff --git a/cmd/root/chat.go b/cmd/root/chat.go index b6003476d7..f3896c69da 100644 --- a/cmd/root/chat.go +++ b/cmd/root/chat.go @@ -1,6 +1,8 @@ package root import ( + "errors" + "fmt" "os" "time" @@ -9,6 +11,8 @@ import ( "github.com/docker/docker-agent/pkg/chatserver" "github.com/docker/docker-agent/pkg/cli" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/telemetry" ) @@ -23,6 +27,8 @@ type chatFlags struct { conversationsMaxItems int conversationTTL time.Duration maxIdleRuntimes int + safety string + insecureNoAuth bool runConfig config.RuntimeConfig } @@ -48,6 +54,8 @@ agent without any custom integration.`, cmd.Flags().StringVar(&flags.corsOrigin, "cors-origin", "", "Allowed CORS origin (e.g. https://example.com); empty disables CORS entirely") cmd.Flags().StringVar(&flags.apiKey, "api-key", "", "Required Bearer token clients must present (Authorization: Bearer ); empty disables auth") cmd.Flags().StringVar(&flags.apiKeyEnv, "api-key-env", "", "Read the API key from this environment variable instead of the command line") + cmd.Flags().StringVar(&flags.safety, "safety", "", "Tool safety policy (strict, balanced, restricted, autonomous)") + cmd.Flags().BoolVar(&flags.insecureNoAuth, "insecure-no-auth", false, "Allow unauthenticated non-loopback binding (insecure)") cmd.Flags().Int64Var(&flags.maxRequestSize, "max-request-size", 1<<20, "Maximum request body size in bytes (default 1 MiB)") cmd.Flags().DurationVar(&flags.requestTimeout, "request-timeout", 5*time.Minute, "Per-request timeout (covers model + tool calls + streaming)") cmd.Flags().IntVar(&flags.conversationsMaxItems, "conversations-max", 0, "Cache up to N conversations server-side, keyed by X-Conversation-Id (0 disables; clients must resend full history)") @@ -65,6 +73,21 @@ func (f *chatFlags) runChatCommand(cmd *cobra.Command, args []string) (commandEr telemetry.TrackCommandError(ctx, "serve", append([]string{"chat"}, args...), commandErr) }() + if err := validateSafetyFlag(f.safety); err != nil { + return err + } + + apiKey := f.apiKey + if f.apiKeyEnv != "" { + apiKey = os.Getenv(f.apiKeyEnv) + if apiKey == "" { + return fmt.Errorf("environment variable %q is empty or not set", f.apiKeyEnv) + } + } + if !isLoopbackListenAddr(f.listenAddr) && apiKey == "" && !f.insecureNoAuth { + return errors.New("non-loopback chat listeners require --api-key, --api-key-env, or --insecure-no-auth") + } + out := cli.NewPrinter(cmd.OutOrStdout()) agentFilename := args[0] @@ -77,13 +100,6 @@ func (f *chatFlags) runChatCommand(cmd *cobra.Command, args []string) (commandEr out.Println("Listening on", ln.Addr().String()) out.Println("OpenAI-compatible chat completions endpoint: http://" + ln.Addr().String() + "/v1/chat/completions") - apiKey := f.apiKey - if f.apiKeyEnv != "" { - if v := os.Getenv(f.apiKeyEnv); v != "" { - apiKey = v - } - } - return chatserver.Run(ctx, agentFilename, chatserver.Options{ AgentName: f.agentName, RunConfig: &f.runConfig, @@ -94,5 +110,9 @@ func (f *chatFlags) runChatCommand(cmd *cobra.Command, args []string) (commandEr ConversationsMaxSessions: f.conversationsMaxItems, ConversationTTL: f.conversationTTL, MaxIdleRuntimes: f.maxIdleRuntimes, + CLISafety: session.SafetyPolicy(f.safety), + OnSafetyPolicy: func(resolved servesafety.Resolved) { + out.Printf("Tool safety policy: %s (source: %s)\n", resolved.Policy, resolved.Source) + }, }, ln) } diff --git a/cmd/root/chat_test.go b/cmd/root/chat_test.go new file mode 100644 index 0000000000..179274383e --- /dev/null +++ b/cmd/root/chat_test.go @@ -0,0 +1,28 @@ +package root + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestChatRejectsUnauthenticatedNonLoopbackBind(t *testing.T) { + t.Parallel() + + cmd := newChatCmd() + cmd.SetArgs([]string{"agent.yaml", "--listen", "0.0.0.0:8083"}) + err := cmd.ExecuteContext(t.Context()) + require.Error(t, err) + assert.Contains(t, err.Error(), "require --api-key, --api-key-env, or --insecure-no-auth") +} + +func TestChatRejectsEmptyAPIKeyEnvironmentVariable(t *testing.T) { + t.Setenv("CHAT_API_KEY", "") + + cmd := newChatCmd() + cmd.SetArgs([]string{"agent.yaml", "--api-key-env", "CHAT_API_KEY"}) + err := cmd.ExecuteContext(t.Context()) + require.Error(t, err) + assert.Contains(t, err.Error(), "CHAT_API_KEY") +} diff --git a/cmd/root/mcp.go b/cmd/root/mcp.go index 3d8d3bf9a3..1aef897d2f 100644 --- a/cmd/root/mcp.go +++ b/cmd/root/mcp.go @@ -3,21 +3,27 @@ package root import ( "context" "errors" + "fmt" "github.com/spf13/cobra" "github.com/docker/docker-agent/pkg/config" "github.com/docker/docker-agent/pkg/mcp" "github.com/docker/docker-agent/pkg/runregistry" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/telemetry" ) type mcpFlags struct { - agentName string - http bool - listenAddr string - attach string - runConfig config.RuntimeConfig + agentName string + http bool + listenAddr string + attach string + safety string + authToken string + insecureNoAuth bool + runConfig config.RuntimeConfig } func newMCPCmd() *cobra.Command { @@ -41,6 +47,9 @@ func newMCPCmd() *cobra.Command { cmd.PersistentFlags().StringVarP(&flags.listenAddr, "listen", "l", "127.0.0.1:8081", "Address to listen on") cmd.PersistentFlags().StringVar(&flags.attach, "attach", "", "Attach to a running TUI run by pid, address, or session id (or empty for the most recent)") cmd.PersistentFlags().Lookup("attach").NoOptDefVal = "latest" + cmd.PersistentFlags().StringVar(&flags.safety, "safety", "", "Tool safety policy (strict, balanced, restricted, autonomous); only valid with --http") + cmd.PersistentFlags().StringVar(&flags.authToken, "auth-token", "", "Bearer token required for HTTP MCP requests; only valid with --http") + cmd.PersistentFlags().BoolVar(&flags.insecureNoAuth, "insecure-no-auth", false, "Allow unauthenticated non-loopback HTTP binding (insecure); only valid with --http") cmd.PersistentFlags().StringVar(&flags.runConfig.MCPToolName, "tool-name", "", "Override the MCP tool identifier clients call (defaults to agent name); only valid when exposing a single agent") cmd.PersistentFlags().DurationVar(&flags.runConfig.MCPKeepAlive, "mcp-keepalive", 0, "Interval between MCP keep-alive pings (e.g. 30s); 0 disables keep-alive") addRuntimeConfigFlags(cmd, &flags.runConfig) @@ -56,9 +65,19 @@ func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) (commandErr }() if f.attach != "" { + if f.http || f.safety != "" || f.authToken != "" || f.insecureNoAuth { + return errors.New("--http-only safety and authentication flags cannot be used with --attach") + } return f.runAttach(ctx) } + if !f.http && (f.safety != "" || f.authToken != "" || f.insecureNoAuth) { + return errors.New("--safety, --auth-token, and --insecure-no-auth require --http") + } + if err := validateSafetyFlag(f.safety); err != nil { + return err + } + if len(args) == 0 { return errors.New("agent file is required (or use --attach)") } @@ -68,13 +87,23 @@ func (f *mcpFlags) runMCPCommand(cmd *cobra.Command, args []string) (commandErr return mcp.StartMCPServer(ctx, agentFilename, f.agentName, &f.runConfig) } + if !isLoopbackListenAddr(f.listenAddr) && f.authToken == "" && !f.insecureNoAuth { + return errors.New("non-loopback MCP HTTP listeners require --auth-token or --insecure-no-auth") + } + ln, cleanup, err := newListener(ctx, f.listenAddr) if err != nil { return err } defer cleanup() - return mcp.StartHTTPServer(ctx, agentFilename, f.agentName, &f.runConfig, ln) + return mcp.StartHTTPServer(ctx, agentFilename, f.agentName, &f.runConfig, ln, mcp.HTTPOptions{ + CLISafety: session.SafetyPolicy(f.safety), + AuthToken: f.authToken, + OnSafetyPolicy: func(resolved servesafety.Resolved) { + fmt.Fprintf(cmd.OutOrStdout(), "Tool safety policy: %s (source: %s)\n", resolved.Policy, resolved.Source) + }, + }) } func (f *mcpFlags) runAttach(ctx context.Context) error { diff --git a/cmd/root/mcp_test.go b/cmd/root/mcp_test.go new file mode 100644 index 0000000000..8f2e35bf11 --- /dev/null +++ b/cmd/root/mcp_test.go @@ -0,0 +1,39 @@ +package root + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestMCPHTTPSafetyAndAuthenticationFlagsRequireHTTP(t *testing.T) { + t.Parallel() + + for _, args := range [][]string{ + {"agent.yaml", "--safety", "restricted"}, + {"agent.yaml", "--auth-token", "secret"}, + {"agent.yaml", "--insecure-no-auth"}, + {"--attach", "--safety", "restricted"}, + } { + cmd := newMCPCmd() + cmd.SetArgs(args) + err := cmd.ExecuteContext(t.Context()) + require.Error(t, err) + if len(args) > 0 && args[0] == "--attach" { + assert.Contains(t, err.Error(), "--http-only") + } else { + assert.Contains(t, err.Error(), "require --http") + } + } +} + +func TestMCPHTTPRejectsUnauthenticatedNonLoopbackBind(t *testing.T) { + t.Parallel() + + cmd := newMCPCmd() + cmd.SetArgs([]string{"agent.yaml", "--http", "--listen", "0.0.0.0:8081"}) + err := cmd.ExecuteContext(t.Context()) + require.Error(t, err) + assert.Contains(t, err.Error(), "require --auth-token or --insecure-no-auth") +} diff --git a/cmd/root/run.go b/cmd/root/run.go index 887ab3d1ef..5d8a170837 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -264,13 +264,8 @@ func (f *runExecFlags) runRunCommand(cmd *cobra.Command, args []string) (command } } - // Same early-failure treatment for --safety: a typo must not - // silently collapse to strict. Legacy policy aliases are a wire - // compat concern; the new flag only takes the canonical modes. - switch session.SafetyPolicy(f.safety) { - case "", session.SafetyPolicyStrict, session.SafetyPolicyBalanced, session.SafetyPolicyRestricted, session.SafetyPolicyAutonomous: - default: - return fmt.Errorf("invalid --safety value %q (valid: strict, balanced, restricted, autonomous)", f.safety) + if err := validateSafetyFlag(f.safety); err != nil { + return err } f.safetyChanged = cmd.Flags().Changed("safety") f.yoloChanged = cmd.Flags().Changed("yolo") @@ -1359,6 +1354,16 @@ func stopToolSets(ctx context.Context, t toolStopper) { } } +// validateSafetyFlag rejects non-canonical --safety values consistently across commands. +func validateSafetyFlag(value string) error { + switch session.SafetyPolicy(value) { + case "", session.SafetyPolicyStrict, session.SafetyPolicyBalanced, session.SafetyPolicyRestricted, session.SafetyPolicyAutonomous: + return nil + default: + return fmt.Errorf("invalid --safety value %q (valid: strict, balanced, restricted, autonomous)", value) + } +} + // validateTheme reports whether ref names a loadable theme. It is used to // fail fast on an explicit --theme value, listing the available themes so the // user can correct a typo. The "auto" sentinel is accepted as-is: it resolves diff --git a/cmd/root/safety_test.go b/cmd/root/safety_test.go new file mode 100644 index 0000000000..159cbbd082 --- /dev/null +++ b/cmd/root/safety_test.go @@ -0,0 +1,20 @@ +package root + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestValidateSafetyFlag(t *testing.T) { + t.Parallel() + + for _, value := range []string{"", "strict", "balanced", "restricted", "autonomous"} { + t.Run(value, func(t *testing.T) { + assert.NoError(t, validateSafetyFlag(value)) + }) + } + + err := validateSafetyFlag("unsafe") + assert.EqualError(t, err, `invalid --safety value "unsafe" (valid: strict, balanced, restricted, autonomous)`) +} diff --git a/docs/community/troubleshooting/index.md b/docs/community/troubleshooting/index.md index ac854ba2c0..1e73169e06 100644 --- a/docs/community/troubleshooting/index.md +++ b/docs/community/troubleshooting/index.md @@ -228,6 +228,13 @@ Docker Agent validates config at startup and reports errors with line numbers. C ## Session & Connectivity Issues +### Downgrade fails with a newer-database error + +If an older Docker Agent binary cannot open the session database after an upgrade, +the database may contain a schema migration that the older binary does not know. +Restore a database created by the older version, or use a binary that includes the +migration. + ### Port conflicts When running Docker Agent as an API server or MCP server, ensure the port is not already in use: diff --git a/docs/features/a2a/index.md b/docs/features/a2a/index.md index 2e82eecd2c..a09f5b0065 100644 --- a/docs/features/a2a/index.md +++ b/docs/features/a2a/index.md @@ -49,6 +49,42 @@ $ docker agent serve a2a myorg/agent:tag | `--hook-session-end ` | (none) | Add a session-end hook (repeatable). | | `--hook-on-user-input ` | (none) | Add an on-user-input hook (repeatable). | | `--hook-stop ` | (none) | Add a stop hook, fired when the model finishes responding (repeatable). | +| `--auth-token ` | (none) | Bearer token required for agent-card and invocation requests. | +| `--cors-origin ` | (none) | Allowed browser origins, comma-separated; empty disables CORS. | +| `--insecure-no-auth` | `false` | Allow an unauthenticated non-loopback listener (unsafe). | +| `--safety ` | `restricted` | Tool safety policy; `autonomous` is permitted only through this explicit CLI flag. | + +## Authentication and network exposure + +Loopback listeners may run without authentication. Non-loopback listeners require +`--auth-token` unless `--insecure-no-auth` explicitly acknowledges the exposure. +Clients must send `Authorization: Bearer ` for both agent-card discovery +and JSON-RPC invocation. Configure browser access with `--cors-origin`; it accepts +comma-separated literal origins or `~`-prefixed regular expressions and permits +credentials only for matching origins. + +```bash +$ docker agent serve a2a ./agent.yaml --auth-token "$A2A_TOKEN" \ + --cors-origin http://localhost:3000 +``` + +## Tool safety and migration + +A2A sessions default to the `restricted` tool safety policy. Precedence is the +`--safety` flag, then agent YAML, then runtime YAML. YAML may select `strict`, +`balanced`, or `restricted`; `safety: autonomous` stops startup and directs the +operator to `--safety autonomous`. That CLI flag is the only deliberate opt-in +to autonomous tool execution. + +Existing deployments should choose an explicit policy before upgrading. Migration +027 labels pre-existing sessions as `run`, so they cannot be resumed through +`/invoke`; clients must start new A2A contexts. An A2A context ID that collides +with another session is rejected without changing that session. + +Downgrading to a binary that predates migration 027 fails because the session +database has a newer schema (`ErrNewerDatabase`). Restore an older database, or +use a binary that includes the migration. Revert changes without removing the +migration catalogue entry. ## Features diff --git a/docs/features/chat-server/index.md b/docs/features/chat-server/index.md index d7e58a31b4..6b99575680 100644 --- a/docs/features/chat-server/index.md +++ b/docs/features/chat-server/index.md @@ -155,7 +155,9 @@ When a request fails — for example because the model returns an error or the ` ## Authentication -The chat server has **no authentication by default**. To require a Bearer +The chat server defaults to loopback binding. A non-loopback `--listen` address requires `--api-key`, `--api-key-env`, or the explicit `--insecure-no-auth` override. An environment variable selected by `--api-key-env` must be set and non-empty. + +To require a Bearer token, pass `--api-key` (literal value) or `--api-key-env` (name of an environment variable that holds the value): @@ -170,7 +172,11 @@ protected once a key is set. > [!WARNING] > **Public exposure** > -> The default listen address is `127.0.0.1:8083`. If you bind to a non-loopback address, always set `--api-key` or `--api-key-env` — there is no other authentication layer. +> The default listen address is `127.0.0.1:8083`. Non-loopback binding is rejected unless `--api-key`, `--api-key-env`, or `--insecure-no-auth` is supplied. Use the insecure override only behind a trusted authentication boundary. + +## Tool safety + +The chat server resolves its safety policy in this order: `--safety`, agent configuration, runtime configuration, then `restricted`. Cached conversations retain the more restrictive of their prior policy and the server policy, so a continuation cannot regain permissions after the server policy becomes stricter. ## CORS @@ -194,7 +200,9 @@ docker agent serve chat | [flags] | `-l, --listen ` | `127.0.0.1:8083` | Address to listen on. | | `--cors-origin ` | (none) | Allowed CORS origin (e.g. `https://example.com`). Empty disables CORS. | | `--api-key ` | (none) | Required Bearer token clients must present (`Authorization: Bearer `). Empty disables auth. | -| `--api-key-env ` | (none) | Read the API key from this environment variable instead of the command line. | +| `--api-key-env ` | (none) | Read the required API key from this non-empty environment variable. | +| `--insecure-no-auth` | `false` | Permit unauthenticated non-loopback binding. Use only behind a trusted authentication boundary. | +| `--safety ` | `restricted` | Tool safety policy. CLI value overrides agent/runtime configuration. | | `--max-request-size ` | `1048576` (1 MiB) | Maximum request body size in bytes. Requests whose body exceeds this limit are rejected with HTTP 413 (Request Entity Too Large) — see [Troubleshooting: HTTP 413](../../community/troubleshooting/index.md#http-413-request-body-too-large) if you hit this. | | `--request-timeout ` | `5m` | Per-request timeout (covers model + tool calls + streaming). | | `--conversations-max ` | `0` | Cache up to N conversations server-side, keyed by `X-Conversation-Id`. `0` disables — clients must resend history. | diff --git a/docs/features/cli/index.md b/docs/features/cli/index.md index 34618873ac..41db671d40 100644 --- a/docs/features/cli/index.md +++ b/docs/features/cli/index.md @@ -318,6 +318,9 @@ $ docker agent serve mcp [flags] | `-a, --agent ` | (all agents) | Name of the agent to expose. If omitted, every agent in the config is exposed as a separate tool. | | `--tool-name ` | (agent name) | Override the MCP tool identifier clients call; only valid when exposing a single agent. | | `--http` | `false` | Use streaming HTTP transport instead of stdio. | +| `--safety ` | `restricted` | HTTP MCP safety policy; no effect on stdio or `--attach`. | +| `--auth-token ` | (none) | Required Bearer token for HTTP MCP requests. | +| `--insecure-no-auth` | `false` | Permit unauthenticated non-loopback HTTP MCP binding. | | `-l, --listen ` | `127.0.0.1:8081` | Address to listen on (only used with `--http`). | | `--mcp-keepalive `| `0` (disabled) | Interval between MCP keep-alive pings (e.g. `30s`). | | `--attach [target]` | (none) | Attach to a running TUI run by pid, address, or session id; given without a value, selects the most recent run. | @@ -394,7 +397,9 @@ $ docker agent serve chat [flags] | `-l, --listen ` | `127.0.0.1:8083` | Address to listen on. | | `--cors-origin ` | (none) | Allowed CORS origin (e.g. `https://example.com`). Empty disables CORS. | | `--api-key ` | (none) | Required Bearer token clients must present (`Authorization: Bearer `). Empty disables auth. | -| `--api-key-env ` | (none) | Read the API key from this environment variable instead of the command line. | +| `--api-key-env ` | (none) | Read the required API key from this non-empty environment variable. | +| `--insecure-no-auth` | `false` | Permit unauthenticated non-loopback binding. | +| `--safety ` | `restricted` | Tool safety policy; CLI value overrides agent/runtime configuration. | | `--max-request-size ` | `1048576` (1 MiB) | Maximum request body size. Requests exceeding this limit are rejected with HTTP 413 — see [Troubleshooting: HTTP 413](../../community/troubleshooting/index.md#http-413-request-body-too-large). | | `--request-timeout ` | `5m` | Per-request timeout (covers model + tool calls + streaming). | | `--conversations-max ` | `0` | Cache up to N conversations server-side, keyed by `X-Conversation-Id`. `0` disables — clients must resend history. | diff --git a/docs/features/evaluation/index.md b/docs/features/evaluation/index.md index e9c6fed4ee..c90701dec2 100644 --- a/docs/features/evaluation/index.md +++ b/docs/features/evaluation/index.md @@ -139,6 +139,10 @@ Docker Agent evaluates agents across three dimensions: | **Relevance** | An LLM judge (configurable via `--judge-model`) evaluates whether each relevance statement is satisfied by the response. | | **Size** | Whether the response length matches the expected size category (S/M/L/XL). | +## Serve-safety verification and rollback + +When changing an agent served over MCP HTTP, chat, or A2A, add an evaluation that attempts an approval-requiring tool call and verifies the resolved safety policy and authentication behavior. Run the evaluation with the same explicit `--safety` setting used in deployment. If a rollout must be reversed, stop the affected listener, restore the prior agent configuration and explicit safety flag, then restart only after confirming non-loopback listeners still require authentication. Do not restore an unauthenticated network listener as a rollback shortcut. + ## Creating Eval Sessions The easiest way to create eval sessions is from real conversations: diff --git a/docs/features/mcp-mode/index.md b/docs/features/mcp-mode/index.md index bc01ef3b1f..9e3ae0bf69 100644 --- a/docs/features/mcp-mode/index.md +++ b/docs/features/mcp-mode/index.md @@ -45,8 +45,8 @@ To expose the MCP server over streaming HTTP instead, pass `--http`: # Streaming HTTP transport on the default 127.0.0.1:8081 $ docker agent serve mcp ./agent.yaml --http -# Override the listen address / port -$ docker agent serve mcp ./agent.yaml --http --listen 0.0.0.0:9090 +# Override the listen address / port; non-loopback HTTP requires authentication +$ docker agent serve mcp ./agent.yaml --http --listen 0.0.0.0:9090 --auth-token "$MCP_BEARER_TOKEN" ``` | Flag | Default | Description | @@ -55,10 +55,17 @@ $ docker agent serve mcp ./agent.yaml --http --listen 0.0.0.0:9090 | `-l`, `--listen` | `127.0.0.1:8081` | Address to listen on when `--http` is enabled. | | `-a`, `--agent` | all agents | Expose a single named agent instead of every agent in the config. | | `--tool-name` | (none) | Override the MCP tool identifier clients call (defaults to agent name); only valid when exposing one agent. | +| `--auth-token` | (none) | Require this Bearer token for HTTP requests. Required for non-loopback HTTP unless explicitly overridden. | +| `--insecure-no-auth` | `false` | Permit unauthenticated non-loopback HTTP. Use only behind a trusted authentication boundary. | +| `--safety` | `restricted` | Tool safety policy for HTTP requests. CLI value overrides agent/runtime configuration. | | `--mcp-keepalive` | `0` | Interval between MCP keep-alive pings (e.g. `30s`); `0` disables keep-alive. | Runtime configuration flags such as `--working-dir`, `--env-from-file`, `--models-gateway`, and hook flags are also available — see the [CLI reference](../cli/index.md). +## HTTP security + +HTTP MCP defaults to loopback binding. A non-loopback `--listen` address requires `--auth-token`; use `--insecure-no-auth` only when a trusted reverse proxy or network boundary authenticates clients. The safety policy is resolved in this order: `--safety`, agent configuration, runtime configuration, then `restricted`. These HTTP-only flags do not affect stdio or `--attach` operation. + ## Using with Claude Desktop Add a configuration to your Claude Desktop MCP settings file: diff --git a/docs/tools/a2a/index.md b/docs/tools/a2a/index.md index 9e405ab494..f07a9fe624 100644 --- a/docs/tools/a2a/index.md +++ b/docs/tools/a2a/index.md @@ -27,6 +27,8 @@ toolsets: X-Tenant: "acme" ``` +The `Authorization` header shown above authenticates to endpoints served with `docker agent serve a2a --auth-token`. + ## Properties | Property | Type | Required | Description | diff --git a/e2e/a2a_test.go b/e2e/a2a_test.go index a79bea5866..759db2f5b5 100644 --- a/e2e/a2a_test.go +++ b/e2e/a2a_test.go @@ -179,7 +179,7 @@ func startA2AServer(t *testing.T, agentFile string, runConfig *config.RuntimeCon done := make(chan struct{}) go func() { defer close(done) - _ = a2aserver.Run(t.Context(), agentFile, "root", sessionDB, runConfig, ln) + _ = a2aserver.Run(t.Context(), agentFile, "root", sessionDB, runConfig, ln, a2aserver.RunOptions{}) }() // Run stops when t.Context() is canceled (just before cleanups run); // closing the listener also covers the window before Serve starts. Wait diff --git a/e2e/mcp_test.go b/e2e/mcp_test.go index e37b58d868..c742ed0781 100644 --- a/e2e/mcp_test.go +++ b/e2e/mcp_test.go @@ -8,6 +8,7 @@ import ( "github.com/docker/docker-agent/pkg/config" "github.com/docker/docker-agent/pkg/mcp" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/teamloader" loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults" ) @@ -26,7 +27,7 @@ func TestMCP_SingleAgent(t *testing.T) { require.NoError(t, team.StopToolSets(ctx)) }) - handler := mcp.CreateToolHandler(team, "root") + handler := mcp.CreateToolHandler(team, "root", session.SafetyPolicyAutonomous) _, output, err := handler(ctx, nil, mcp.ToolInput{ Message: "What is 2+2? Answer in one sentence.", }) @@ -49,7 +50,7 @@ func TestMCP_MultiAgent(t *testing.T) { require.NoError(t, team.StopToolSets(ctx)) }) - handler := mcp.CreateToolHandler(team, "web") + handler := mcp.CreateToolHandler(team, "web", session.SafetyPolicyAutonomous) _, output, err := handler(ctx, nil, mcp.ToolInput{ Message: "Say hello in one sentence.", }) diff --git a/pkg/a2a/adapter.go b/pkg/a2a/adapter.go index 89343f0e10..8f0663b80c 100644 --- a/pkg/a2a/adapter.go +++ b/pkg/a2a/adapter.go @@ -2,6 +2,7 @@ package a2a import ( "cmp" + "errors" "fmt" "iter" "log/slog" @@ -18,6 +19,7 @@ import ( dagent "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/servesafety" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/team" cgenai "github.com/docker/docker-agent/pkg/telemetry/genai" @@ -27,7 +29,7 @@ import ( // newDockerAgentAdapter creates a new ADK agent adapter from a docker agent team and agent name. // When agentName is empty, the team's default agent (one explicitly named "root" if it // exists, otherwise the first agent declared) is used. -func newDockerAgentAdapter(t *team.Team, agentName string, sessStore session.Store) (agent.Agent, error) { +func newDockerAgentAdapter(t *team.Team, agentName string, sessStore session.Store, safety servesafety.Resolved) (agent.Agent, error) { a, err := t.AgentOrDefault(agentName) if err != nil { return nil, fmt.Errorf("failed to get agent %s: %w", agentName, err) @@ -40,13 +42,13 @@ func newDockerAgentAdapter(t *team.Team, agentName string, sessStore session.Sto Name: agentName, Description: desc, Run: func(ctx agent.InvocationContext) iter.Seq2[*adksession.Event, error] { - return runDockerAgent(ctx, t, agentName, a, sessStore) + return runDockerAgent(ctx, t, agentName, a, sessStore, safety) }, }) } // runDockerAgent executes a docker agent and returns ADK session events -func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string, a *dagent.Agent, sessStore session.Store) iter.Seq2[*adksession.Event, error] { +func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string, a *dagent.Agent, sessStore session.Store, safety servesafety.Resolved) iter.Seq2[*adksession.Event, error] { return func(yield func(*adksession.Event, error) bool) { // Decorate the inbound `a2a.message` SERVER span (created by // otelhttp.NewHandler in server.go) with the GenAI semconv @@ -67,31 +69,46 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string, userContent := ctx.UserContent() message := contentToMessage(userContent) - // Use the A2A contextID (exposed as the ADK session ID) as the - // docker-agent session ID so subsequent `run --session ` - // invocations can resume the same conversation. + // Use the A2A context ID as the docker-agent session ID so only future + // A2A invocations with that ID can resume the conversation. sessionID := ctx.Session().ID() var sess *session.Session - if existing, err := sessStore.GetSession(ctx, sessionID); err == nil && existing != nil { + existing, err := sessStore.GetSessionByOrigin(ctx, sessionID, "a2a") + switch { + case err == nil: sess = existing sess.AddMessage(session.UserMessage(message)) - sess.SetSafetyPolicy(session.SafetyPolicyAutonomous) + sess.SetSafetyPolicy(servesafety.ResumeCeiling(sess.GetSafetyPolicy(), safety.Policy)) sess.NonInteractive = true - } else { - workingDir, _ := os.Getwd() - sess = session.New( - session.WithID(sessionID), - session.WithUserMessage(message), - session.WithMaxIterations(a.MaxIterations()), - session.WithMaxConsecutiveToolCalls(a.MaxConsecutiveToolCalls()), - session.WithMaxOldToolCallTokens(a.MaxOldToolCallTokens()), - session.WithMaxToolResultTokens(a.MaxToolResultTokens()), - session.WithToolsApproved(true), - session.WithNonInteractive(true), - session.WithWorkingDir(workingDir), - ) - sess.SetTitle("A2A Session " + sessionID) + case !errors.Is(err, session.ErrNotFound): + yield(nil, fmt.Errorf("look up A2A session: %w", err)) + return + default: + _, err := sessStore.GetSession(ctx, sessionID) + switch { + case err == nil: + yield(nil, errors.New("context ID is not available")) + return + case !errors.Is(err, session.ErrNotFound): + yield(nil, fmt.Errorf("check A2A context ID: %w", err)) + return + default: + workingDir, _ := os.Getwd() + sess = session.New( + session.WithID(sessionID), + session.WithOrigin("a2a"), + session.WithUserMessage(message), + session.WithMaxIterations(a.MaxIterations()), + session.WithMaxConsecutiveToolCalls(a.MaxConsecutiveToolCalls()), + session.WithMaxOldToolCallTokens(a.MaxOldToolCallTokens()), + session.WithMaxToolResultTokens(a.MaxToolResultTokens()), + session.WithSafetyPolicy(safety.Policy), + session.WithNonInteractive(true), + session.WithWorkingDir(workingDir), + ) + sess.SetTitle("A2A Session " + sessionID) + } } // Create runtime diff --git a/pkg/a2a/adapter_run_test.go b/pkg/a2a/adapter_run_test.go index f1ad7ce879..a8f3600d60 100644 --- a/pkg/a2a/adapter_run_test.go +++ b/pkg/a2a/adapter_run_test.go @@ -22,6 +22,7 @@ import ( "github.com/docker/docker-agent/pkg/model/provider" "github.com/docker/docker-agent/pkg/model/provider/base" "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/servesafety" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/tools" @@ -179,9 +180,9 @@ type yieldedEvent struct { err error } -func collectRunEvents(ctx agent.InvocationContext, tm *team.Team, a *dagent.Agent, store session.Store) []yieldedEvent { +func collectRunEvents(ctx agent.InvocationContext, tm *team.Team, a *dagent.Agent, store session.Store, policy session.SafetyPolicy) []yieldedEvent { var out []yieldedEvent - for ev, err := range runDockerAgent(ctx, tm, a.Name(), a, store) { + for ev, err := range runDockerAgent(ctx, tm, a.Name(), a, store, servesafety.Resolved{Policy: policy}) { out = append(out, yieldedEvent{event: ev, err: err}) } return out @@ -203,7 +204,7 @@ func TestRunDockerAgent_StreamsPartialAndFinalEvents(t *testing.T) { store := session.NewInMemorySessionStore() ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-events", "Hi there") - events := collectRunEvents(ctx, tm, root, store) + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) require.Len(t, events, 3) for _, e := range events { @@ -241,7 +242,7 @@ func TestRunDockerAgent_ErrorEventStopsIteration(t *testing.T) { store := session.NewInMemorySessionStore() ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-error", "Hi") - events := collectRunEvents(ctx, tm, root, store) + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) require.Len(t, events, 1) assert.Nil(t, events[0].event) @@ -258,7 +259,7 @@ func TestRunDockerAgent_EmptyStreamEmitsNoFinalEvent(t *testing.T) { store := session.NewInMemorySessionStore() ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-empty", "Hi") - events := collectRunEvents(ctx, tm, root, store) + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) assert.Empty(t, events) } @@ -271,7 +272,7 @@ func TestRunDockerAgent_ConsumerStopsEarly(t *testing.T) { ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-early-stop", "Hi") var events []*adksession.Event - for ev, err := range runDockerAgent(ctx, tm, root.Name(), root, store) { + for ev, err := range runDockerAgent(ctx, tm, root.Name(), root, store, servesafety.Resolved{Policy: session.SafetyPolicyRestricted}) { require.NoError(t, err) events = append(events, ev) break @@ -290,7 +291,7 @@ func TestRunDockerAgent_EndedInvocationStopsIteration(t *testing.T) { ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-ended", "Hi") var events []*adksession.Event - for ev, err := range runDockerAgent(ctx, tm, root.Name(), root, store) { + for ev, err := range runDockerAgent(ctx, tm, root.Name(), root, store, servesafety.Resolved{Policy: session.SafetyPolicyRestricted}) { require.NoError(t, err) events = append(events, ev) // Ending the invocation after the first chunk must stop the @@ -309,7 +310,7 @@ func TestRunDockerAgent_NewSessionUsesA2ASettings(t *testing.T) { store := newRecordingStore() ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-new", "What is Docker?") - events := collectRunEvents(ctx, tm, root, store) + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) require.Len(t, events, 2) updated := store.updatedSessions() @@ -317,8 +318,10 @@ func TestRunDockerAgent_NewSessionUsesA2ASettings(t *testing.T) { sess := updated[0] assert.Equal(t, "a2a-ctx-new", sess.ID) + assert.Equal(t, "a2a", sess.Origin) assert.Equal(t, "A2A Session a2a-ctx-new", sess.Title) - assert.True(t, sess.ToolsApproved) + assert.Equal(t, session.SafetyPolicyRestricted, sess.GetSafetyPolicy()) + assert.False(t, sess.ToolsApproved) assert.True(t, sess.NonInteractive) // runDockerAgent stamps new sessions with the process working directory @@ -336,21 +339,124 @@ func TestRunDockerAgent_NewSessionUsesA2ASettings(t *testing.T) { stored, err := store.GetSession(t.Context(), "a2a-ctx-new") require.NoError(t, err) assert.Equal(t, "a2a-ctx-new", stored.ID) + assert.Equal(t, "a2a", stored.Origin) assert.Equal(t, "A2A Session a2a-ctx-new", stored.Title) } +func TestRunDockerAgent_RejectsNonA2ASessionCollision(t *testing.T) { + t.Parallel() + + for _, origin := range []string{"run", "", "acp"} { + t.Run(origin, func(t *testing.T) { + tm, root := newMockTeam("answer") + store := newRecordingStore() + existing := session.New( + session.WithID("a2a-ctx-collision"), + session.WithOrigin(origin), + session.WithTitle("Private Session"), + session.WithUserMessage("private history"), + ) + require.NoError(t, store.AddSession(t.Context(), existing)) + + ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-collision", "A2A request") + for range 2 { + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) + require.Len(t, events, 1) + require.ErrorContains(t, events[0].err, "context ID is not available") + } + + stored, err := store.GetSession(t.Context(), "a2a-ctx-collision") + require.NoError(t, err) + assert.Equal(t, origin, stored.Origin) + assert.Equal(t, "Private Session", stored.Title) + assert.Len(t, stored.GetAllMessages(), 1) + assert.Empty(t, store.updatedSessions()) + }) + } +} + +func TestRunDockerAgent_ExplicitSafety(t *testing.T) { + t.Parallel() + + for _, policy := range []session.SafetyPolicy{ + session.SafetyPolicyStrict, + session.SafetyPolicyBalanced, + session.SafetyPolicyRestricted, + session.SafetyPolicyAutonomous, + } { + t.Run(string(policy), func(t *testing.T) { + t.Parallel() + + tm, root := newMockTeam("answer") + store := newRecordingStore() + ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-"+string(policy), "What is Docker?") + + collectRunEvents(ctx, tm, root, store, policy) + + updated := store.updatedSessions() + require.NotEmpty(t, updated) + assert.Equal(t, policy, updated[0].GetSafetyPolicy()) + assert.Equal(t, policy == session.SafetyPolicyAutonomous, updated[0].ToolsApproved) + }) + } +} + +func TestRunDockerAgent_ResumedSessionDoesNotExceedServerSafety(t *testing.T) { + t.Parallel() + + tm, root := newMockTeam("answer") + store := newRecordingStore() + existing := session.New( + session.WithID("a2a-ctx-ceiling"), + session.WithOrigin("a2a"), + session.WithSafetyPolicy(session.SafetyPolicyAutonomous), + ) + require.NoError(t, store.AddSession(t.Context(), existing)) + + ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-ceiling", "follow-up question") + collectRunEvents(ctx, tm, root, store, session.SafetyPolicyBalanced) + + assert.Equal(t, session.SafetyPolicyBalanced, existing.GetSafetyPolicy()) + assert.False(t, existing.ToolsApproved) + assert.True(t, existing.NonInteractive) +} + +func TestRunDockerAgent_ResumedSaferSessionIsPreserved(t *testing.T) { + t.Parallel() + + tm, root := newMockTeam("answer") + store := newRecordingStore() + existing := session.New( + session.WithID("a2a-ctx-preserve"), + session.WithOrigin("a2a"), + session.WithSafetyPolicy(session.SafetyPolicyStrict), + ) + require.NoError(t, store.AddSession(t.Context(), existing)) + + ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-preserve", "follow-up question") + collectRunEvents(ctx, tm, root, store, session.SafetyPolicyAutonomous) + + assert.Equal(t, session.SafetyPolicyStrict, existing.GetSafetyPolicy()) + assert.False(t, existing.ToolsApproved) + assert.True(t, existing.NonInteractive) +} + func TestRunDockerAgent_ResumesExistingSession(t *testing.T) { t.Parallel() tm, root := newMockTeam("resumed answer") store := newRecordingStore() - existing := session.New(session.WithID("a2a-ctx-resume"), session.WithTitle("Existing Title")) + existing := session.New( + session.WithID("a2a-ctx-resume"), + session.WithOrigin("a2a"), + session.WithTitle("Existing Title"), + ) require.NoError(t, store.AddSession(t.Context(), existing)) ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-resume", "follow-up question") - events := collectRunEvents(ctx, tm, root, store) + events := collectRunEvents(ctx, tm, root, store, session.SafetyPolicyRestricted) require.Len(t, events, 2) updated := store.updatedSessions() @@ -358,7 +464,8 @@ func TestRunDockerAgent_ResumesExistingSession(t *testing.T) { assert.Same(t, existing, updated[0], "the stored session should be resumed, not recreated") assert.Equal(t, "Existing Title", existing.Title) - assert.True(t, existing.ToolsApproved) + assert.Equal(t, session.SafetyPolicyRestricted, existing.GetSafetyPolicy()) + assert.False(t, existing.ToolsApproved) assert.True(t, existing.NonInteractive) msgs := existing.GetAllMessages() @@ -383,7 +490,7 @@ func TestRunDockerAgent_RuntimeCreationError(t *testing.T) { store := session.NewInMemorySessionStore() ctx := newFakeInvocationContext(t.Context(), "a2a-ctx-no-team", "Hi") - events := collectRunEvents(ctx, emptyTeam, root, store) + events := collectRunEvents(ctx, emptyTeam, root, store, session.SafetyPolicyRestricted) require.Len(t, events, 1) assert.Nil(t, events[0].event) diff --git a/pkg/a2a/adapter_test.go b/pkg/a2a/adapter_test.go index d6eaad9ebf..88298b6b78 100644 --- a/pkg/a2a/adapter_test.go +++ b/pkg/a2a/adapter_test.go @@ -8,6 +8,8 @@ import ( "google.golang.org/genai" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/teamloader" loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults" ) @@ -24,7 +26,7 @@ func TestNewDockerAgentAdapter(t *testing.T) { require.NoError(t, team.StopToolSets(t.Context())) }() - adapter, err := newDockerAgentAdapter(team, "root", nil) + adapter, err := newDockerAgentAdapter(team, "root", nil, servesafety.Resolved{Policy: session.SafetyPolicyRestricted}) require.NoError(t, err) assert.Equal(t, "root", adapter.Name()) @@ -43,7 +45,7 @@ func TestNewCAgentAdapter_NonExistent(t *testing.T) { require.NoError(t, team.StopToolSets(t.Context())) }() - _, err = newDockerAgentAdapter(team, "nonexistent", nil) + _, err = newDockerAgentAdapter(team, "nonexistent", nil, servesafety.Resolved{Policy: session.SafetyPolicyRestricted}) assert.Contains(t, err.Error(), "failed to get agent") } diff --git a/pkg/a2a/server.go b/pkg/a2a/server.go index 44c3f86ce1..4282987ff6 100644 --- a/pkg/a2a/server.go +++ b/pkg/a2a/server.go @@ -20,8 +20,12 @@ import ( adksession "google.golang.org/adk/session" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/httpsec" pathx "github.com/docker/docker-agent/pkg/path" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/session/sqlitestore" + "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/teamloader" loaderdefaults "github.com/docker/docker-agent/pkg/teamloader/defaults" "github.com/docker/docker-agent/pkg/version" @@ -40,7 +44,14 @@ func routableAddr(addr string) string { return addr } -func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runConfig *config.RuntimeConfig, ln net.Listener) error { +type RunOptions struct { + CLISafety session.SafetyPolicy + OnSafetyPolicy func(servesafety.Resolved) + AuthToken string + CORSOrigin string +} + +func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runConfig *config.RuntimeConfig, ln net.Listener, options RunOptions) error { slog.DebugContext(ctx, "Starting A2A server", "source", agentFilename, "agent", agentName, "addr", ln.Addr().String()) agentSource, err := config.Resolve(agentFilename, nil) @@ -58,6 +69,18 @@ func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runCon } }() + selectedAgent, err := t.AgentOrDefault(agentName) + if err != nil { + return fmt.Errorf("failed to get agent: %w", err) + } + resolvedSafety, err := servesafety.Resolve(options.CLISafety, string(selectedAgent.Safety()), string(t.RuntimeSafety())) + if err != nil { + return fmt.Errorf("resolve serve safety policy: %w", err) + } + if options.OnSafetyPolicy != nil { + options.OnSafetyPolicy(resolvedSafety) + } + expandedSessionDB, err := pathx.ExpandHomeDir(sessionDB) if err != nil { return fmt.Errorf("failed to expand session db path: %w", err) @@ -72,15 +95,36 @@ func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runCon } }() - adkAgent, err := newDockerAgentAdapter(t, agentName, sessStore) + baseURL := &url.URL{Scheme: "http", Host: routableAddr(ln.Addr().String())} + slog.DebugContext(ctx, "A2A server listening", "url", baseURL.String()) + + e, err := newServer(t, agentFilename, agentName, sessStore, resolvedSafety, ln.Addr().String(), options) if err != nil { - return fmt.Errorf("failed to create ADK agent adapter: %w", err) + return fmt.Errorf("failed to create A2A server: %w", err) } - baseURL := &url.URL{Scheme: "http", Host: routableAddr(ln.Addr().String())} + // Stop serving when ctx is canceled so Run returns and the deferred + // cleanups (session store, tool sets) release their resources. + stop := context.AfterFunc(ctx, func() { + _ = e.Server.Close() + }) + defer stop() - slog.DebugContext(ctx, "A2A server listening", "url", baseURL.String()) + if err := e.Server.Serve(ln); err != nil && ctx.Err() == nil { + slog.ErrorContext(ctx, "Failed to start server", "error", err) + return err + } + + return nil +} +func newServer(t *team.Team, agentFilename, agentName string, sessStore session.Store, safety servesafety.Resolved, listenAddr string, options RunOptions) (*echo.Echo, error) { + adkAgent, err := newDockerAgentAdapter(t, agentName, sessStore, safety) + if err != nil { + return nil, err + } + + baseURL := &url.URL{Scheme: "http", Host: routableAddr(listenAddr)} name := strings.TrimSuffix(filepath.Base(agentFilename), filepath.Ext(agentFilename)) agentPath := "/invoke" @@ -114,12 +158,16 @@ func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runCon e.HideBanner = true e.HidePort = true - e.Use(middleware.CORSWithConfig(middleware.CORSConfig{ - AllowOrigins: []string{"*"}, - AllowMethods: []string{http.MethodPost, http.MethodOptions}, - AllowHeaders: []string{"Content-Type", "Accept"}, - MaxAge: 86400, - })) + if options.CORSOrigin != "" { + cfg, err := corsMiddlewareConfig(options.CORSOrigin) + if err != nil { + return nil, fmt.Errorf("invalid CORS origin: %w", err) + } + e.Use(middleware.CORSWithConfig(cfg)) + } + if options.AuthToken != "" { + e.Use(bearerAuthMiddleware(options.AuthToken)) + } e.Use(middleware.RequestLogger()) // Wrap both A2A endpoints with otelhttp so the configured W3C @@ -143,17 +191,34 @@ func Run(ctx context.Context, agentFilename, agentName, sessionDB string, runCon e.GET(a2asrv.WellKnownAgentCardPath, echo.WrapHandler(cardHandler)) e.POST(agentPath, echo.WrapHandler(jsonrpcHandler)) - // Stop serving when ctx is canceled so Run returns and the deferred - // cleanups (session store, tool sets) release their resources. - stop := context.AfterFunc(ctx, func() { - _ = e.Server.Close() - }) - defer stop() + return e, nil +} - if err := e.Server.Serve(ln); err != nil && ctx.Err() == nil { - slog.ErrorContext(ctx, "Failed to start server", "error", err) - return err +func corsMiddlewareConfig(spec string) (middleware.CORSConfig, error) { + origins, err := httpsec.ParseOrigins(spec) + if err != nil { + return middleware.CORSConfig{}, err } + cfg := middleware.CORSConfig{ + AllowOrigins: origins.Literals(), + AllowMethods: []string{http.MethodPost, http.MethodOptions}, + AllowHeaders: []string{"Authorization", "Content-Type", "Accept"}, + MaxAge: 86400, + } + if origins.HasPatterns() { + cfg.AllowOriginFunc = func(origin string) (bool, error) { return origins.MatchPattern(origin), nil } + } + return cfg, nil +} - return nil +func bearerAuthMiddleware(token string) echo.MiddlewareFunc { + auth := httpsec.BearerAuth(token) + return func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + if c.Request().Method == http.MethodOptions { + return next(c) + } + return echo.WrapMiddleware(auth)(next)(c) + } + } } diff --git a/pkg/a2a/server_invoke_test.go b/pkg/a2a/server_invoke_test.go new file mode 100644 index 0000000000..9e24e43a32 --- /dev/null +++ b/pkg/a2a/server_invoke_test.go @@ -0,0 +1,172 @@ +package a2a + +import ( + "context" + "fmt" + "net" + "sync" + "sync/atomic" + "testing" + + "github.com/a2aproject/a2a-go/a2a" + "github.com/a2aproject/a2a-go/a2aclient" + "github.com/stretchr/testify/require" + + dagent "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/modelsdev" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" +) + +type sequentialMockProvider struct { + id modelsdev.ID + streams []chat.MessageStream + + mu sync.Mutex + next int +} + +func (p *sequentialMockProvider) ID() modelsdev.ID { return p.id } + +func (p *sequentialMockProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + p.mu.Lock() + defer p.mu.Unlock() + + if p.next >= len(p.streams) { + return &mockStream{responses: []chat.MessageStreamResponse{{ + Choices: []chat.MessageStreamChoice{{Index: 0, FinishReason: chat.FinishReasonStop}}, + }}}, nil + } + stream := p.streams[p.next] + p.next++ + return stream, nil +} + +func (p *sequentialMockProvider) BaseConfig() base.Config { return base.Config{} } +func (p *sequentialMockProvider) MaxTokens() int { return 0 } + +func toolCallStream(name string) chat.MessageStream { + return &mockStream{responses: []chat.MessageStreamResponse{{ + Choices: []chat.MessageStreamChoice{{ + Index: 0, + Delta: chat.MessageDelta{ToolCalls: []tools.ToolCall{{ + ID: "tool-call", + Type: "function", + Function: tools.FunctionCall{Name: name, Arguments: "{}"}, + }}}, + }}, + }}} +} + +func stopStream() chat.MessageStream { + return &mockStream{responses: []chat.MessageStreamResponse{{ + Choices: []chat.MessageStreamChoice{{Index: 0, FinishReason: chat.FinishReasonStop}}, + }}} +} + +func TestServer_ToolPolicyOverInvoke(t *testing.T) { + t.Parallel() + + for _, tc := range []struct { + name string + policy session.SafetyPolicy + executed int32 + }{ + {name: "restricted", policy: session.SafetyPolicyRestricted}, + {name: "autonomous", policy: session.SafetyPolicyAutonomous, executed: 1}, + } { + t.Run(tc.name, func(t *testing.T) { + var executions atomic.Int32 + tool := tools.Tool{ + Name: "unsafe_tool", + Description: "test tool", + Parameters: map[string]any{}, + Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + executions.Add(1) + return tools.ResultSuccess("done"), nil + }, + } + provider := &sequentialMockProvider{ + id: modelsdev.NewID("test", "mock-model"), + streams: []chat.MessageStream{toolCallStream(tool.Name), stopStream()}, + } + root := dagent.New("root", "You are a test agent", dagent.WithModel(provider), dagent.WithTools(tool)) + server := startInvokeServer(t, team.New(team.WithAgents(root)), session.NewInMemorySessionStore(), servesafety.Resolved{Policy: tc.policy}) + + client, err := a2aclient.NewFromEndpoints(t.Context(), []a2a.AgentInterface{{ + Transport: a2a.TransportProtocolJSONRPC, + URL: fmt.Sprintf("http://%s/invoke", server.Addr()), + }}) + require.NoError(t, err) + _, err = client.SendMessage(t.Context(), &a2a.MessageSendParams{ + Message: a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "run the tool"}), + }) + require.NoError(t, err) + require.Equal(t, tc.executed, executions.Load()) + }) + } +} + +func TestServer_RejectsNonA2AContextCollision(t *testing.T) { + t.Parallel() + + var executions atomic.Int32 + tool := tools.Tool{ + Name: "collision_canary", + Description: "test tool", + Parameters: map[string]any{}, + Handler: func(context.Context, tools.ToolCall, tools.Runtime) (*tools.ToolCallResult, error) { + executions.Add(1) + return tools.ResultSuccess("done"), nil + }, + } + provider := &sequentialMockProvider{ + id: modelsdev.NewID("test", "mock-model"), + streams: []chat.MessageStream{toolCallStream(tool.Name), stopStream()}, + } + root := dagent.New("root", "You are a test agent", dagent.WithModel(provider), dagent.WithTools(tool)) + store := session.NewInMemorySessionStore() + existing := session.New(session.WithID("colliding-context"), session.WithOrigin("run"), session.WithTitle("private")) + require.NoError(t, store.AddSession(t.Context(), existing)) + server := startInvokeServer(t, team.New(team.WithAgents(root)), store, servesafety.Resolved{Policy: session.SafetyPolicyAutonomous}) + + client, err := a2aclient.NewFromEndpoints(t.Context(), []a2a.AgentInterface{{ + Transport: a2a.TransportProtocolJSONRPC, + URL: fmt.Sprintf("http://%s/invoke", server.Addr()), + }}) + require.NoError(t, err) + message := a2a.NewMessage(a2a.MessageRoleUser, a2a.TextPart{Text: "run the tool"}) + message.ContextID = "colliding-context" + got, err := client.SendMessage(t.Context(), &a2a.MessageSendParams{Message: message}) + require.NoError(t, err) + task, ok := got.(*a2a.Task) + require.True(t, ok) + require.Equal(t, a2a.TaskStateFailed, task.Status.State) + require.NotNil(t, task.Status.Message) + require.Len(t, task.Status.Message.Parts, 1) + failure, ok := task.Status.Message.Parts[0].(a2a.TextPart) + require.True(t, ok) + require.Equal(t, "agent run failed: context ID is not available", failure.Text) + require.Zero(t, executions.Load()) + + stored, err := store.GetSession(t.Context(), existing.ID) + require.NoError(t, err) + require.Equal(t, "run", stored.Origin) + require.Equal(t, "private", stored.Title) +} + +func startInvokeServer(t *testing.T, tm *team.Team, store session.Store, safety servesafety.Resolved) net.Listener { + t.Helper() + + ln, err := (&net.ListenConfig{}).Listen(t.Context(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + e, err := newServer(tm, "test.yaml", "root", store, safety, ln.Addr().String(), RunOptions{}) + require.NoError(t, err) + go func() { _ = e.Server.Serve(ln) }() + t.Cleanup(func() { require.NoError(t, e.Server.Close()) }) + return ln +} diff --git a/pkg/a2a/server_test.go b/pkg/a2a/server_test.go index eb2cd02e0e..6aa0b36221 100644 --- a/pkg/a2a/server_test.go +++ b/pkg/a2a/server_test.go @@ -4,6 +4,7 @@ import ( "context" "net" "net/http" + "net/http/httptest" "path/filepath" "testing" "time" @@ -12,7 +13,11 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/servesafety" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" ) func TestRoutableAddr(t *testing.T) { @@ -41,6 +46,18 @@ func TestRoutableAddr(t *testing.T) { } } +func TestRun_RejectsAutonomousYAMLSafety(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "DUMMY") + + var lc net.ListenConfig + ln, err := lc.Listen(t.Context(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + + err = Run(t.Context(), "testdata/autonomous.yaml", "root", filepath.Join(t.TempDir(), "session.db"), &config.RuntimeConfig{}, ln, RunOptions{}) + require.ErrorContains(t, err, "--safety autonomous") +} + // TestRun_StopsOnContextCancel: canceling the context must make Run stop // serving and return, releasing the session store's file handles (otherwise // t.TempDir cleanup fails on Windows with an open session.db). @@ -58,7 +75,7 @@ func TestRun_StopsOnContextCancel(t *testing.T) { done := make(chan error, 1) go func() { - done <- Run(ctx, "testdata/basic.yaml", "root", sessionDB, &config.RuntimeConfig{}, ln) + done <- Run(ctx, "testdata/basic.yaml", "root", sessionDB, &config.RuntimeConfig{}, ln, RunOptions{}) }() // Cancel only once the server actually serves, so the test exercises a @@ -86,3 +103,44 @@ func TestRun_StopsOnContextCancel(t *testing.T) { t.Fatal("Run did not return after context cancellation") } } + +func TestServerSecurity(t *testing.T) { + t.Parallel() + + tm := team.New(team.WithAgents(agent.New("root", "test"))) + store := session.NewInMemorySessionStore() + server, err := newServer(tm, "test.yaml", "root", store, servesafety.Resolved{}, "127.0.0.1:0", RunOptions{AuthToken: "secret", CORSOrigin: "https://app.example.com"}) + require.NoError(t, err) + + request := func(method, path string, headers map[string]string) *httptest.ResponseRecorder { + r := httptest.NewRequestWithContext(t.Context(), method, path, http.NoBody) + for key, value := range headers { + r.Header.Set(key, value) + } + w := httptest.NewRecorder() + server.ServeHTTP(w, r) + return w + } + + for _, path := range []string{a2asrv.WellKnownAgentCardPath, "/invoke"} { + response := request(http.MethodGet, path, nil) + require.Equal(t, http.StatusUnauthorized, response.Code) + require.Equal(t, "Bearer", response.Header().Get("WWW-Authenticate")) + } + response := request(http.MethodGet, a2asrv.WellKnownAgentCardPath, map[string]string{"Authorization": "Bearer secret"}) + require.Equal(t, http.StatusOK, response.Code) + + response = request(http.MethodOptions, "/invoke", map[string]string{ + "Origin": "https://app.example.com", "Access-Control-Request-Method": http.MethodPost, + "Access-Control-Request-Headers": "authorization,content-type", + }) + require.Equal(t, http.StatusNoContent, response.Code) + require.Equal(t, "https://app.example.com", response.Header().Get("Access-Control-Allow-Origin")) + require.Contains(t, response.Header().Get("Access-Control-Allow-Headers"), "Authorization") +} + +func TestCorsMiddlewareConfigRejectsInvalidOrigin(t *testing.T) { + t.Parallel() + _, err := corsMiddlewareConfig("not an origin") + require.Error(t, err) +} diff --git a/pkg/a2a/testdata/autonomous.yaml b/pkg/a2a/testdata/autonomous.yaml new file mode 100644 index 0000000000..25e2966160 --- /dev/null +++ b/pkg/a2a/testdata/autonomous.yaml @@ -0,0 +1,7 @@ +version: "15" + +agents: + root: + model: openai/gpt-3.5-turbo + safety: autonomous + instruction: Be helpful. diff --git a/pkg/chatserver/agent.go b/pkg/chatserver/agent.go index 4673bccf3a..1d33542c10 100644 --- a/pkg/chatserver/agent.go +++ b/pkg/chatserver/agent.go @@ -59,14 +59,13 @@ func (p agentPolicy) pick(model string) string { // assistant/tool turns are replayed verbatim so the agent sees the full // conversation, and the latest user message becomes the prompt. // -// Tool approval and non-interactive mode are forced on: this is a headless -// HTTP endpoint, there's no human in the loop to approve anything. +// Tool approval is determined by the resolved server policy because this +// endpoint has no interactive approval channel. // // Returns nil when the history contains no usable user message, in which // case the caller should reject the request. func buildSession(messages []ChatCompletionMessage) *session.Session { sess := session.New( - session.WithToolsApproved(true), session.WithNonInteractive(true), ) @@ -202,13 +201,11 @@ type agentEmit struct { // runAgentLoop drives the runtime to completion, forwarding events to // the supplied callbacks. // -// The session is built with ToolsApproved=true and NonInteractive=true, -// which means the runtime auto-approves tool calls and auto-stops on -// max-iterations. The handler cases below are intentionally kept as -// defence-in-depth: if those session settings ever drift, this handler -// still won't hang the request. Elicitation is the exception — the -// runtime always blocks until we respond, so its case is required for -// correctness, not just defence. +// The session is built with NonInteractive=true and a resolved safety policy. +// The handler cases below are intentionally kept as defence-in-depth: if +// those session settings ever drift, this handler still won't hang the +// request. Elicitation is the exception — the runtime always blocks until we +// respond, so its case is required for correctness, not just defence. // // All ErrorEvents seen in the run are joined into the returned error so // callers can see the full picture; the loop keeps draining until the @@ -233,7 +230,7 @@ func runAgentLoop(ctx context.Context, rt runtime.Runtime, sess *session.Session toolIndex++ } case *runtime.ToolCallConfirmationEvent: - // Defensive: should never fire while ToolsApproved=true. + // Defensive: resolved policy normally handles tool approval before this point. rt.Resume(ctx, runtime.ResumeApprove()) case *runtime.ElicitationRequestEvent: // Required: the runtime blocks until we respond, regardless diff --git a/pkg/chatserver/conversations_transaction_test.go b/pkg/chatserver/conversations_transaction_test.go index f55e2762ef..6974b1e874 100644 --- a/pkg/chatserver/conversations_transaction_test.go +++ b/pkg/chatserver/conversations_transaction_test.go @@ -23,6 +23,31 @@ func newConvServer(t *testing.T) *server { } } +func TestResolveSession_AppliesSafetyCeiling(t *testing.T) { + t.Parallel() + s := newConvServer(t) + s.safety = session.SafetyPolicyRestricted + + seed := session.New(session.WithSafetyPolicy(session.SafetyPolicyAutonomous)) + seed.AddMessage(session.UserMessage("first")) + s.conversations.Put("conv-1", seed) + + working, err := s.resolveSession("conv-1", []ChatCompletionMessage{{Role: "user", Content: "second"}}) + require.NoError(t, err) + assert.Equal(t, session.SafetyPolicyRestricted, working.GetSafetyPolicy()) + assert.Equal(t, session.SafetyPolicyAutonomous, seed.GetSafetyPolicy(), "the cached session remains untouched until commit") +} + +func TestResolveSession_AppliesSafetyToNewSession(t *testing.T) { + t.Parallel() + s := newConvServer(t) + s.safety = session.SafetyPolicyRestricted + + working, err := s.resolveSession("", []ChatCompletionMessage{{Role: "user", Content: "first"}}) + require.NoError(t, err) + assert.Equal(t, session.SafetyPolicyRestricted, working.GetSafetyPolicy()) +} + // TestResolveSession_WorksOnClone verifies that continuing a cached // conversation mutates a copy, leaving the cached session untouched until // the caller commits. diff --git a/pkg/chatserver/openapi.json b/pkg/chatserver/openapi.json index 2856d5dc25..8702554f61 100644 --- a/pkg/chatserver/openapi.json +++ b/pkg/chatserver/openapi.json @@ -17,7 +17,7 @@ "bearerAuth": { "type": "http", "scheme": "bearer", - "description": "Static token configured via --api-key. When --api-key is not set the server is unauthenticated." + "description": "Static token configured via --api-key or --api-key-env. Non-loopback listeners require authentication unless --insecure-no-auth is explicitly set." } }, "schemas": { @@ -217,7 +217,7 @@ } }, "401": { - "description": "Missing or invalid bearer token (only when --api-key is set).", + "description": "Missing or invalid bearer token (when authentication is configured).", "content": { "application/json": { "schema": { "$ref": "#/components/schemas/ErrorResponse" } diff --git a/pkg/chatserver/server.go b/pkg/chatserver/server.go index c466108f54..6482dc002e 100644 --- a/pkg/chatserver/server.go +++ b/pkg/chatserver/server.go @@ -25,8 +25,6 @@ import ( "math" "net" "net/http" - "net/url" - "regexp" "slices" "strconv" "strings" @@ -40,7 +38,9 @@ import ( "github.com/docker/docker-agent/pkg/config" "github.com/docker/docker-agent/pkg/echolog" + "github.com/docker/docker-agent/pkg/httpsec" "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/servesafety" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/teamloader" @@ -56,6 +56,10 @@ type Options struct { AgentName string // RunConfig is the runtime configuration used to load the team. RunConfig *config.RuntimeConfig + // CLISafety selects the server safety policy before agent and runtime YAML. + CLISafety session.SafetyPolicy + // OnSafetyPolicy receives the resolved policy after the team loads. + OnSafetyPolicy func(servesafety.Resolved) // CORSOrigin is the allowed value for the Access-Control-Allow-Origin // header. When empty, the CORS middleware is not registered at all // (the server never emits any Access-Control-* response header). @@ -127,6 +131,17 @@ func Run(ctx context.Context, agentFilename string, opts Options, ln net.Listene if err != nil { return err } + selectedAgent, err := t.AgentOrDefault(opts.AgentName) + if err != nil { + return fmt.Errorf("failed to get agent: %w", err) + } + resolvedSafety, err := servesafety.Resolve(opts.CLISafety, string(selectedAgent.Safety()), string(t.RuntimeSafety())) + if err != nil { + return fmt.Errorf("resolve serve safety policy: %w", err) + } + if opts.OnSafetyPolicy != nil { + opts.OnSafetyPolicy(resolvedSafety) + } // Wrap with otelhttp so incoming /v1/chat/completions requests // (including SSE streams) extract the caller's trace context. @@ -137,6 +152,7 @@ func Run(ctx context.Context, agentFilename string, opts Options, ln net.Listene newRouter(&server{ team: t, policy: policy, + safety: resolvedSafety.Policy, conversations: newConversationStore(opts.ConversationsMaxSessions, conversationTTL(opts)), conversationLocks: newConversationLockSet(), runtimes: newRuntimePool(ctx, t, opts.MaxIdleRuntimes), @@ -197,6 +213,7 @@ func serve(ctx context.Context, httpServer *http.Server, ln net.Listener) error type server struct { team *team.Team policy agentPolicy + safety session.SafetyPolicy conversations *conversationStore conversationLocks *conversationLockSet runtimes *runtimePool @@ -268,73 +285,25 @@ func requestTimeoutMiddleware(d time.Duration) echo.MiddlewareFunc { // Returns an error when no entry parses successfully, in which case the // caller leaves the middleware unregistered. func corsMiddlewareConfig(spec string) (middleware.CORSConfig, error) { - var literals []string - var patterns []*regexp.Regexp - for raw := range strings.SplitSeq(spec, ",") { - entry := strings.TrimSpace(raw) - if entry == "" { - continue - } - if rest, ok := strings.CutPrefix(entry, "~"); ok { - re, err := regexp.Compile(rest) - if err != nil { - return middleware.CORSConfig{}, fmt.Errorf("invalid CORS regex %q: %w", rest, err) - } - patterns = append(patterns, re) - continue - } - if err := validateCORSOrigin(entry); err != nil { - return middleware.CORSConfig{}, err - } - literals = append(literals, entry) - } - if len(literals) == 0 && len(patterns) == 0 { - return middleware.CORSConfig{}, errors.New("no usable CORS origins") + origins, err := httpsec.ParseOrigins(spec) + if err != nil { + return middleware.CORSConfig{}, err } cfg := middleware.CORSConfig{ - AllowOrigins: literals, + AllowOrigins: origins.Literals(), AllowMethods: []string{http.MethodGet, http.MethodPost, http.MethodOptions}, AllowHeaders: []string{"Authorization", "Content-Type", "Accept"}, MaxAge: 86400, } - if len(patterns) > 0 { + if origins.HasPatterns() { cfg.AllowOriginFunc = func(origin string) (bool, error) { - for _, re := range patterns { - if re.MatchString(origin) { - return true, nil - } - } - return false, nil + return origins.MatchPattern(origin), nil } } return cfg, nil } -// validateCORSOrigin sanity-checks a literal origin entry. The aim is to -// reject obvious typos early ("http//foo.com", "https://foo.com/bar") -// rather than to be a full URL parser — the echo middleware will still -// do its own matching at request time. -func validateCORSOrigin(o string) error { - if o == "*" { - return nil - } - u, err := url.Parse(o) - if err != nil { - return fmt.Errorf("invalid CORS origin %q: %w", o, err) - } - if u.Scheme != "http" && u.Scheme != "https" { - return fmt.Errorf("invalid CORS origin %q: scheme must be http or https", o) - } - if u.Host == "" { - return fmt.Errorf("invalid CORS origin %q: missing host", o) - } - if u.Path != "" || u.RawQuery != "" || u.Fragment != "" { - return fmt.Errorf("invalid CORS origin %q: must not include path, query, or fragment", o) - } - return nil -} - // bearerAuthMiddleware enforces the static `Authorization: Bearer ` // header. CORS preflight requests (OPTIONS) are exempted so that browsers // can negotiate before sending the auth header. @@ -440,6 +409,7 @@ func (s *server) resolveSession(id string, msgs []ChatCompletionMessage) (*sessi if !appendLatestUser(working, msgs) { return nil, errors.New("no user message provided") } + working.SetSafetyPolicy(servesafety.ResumeCeiling(working.GetSafetyPolicy(), s.safety)) return working, nil } } @@ -447,6 +417,7 @@ func (s *server) resolveSession(id string, msgs []ChatCompletionMessage) (*sessi if sess == nil { return nil, errors.New("no user message provided") } + sess.SetSafetyPolicy(s.safety) return sess, nil } diff --git a/pkg/chatserver/server_test.go b/pkg/chatserver/server_test.go index 084a45f513..13f347b3f6 100644 --- a/pkg/chatserver/server_test.go +++ b/pkg/chatserver/server_test.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "errors" + "net" "net/http" "net/http/httptest" "strings" @@ -15,8 +16,21 @@ import ( "github.com/stretchr/testify/require" "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/config" ) +func TestRun_RejectsAutonomousYAMLSafety(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "DUMMY") + + var lc net.ListenConfig + ln, err := lc.Listen(t.Context(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + + err = Run(t.Context(), "testdata/autonomous.yaml", Options{RunConfig: &config.RuntimeConfig{}}, ln) + require.ErrorContains(t, err, "--safety autonomous") +} + func TestBuildSession_RequiresUserMessage(t *testing.T) { t.Parallel() tests := []struct { @@ -58,7 +72,6 @@ func TestBuildSession_RequiresUserMessage(t *testing.T) { return } require.NotNil(t, sess) - assert.True(t, sess.ToolsApproved) assert.True(t, sess.NonInteractive) }) } diff --git a/pkg/chatserver/testdata/autonomous.yaml b/pkg/chatserver/testdata/autonomous.yaml new file mode 100644 index 0000000000..25e2966160 --- /dev/null +++ b/pkg/chatserver/testdata/autonomous.yaml @@ -0,0 +1,7 @@ +version: "15" + +agents: + root: + model: openai/gpt-3.5-turbo + safety: autonomous + instruction: Be helpful. diff --git a/pkg/httpsec/httpsec.go b/pkg/httpsec/httpsec.go new file mode 100644 index 0000000000..4204bee16b --- /dev/null +++ b/pkg/httpsec/httpsec.go @@ -0,0 +1,103 @@ +// Package httpsec provides HTTP security primitives for serving commands. +package httpsec + +import ( + "crypto/subtle" + "errors" + "fmt" + "net/http" + "net/url" + "regexp" + "slices" + "strings" +) + +// BearerAuth authenticates requests with a static bearer token. +func BearerAuth(token string) func(http.Handler) http.Handler { + expected := []byte(token) + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + got, ok := strings.CutPrefix(r.Header.Get("Authorization"), "Bearer ") + if !ok || subtle.ConstantTimeCompare([]byte(got), expected) != 1 { + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + return + } + next.ServeHTTP(w, r) + }) + } +} + +// OriginMatcher holds validated literal origins and compiled origin patterns. +type OriginMatcher struct { + literals []string + patterns []*regexp.Regexp +} + +// ParseOrigins parses a comma-separated list of literal origins and regular +// expression patterns prefixed with "~". +func ParseOrigins(spec string) (*OriginMatcher, error) { + matcher := &OriginMatcher{} + for raw := range strings.SplitSeq(spec, ",") { + entry := strings.TrimSpace(raw) + if entry == "" { + continue + } + if pattern, ok := strings.CutPrefix(entry, "~"); ok { + re, err := regexp.Compile(pattern) + if err != nil { + return nil, fmt.Errorf("invalid CORS regex %q: %w", pattern, err) + } + matcher.patterns = append(matcher.patterns, re) + continue + } + if err := validateOrigin(entry); err != nil { + return nil, err + } + matcher.literals = append(matcher.literals, entry) + } + if len(matcher.literals) == 0 && len(matcher.patterns) == 0 { + return nil, errors.New("no usable CORS origins") + } + return matcher, nil +} + +// Literals returns the configured literal origins. +func (m *OriginMatcher) Literals() []string { + return slices.Clone(m.literals) +} + +// HasPatterns reports whether the matcher has regular expression patterns. +func (m *OriginMatcher) HasPatterns() bool { + return len(m.patterns) > 0 +} + +// MatchPattern reports whether origin matches a configured regular expression. +func (m *OriginMatcher) MatchPattern(origin string) bool { + for _, pattern := range m.patterns { + if pattern.MatchString(origin) { + return true + } + } + return false +} + +func validateOrigin(origin string) error { + if origin == "*" { + return nil + } + u, err := url.Parse(origin) + if err != nil { + return fmt.Errorf("invalid CORS origin %q: %w", origin, err) + } + if u.Scheme != "http" && u.Scheme != "https" { + return fmt.Errorf("invalid CORS origin %q: scheme must be http or https", origin) + } + if u.Host == "" { + return fmt.Errorf("invalid CORS origin %q: missing host", origin) + } + if u.Path != "" || u.RawQuery != "" || u.Fragment != "" { + return fmt.Errorf("invalid CORS origin %q: must not include path, query, or fragment", origin) + } + return nil +} diff --git a/pkg/httpsec/httpsec_test.go b/pkg/httpsec/httpsec_test.go new file mode 100644 index 0000000000..4e284807e1 --- /dev/null +++ b/pkg/httpsec/httpsec_test.go @@ -0,0 +1,107 @@ +package httpsec + +import ( + "net/http" + "net/http/httptest" + "os/exec" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestBearerAuth(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + authorization string + wantStatus int + wantCalled bool + }{ + {name: "missing", wantStatus: http.StatusUnauthorized}, + {name: "wrong scheme", authorization: "Basic secret", wantStatus: http.StatusUnauthorized}, + {name: "wrong token", authorization: "Bearer wrong", wantStatus: http.StatusUnauthorized}, + {name: "empty token", authorization: "Bearer ", wantStatus: http.StatusUnauthorized}, + {name: "matching token", authorization: "Bearer secret", wantStatus: http.StatusNoContent, wantCalled: true}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + called := false + handler := BearerAuth("secret")(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + called = true + w.WriteHeader(http.StatusNoContent) + })) + req := httptest.NewRequestWithContext(t.Context(), http.MethodGet, "/", http.NoBody) + req.Header.Set("Authorization", tc.authorization) + rec := httptest.NewRecorder() + + handler.ServeHTTP(rec, req) + + assert.Equal(t, tc.wantStatus, rec.Code) + assert.Equal(t, tc.wantCalled, called) + if !tc.wantCalled { + assert.Equal(t, "Bearer", rec.Header().Get("WWW-Authenticate")) + } + }) + } +} + +func TestParseOrigins(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + spec string + wantLiterals []string + wantPatterns bool + matches map[string]bool + wantErr string + }{ + {name: "literal", spec: "https://example.com", wantLiterals: []string{"https://example.com"}}, + {name: "literals with whitespace", spec: " https://example.com, http://localhost:3000 ", wantLiterals: []string{"https://example.com", "http://localhost:3000"}}, + {name: "wildcard", spec: "*", wantLiterals: []string{"*"}}, + {name: "pattern", spec: "~^https://[a-z]+\\.example\\.com$", wantPatterns: true, matches: map[string]bool{"https://app.example.com": true, "https://example.com": false}}, + {name: "literal and pattern", spec: "https://example.com,~^https://[a-z]+\\.example\\.com$", wantLiterals: []string{"https://example.com"}, wantPatterns: true}, + {name: "blank entries", spec: ",, https://example.com, ,", wantLiterals: []string{"https://example.com"}}, + {name: "empty", spec: " , ", wantErr: "no usable CORS origins"}, + {name: "invalid pattern", spec: "~[", wantErr: "invalid CORS regex"}, + {name: "missing scheme", spec: "example.com", wantErr: "scheme must be http or https"}, + {name: "unsupported scheme", spec: "ftp://example.com", wantErr: "scheme must be http or https"}, + {name: "missing host", spec: "https:", wantErr: "missing host"}, + {name: "path", spec: "https://example.com/path", wantErr: "must not include path"}, + {name: "query", spec: "https://example.com?query=value", wantErr: "must not include path"}, + {name: "fragment", spec: "https://example.com#fragment", wantErr: "must not include path"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + matcher, err := ParseOrigins(tc.spec) + if tc.wantErr != "" { + require.Error(t, err) + assert.Contains(t, err.Error(), tc.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tc.wantLiterals, matcher.Literals()) + assert.Equal(t, tc.wantPatterns, matcher.HasPatterns()) + for origin, want := range tc.matches { + assert.Equal(t, want, matcher.MatchPattern(origin)) + } + }) + } +} + +func TestPackageDoesNotDependOnEcho(t *testing.T) { + cmd := exec.CommandContext(t.Context(), "go", "list", "-deps", "-test", ".") + output, err := cmd.Output() + require.NoError(t, err) + + for dependency := range strings.FieldsSeq(string(output)) { + assert.Falsef(t, dependency == "github.com/labstack/echo" || strings.HasPrefix(dependency, "github.com/labstack/echo/"), "httpsec depends on Echo package %s", dependency) + } +} diff --git a/pkg/mcp/server.go b/pkg/mcp/server.go index ca07389df2..7d0a548abd 100644 --- a/pkg/mcp/server.go +++ b/pkg/mcp/server.go @@ -17,7 +17,9 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/httpsec" "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/servesafety" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/teamloader" @@ -35,6 +37,12 @@ type ToolOutput struct { Response string `json:"response" jsonschema:"the response from the agent"` } +type HTTPOptions struct { + CLISafety session.SafetyPolicy + AuthToken string + OnSafetyPolicy func(servesafety.Resolved) +} + func StartMCPServer(ctx context.Context, agentFilename, agentName string, runConfig *config.RuntimeConfig) error { slog.DebugContext(ctx, "Starting MCP server", "agent", agentFilename) @@ -54,28 +62,55 @@ func StartMCPServer(ctx context.Context, agentFilename, agentName string, runCon } // StartHTTPServer starts a streaming HTTP MCP server on the given listener -func StartHTTPServer(ctx context.Context, agentFilename, agentName string, runConfig *config.RuntimeConfig, ln net.Listener) error { +func StartHTTPServer(ctx context.Context, agentFilename, agentName string, runConfig *config.RuntimeConfig, ln net.Listener, options HTTPOptions) error { slog.DebugContext(ctx, "Starting HTTP MCP server", "agent", agentFilename, "addr", ln.Addr()) - server, cleanup, err := createMCPServer(ctx, agentFilename, agentName, runConfig) + agentSource, err := config.Resolve(agentFilename, nil) + if err != nil { + return err + } + t, err := teamloader.Load(ctx, agentSource, runConfig, loaderdefaults.Opts()...) + if err != nil { + return fmt.Errorf("failed to load agents: %w", err) + } + defer func() { + if err := t.StopToolSets(ctx); err != nil { + slog.ErrorContext(ctx, "Failed to stop tool sets", "error", err) + } + }() + + selectedAgent, err := t.AgentOrDefault(agentName) + if err != nil { + return fmt.Errorf("failed to get agent: %w", err) + } + resolvedSafety, err := servesafety.Resolve(options.CLISafety, string(selectedAgent.Safety()), string(t.RuntimeSafety())) + if err != nil { + return fmt.Errorf("resolve serve safety policy: %w", err) + } + if options.OnSafetyPolicy != nil { + options.OnSafetyPolicy(resolvedSafety) + } + + server, err := createMCPServerForTeam(ctx, t, agentFilename, agentName, runConfig, resolvedSafety.Policy) if err != nil { return err } - defer cleanup() fmt.Printf("MCP HTTP server listening on http://%s\n", ln.Addr()) + handler := http.Handler(mcp.NewStreamableHTTPHandler(func(_ *http.Request) *mcp.Server { + return server + }, nil)) + if options.AuthToken != "" { + handler = httpsec.BearerAuth(options.AuthToken)(handler) + } + // Wrap with otelhttp so the MCP-over-HTTP transport extracts // `traceparent` / `baggage` from incoming requests just like the // stdio transport extracts them from `params._meta`. Without this // HTTP-mode MCP clients lose trace context at the boundary. httpServer := &http.Server{ - Handler: otelhttp.NewHandler( - mcp.NewStreamableHTTPHandler(func(_ *http.Request) *mcp.Server { - return server - }, nil), - "mcp.http", - ), + Handler: otelhttp.NewHandler(handler, "mcp.http"), ReadHeaderTimeout: 10 * time.Second, } @@ -117,6 +152,15 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo } } + server, err := createMCPServerForTeam(ctx, t, agentFilename, agentName, runConfig, session.SafetyPolicyAutonomous) + if err != nil { + cleanup() + return nil, nil, err + } + return server, cleanup, nil +} + +func createMCPServerForTeam(ctx context.Context, t *team.Team, agentFilename, agentName string, runConfig *config.RuntimeConfig, safety session.SafetyPolicy) (*mcp.Server, error) { // The SDK only starts keep-alive when KeepAlive > 0. server := mcp.NewServer(&mcp.Implementation{ Name: "docker agent", @@ -128,15 +172,13 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo agentNames := t.AgentNames() if agentName != "" { if !slices.Contains(agentNames, agentName) { - cleanup() - return nil, nil, fmt.Errorf("agent %s not found in %s", agentName, agentFilename) + return nil, fmt.Errorf("agent %s not found in %s", agentName, agentFilename) } agentNames = []string{agentName} } if runConfig.MCPToolName != "" && len(agentNames) > 1 { - cleanup() - return nil, nil, errors.New("--tool-name can only be used when exactly one agent is exposed") + return nil, errors.New("--tool-name can only be used when exactly one agent is exposed") } slog.DebugContext(ctx, "Adding MCP tools for agents", "count", len(agentNames)) @@ -144,8 +186,7 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo for _, agentName := range agentNames { ag, err := t.Agent(agentName) if err != nil { - cleanup() - return nil, nil, fmt.Errorf("failed to get agent %s: %w", agentName, err) + return nil, fmt.Errorf("failed to get agent %s: %w", agentName, err) } description := cmp.Or(ag.Description(), fmt.Sprintf("Run the %s agent", agentName)) @@ -154,8 +195,7 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo annotations, err := agentToolAnnotations(ctx, ag) if err != nil { - cleanup() - return nil, nil, fmt.Errorf("failed to compute annotations for agent %s: %w", agentName, err) + return nil, fmt.Errorf("failed to compute annotations for agent %s: %w", agentName, err) } annotations.Title = description @@ -168,13 +208,17 @@ func createMCPServer(ctx context.Context, agentFilename, agentName string, runCo OutputSchema: tools.MustSchemaFor[ToolOutput](), } - mcp.AddTool(server, toolDef, CreateToolHandler(t, agentName)) + mcp.AddTool(server, toolDef, createToolHandler(t, agentName, safety)) } - return server, cleanup, nil + return server, nil } -func CreateToolHandler(t *team.Team, agentName string) func(context.Context, *mcp.CallToolRequest, ToolInput) (*mcp.CallToolResult, ToolOutput, error) { +func CreateToolHandler(t *team.Team, agentName string, safety session.SafetyPolicy) func(context.Context, *mcp.CallToolRequest, ToolInput) (*mcp.CallToolResult, ToolOutput, error) { + return createToolHandler(t, agentName, safety) +} + +func createToolHandler(t *team.Team, agentName string, safety session.SafetyPolicy) func(context.Context, *mcp.CallToolRequest, ToolInput) (*mcp.CallToolResult, ToolOutput, error) { return func(ctx context.Context, req *mcp.CallToolRequest, input ToolInput) (result *mcp.CallToolResult, output ToolOutput, err error) { // Extract W3C trace context from `params._meta` (per the OTel // MCP semconv) so the SERVER span chains onto the calling @@ -201,16 +245,17 @@ func CreateToolHandler(t *team.Team, agentName string) func(context.Context, *mc return nil, ToolOutput{}, fmt.Errorf("failed to get agent: %w", err) } - sess := session.New( + sessionOptions := []session.Opt{ session.WithTitle("MCP tool call"), session.WithMaxIterations(ag.MaxIterations()), session.WithMaxConsecutiveToolCalls(ag.MaxConsecutiveToolCalls()), session.WithMaxOldToolCallTokens(ag.MaxOldToolCallTokens()), session.WithMaxToolResultTokens(ag.MaxToolResultTokens()), session.WithUserMessage(input.Message), - session.WithToolsApproved(true), session.WithNonInteractive(true), - ) + session.WithSafetyPolicy(safety), + } + sess := session.New(sessionOptions...) rt, err := runtime.New(ctx, t, runtime.WithCurrentAgent(agentName), diff --git a/pkg/mcp/server_test.go b/pkg/mcp/server_test.go index 0cfb581eab..543ce42efe 100644 --- a/pkg/mcp/server_test.go +++ b/pkg/mcp/server_test.go @@ -1,6 +1,9 @@ package mcp import ( + "net" + "net/http" + "net/http/httptest" "testing" "github.com/stretchr/testify/assert" @@ -8,6 +11,7 @@ import ( "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/config" + "github.com/docker/docker-agent/pkg/httpsec" "github.com/docker/docker-agent/pkg/tools" ) @@ -21,6 +25,52 @@ func annot(readOnly, idempotent bool, destructive, openWorld *bool) tools.ToolAn } } +func TestStartHTTPServer_RejectsAutonomousYAMLSafety(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "DUMMY") + + var lc net.ListenConfig + ln, err := lc.Listen(t.Context(), "tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + + err = StartHTTPServer(t.Context(), "testdata/autonomous.yaml", "root", &config.RuntimeConfig{}, ln, HTTPOptions{}) + require.ErrorContains(t, err, "--safety autonomous") +} + +func TestCreateMCPServer_AcceptsAutonomousYAMLSafetyForStdio(t *testing.T) { + t.Setenv("OPENAI_API_KEY", "DUMMY") + + server, cleanup, err := createMCPServer(t.Context(), "testdata/autonomous.yaml", "root", &config.RuntimeConfig{}) + require.NoError(t, err) + require.NotNil(t, server) + cleanup() +} + +func TestHTTPBearerAuth(t *testing.T) { + t.Parallel() + + handler := httpsec.BearerAuth("secret")(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + for _, tc := range []struct { + name string + header string + want int + }{ + {name: "missing", want: http.StatusUnauthorized}, + {name: "wrong", header: "Bearer wrong", want: http.StatusUnauthorized}, + {name: "correct", header: "Bearer secret", want: http.StatusNoContent}, + } { + t.Run(tc.name, func(t *testing.T) { + req := httptest.NewRequestWithContext(t.Context(), http.MethodPost, "/mcp", http.NoBody) + req.Header.Set("Authorization", tc.header) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + assert.Equal(t, tc.want, rec.Code) + }) + } +} + func TestAgentToolAnnotations(t *testing.T) { t.Parallel() diff --git a/pkg/mcp/testdata/autonomous.yaml b/pkg/mcp/testdata/autonomous.yaml new file mode 100644 index 0000000000..25e2966160 --- /dev/null +++ b/pkg/mcp/testdata/autonomous.yaml @@ -0,0 +1,7 @@ +version: "15" + +agents: + root: + model: openai/gpt-3.5-turbo + safety: autonomous + instruction: Be helpful. diff --git a/pkg/runtime/remote_runtime.go b/pkg/runtime/remote_runtime.go index df871a8e42..bc515982a5 100644 --- a/pkg/runtime/remote_runtime.go +++ b/pkg/runtime/remote_runtime.go @@ -818,6 +818,10 @@ func (s *RemoteSessionStore) GetSession(context.Context, string) (*session.Sessi return nil, fmt.Errorf("get session: %w", ErrUnsupported) } +func (s *RemoteSessionStore) GetSessionByOrigin(context.Context, string, string) (*session.Session, error) { + return nil, fmt.Errorf("get session by origin: %w", ErrUnsupported) +} + func (s *RemoteSessionStore) GetSessions(ctx context.Context) ([]*session.Session, error) { sessions, err := s.client.GetAllSessions(ctx) if err != nil { diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 1b77dab610..ec9ca595c2 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -2011,6 +2011,7 @@ func (sm *SessionManager) SetSessionAgentModel(ctx context.Context, sessionID, m ID: sess.ID, Title: title, CreatedAt: sess.CreatedAt, + Origin: sess.Origin, WorkingDir: sess.WorkingDir, // SafetyPolicy must travel with ToolsApproved: omitting it would // reset a strict/balanced session to the legacy default on reload. diff --git a/pkg/servesafety/safety.go b/pkg/servesafety/safety.go new file mode 100644 index 0000000000..1de8ab48a5 --- /dev/null +++ b/pkg/servesafety/safety.go @@ -0,0 +1,70 @@ +package servesafety + +import ( + "errors" + "fmt" + + "github.com/docker/docker-agent/pkg/session" +) + +type Source string + +const ( + SourceCLI Source = "command line" + SourceAgentYAML Source = "agent configuration" + SourceRuntimeYAML Source = "runtime configuration" + SourceDefault Source = "serve default" +) + +type Resolved struct { + Policy session.SafetyPolicy + Source Source +} + +// Resolve selects a serve safety policy without consulting local user settings. +func Resolve(cli session.SafetyPolicy, agentYAML, runtimeYAML string) (Resolved, error) { + if cli != "" { + policy := cli.Normalize() + if policy != session.SafetyPolicyAutonomous && policy != session.SafetyPolicyStrict && policy != session.SafetyPolicyBalanced && policy != session.SafetyPolicyRestricted { + return Resolved{}, fmt.Errorf("invalid safety value %q (valid: strict, balanced, restricted, autonomous)", cli) + } + return Resolved{Policy: policy, Source: SourceCLI}, nil + } + if agentYAML != "" { + policy, err := yamlPolicy(agentYAML) + if err != nil { + return Resolved{}, fmt.Errorf("agent safety: %w", err) + } + return Resolved{Policy: policy, Source: SourceAgentYAML}, nil + } + if runtimeYAML != "" { + policy, err := yamlPolicy(runtimeYAML) + if err != nil { + return Resolved{}, fmt.Errorf("runtime safety: %w", err) + } + return Resolved{Policy: policy, Source: SourceRuntimeYAML}, nil + } + return Resolved{Policy: session.SafetyPolicyRestricted, Source: SourceDefault}, nil +} + +func yamlPolicy(value string) (session.SafetyPolicy, error) { + policy := session.SafetyPolicy(value) + switch policy { + case session.SafetyPolicyStrict, session.SafetyPolicyBalanced, session.SafetyPolicyRestricted: + return policy, nil + case session.SafetyPolicyAutonomous: + return "", errors.New("autonomous safety in configuration is not supported; use --safety autonomous to opt in") + default: + return "", fmt.Errorf("invalid safety value %q (valid: strict, balanced, restricted)", value) + } +} + +// ResumeCeiling caps an existing session at the server's resolved policy. +// An empty persisted policy is a legacy unset state and adopts the server policy. +func ResumeCeiling(persisted, serve session.SafetyPolicy) session.SafetyPolicy { + persisted = persisted.Normalize() + if persisted == "" { + return serve.Normalize() + } + return session.MinSafetyPolicy(persisted, serve) +} diff --git a/pkg/servesafety/safety_test.go b/pkg/servesafety/safety_test.go new file mode 100644 index 0000000000..60a46caccb --- /dev/null +++ b/pkg/servesafety/safety_test.go @@ -0,0 +1,77 @@ +package servesafety + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/session" +) + +func TestResolve(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + cli session.SafetyPolicy + agentYAML, runtimeYAML string + want Resolved + }{ + {"default", "", "", "", Resolved{session.SafetyPolicyRestricted, SourceDefault}}, + {"runtime", "", "", "balanced", Resolved{session.SafetyPolicyBalanced, SourceRuntimeYAML}}, + {"agent over runtime", "", "strict", "balanced", Resolved{session.SafetyPolicyStrict, SourceAgentYAML}}, + {"CLI over YAML", session.SafetyPolicyAutonomous, "strict", "balanced", Resolved{session.SafetyPolicyAutonomous, SourceCLI}}, + {"CLI aliases normalize", "unsafe", "strict", "balanced", Resolved{session.SafetyPolicyAutonomous, SourceCLI}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := Resolve(test.cli, test.agentYAML, test.runtimeYAML) + require.NoError(t, err) + assert.Equal(t, test.want, got) + }) + } +} + +func TestResolveRejectsInvalidYAML(t *testing.T) { + t.Parallel() + + for _, test := range []struct { + name string + agent string + runtime string + message string + }{ + {"agent autonomous", "autonomous", "", "--safety autonomous"}, + {"runtime autonomous", "", "autonomous", "--safety autonomous"}, + {"invalid agent", "unknown", "", `invalid safety value "unknown"`}, + {"invalid runtime", "", "unknown", `invalid safety value "unknown"`}, + } { + t.Run(test.name, func(t *testing.T) { + _, err := Resolve("", test.agent, test.runtime) + require.ErrorContains(t, err, test.message) + }) + } +} + +func TestResumeCeiling(t *testing.T) { + t.Parallel() + + policies := []session.SafetyPolicy{ + session.SafetyPolicyStrict, + session.SafetyPolicyBalanced, + session.SafetyPolicyRestricted, + session.SafetyPolicyAutonomous, + } + for _, persisted := range policies { + for _, serve := range policies { + t.Run(string(persisted)+"/"+string(serve), func(t *testing.T) { + assert.Equal(t, session.MinSafetyPolicy(persisted, serve), ResumeCeiling(persisted, serve)) + }) + } + } + + assert.Equal(t, session.SafetyPolicyRestricted, ResumeCeiling("", session.SafetyPolicyRestricted)) + assert.Equal(t, session.SafetyPolicyAutonomous, ResumeCeiling("unsafe", session.SafetyPolicyAutonomous)) +} diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 8534d1c4a8..63c20d2b94 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -435,6 +435,15 @@ func getAllMigrations() []Migration { UpSQL: `ALTER TABLE sessions ADD COLUMN attributes TEXT DEFAULT '{}'`, DownSQL: `ALTER TABLE sessions DROP COLUMN attributes`, }, + // Revert PRs must retain this entry: upgraded databases keep the column, + // and removing it makes older binaries reject them as newer databases. + { + ID: 27, + Name: "027_add_session_origin_column", + Description: "Record the protocol surface that created each session", + UpSQL: `ALTER TABLE sessions ADD COLUMN origin TEXT NOT NULL DEFAULT 'run'`, + DownSQL: `ALTER TABLE sessions DROP COLUMN origin`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index e4dc5a45ad..5422ffdfed 100644 --- a/pkg/session/migrations_pinned_test.go +++ b/pkg/session/migrations_pinned_test.go @@ -39,7 +39,7 @@ func TestMigrationCatalogIsContentPinned(t *testing.T) { got := digestMigrationCatalog(getAllMigrations()) - const wantDigest = "71e0d68cf3a8439361a8339bad4ac543ffe3fda02fcce480ca0adc4e0f4017ce" + const wantDigest = "73643834fd1cc3b0dfd2a2ba52593c0b79ba773d364045a8e3b99ba890b55476" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/session/safety_policy_test.go b/pkg/session/safety_policy_test.go index 21be3456ec..95e8d5cb7b 100644 --- a/pkg/session/safety_policy_test.go +++ b/pkg/session/safety_policy_test.go @@ -47,6 +47,44 @@ func TestSafetyPolicy_Normalize(t *testing.T) { } } +func TestMinSafetyPolicy(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + a SafetyPolicy + b SafetyPolicy + want SafetyPolicy + }{ + {"strict and strict", SafetyPolicyStrict, SafetyPolicyStrict, SafetyPolicyStrict}, + {"strict and balanced", SafetyPolicyStrict, SafetyPolicyBalanced, SafetyPolicyStrict}, + {"strict and restricted", SafetyPolicyStrict, SafetyPolicyRestricted, SafetyPolicyStrict}, + {"strict and autonomous", SafetyPolicyStrict, SafetyPolicyAutonomous, SafetyPolicyStrict}, + {"balanced and strict", SafetyPolicyBalanced, SafetyPolicyStrict, SafetyPolicyStrict}, + {"balanced and balanced", SafetyPolicyBalanced, SafetyPolicyBalanced, SafetyPolicyBalanced}, + {"balanced and restricted", SafetyPolicyBalanced, SafetyPolicyRestricted, SafetyPolicyBalanced}, + {"balanced and autonomous", SafetyPolicyBalanced, SafetyPolicyAutonomous, SafetyPolicyBalanced}, + {"restricted and strict", SafetyPolicyRestricted, SafetyPolicyStrict, SafetyPolicyStrict}, + {"restricted and balanced", SafetyPolicyRestricted, SafetyPolicyBalanced, SafetyPolicyBalanced}, + {"restricted and restricted", SafetyPolicyRestricted, SafetyPolicyRestricted, SafetyPolicyRestricted}, + {"restricted and autonomous", SafetyPolicyRestricted, SafetyPolicyAutonomous, SafetyPolicyRestricted}, + {"autonomous and strict", SafetyPolicyAutonomous, SafetyPolicyStrict, SafetyPolicyStrict}, + {"autonomous and balanced", SafetyPolicyAutonomous, SafetyPolicyBalanced, SafetyPolicyBalanced}, + {"autonomous and restricted", SafetyPolicyAutonomous, SafetyPolicyRestricted, SafetyPolicyRestricted}, + {"autonomous and autonomous", SafetyPolicyAutonomous, SafetyPolicyAutonomous, SafetyPolicyAutonomous}, + {"legacy alias", "unsafe", SafetyPolicyRestricted, SafetyPolicyRestricted}, + {"unknown policy", "unknown", SafetyPolicyAutonomous, SafetyPolicyStrict}, + {"unset left", "", SafetyPolicyRestricted, ""}, + {"unset right", SafetyPolicyRestricted, "", ""}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + assert.Equal(t, test.want, MinSafetyPolicy(test.a, test.b)) + }) + } +} + // WithSafetyPolicy keeps ToolsApproved in sync both ways so legacy // readers of the flag always agree with the mode. func TestWithSafetyPolicy_SyncsToolsApproved(t *testing.T) { diff --git a/pkg/session/session.go b/pkg/session/session.go index 470c8451d2..982630e9ee 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -90,6 +90,39 @@ func (p SafetyPolicy) Normalize() SafetyPolicy { } } +// MinSafetyPolicy returns the more restrictive of two concrete safety modes. +// Under NonInteractive, strict and balanced deny calls that need confirmation, +// restricted permits only classifier-safe calls, and autonomous never asks. +// Empty is an unset legacy mode rather than an ordered policy, so an empty input +// yields empty; callers applying a policy ceiling must handle it explicitly. +func MinSafetyPolicy(a, b SafetyPolicy) SafetyPolicy { + a = a.Normalize() + b = b.Normalize() + if a == "" || b == "" { + return "" + } + + if safetyPolicyRank(a) <= safetyPolicyRank(b) { + return a + } + return b +} + +func safetyPolicyRank(policy SafetyPolicy) int { + switch policy { + case SafetyPolicyStrict: + return 0 + case SafetyPolicyBalanced: + return 1 + case SafetyPolicyRestricted: + return 2 + case SafetyPolicyAutonomous: + return 3 + default: + return 0 + } +} + // IsValid accepts current values, the legacy aliases, and empty. func (p SafetyPolicy) IsValid() bool { switch p { @@ -235,6 +268,9 @@ type Session struct { // ID is the unique identifier for the session ID string `json:"id"` + // Origin identifies the protocol surface that created this session. + Origin string `json:"origin,omitempty"` + // InputID is an optional caller-supplied correlation ID read from the eval // input file's "input_id" field. It is carried through to the output as-is // and never used internally. The session's own "id" is always a fresh UUID. @@ -1402,6 +1438,12 @@ func WithMaxToolResultTokens(n int) Opt { } } +func WithOrigin(origin string) Opt { + return func(s *Session) { + s.Origin = origin + } +} + func WithWorkingDir(workingDir string) Opt { return func(s *Session) { s.WorkingDir = workingDir @@ -1816,6 +1858,7 @@ func (s *Session) newID() string { // New creates a new agent session func New(opts ...Opt) *Session { s := &Session{ + Origin: "run", SendUserMessage: true, } diff --git a/pkg/session/store.go b/pkg/session/store.go index 966ac2ba27..72ee4ed211 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -19,9 +19,11 @@ import ( ) var ( - ErrEmptyID = errors.New("session ID cannot be empty") - ErrNotFound = errors.New("session not found") - ErrNewerDatabase = errors.New("session database was created by a newer version of docker-agent") + ErrEmptyID = errors.New("session ID cannot be empty") + ErrNotFound = errors.New("session not found") + ErrAlreadyExists = errors.New("session already exists") + ErrOriginMismatch = errors.New("session origin cannot be changed") + ErrNewerDatabase = errors.New("session database was created by a newer version of docker-agent") ) // IsRelativeSessionRef reports whether ref is a relative session reference @@ -91,7 +93,10 @@ type Summary struct { type Store interface { // === Core session operations === AddSession(ctx context.Context, session *Session) error + // GetSession retrieves a session by ID. GetSession(ctx context.Context, id string) (*Session, error) + // GetSessionByOrigin retrieves a session by ID only when it belongs to origin. + GetSessionByOrigin(ctx context.Context, id, origin string) (*Session, error) GetSessions(ctx context.Context) ([]*Session, error) GetSessionSummaries(ctx context.Context) ([]Summary, error) DeleteSession(ctx context.Context, id string) error @@ -151,7 +156,9 @@ func (s *InMemorySessionStore) AddSession(_ context.Context, session *Session) e if session.ID == "" { return ErrEmptyID } - s.sessions.Store(session.ID, session) + if _, loaded := s.sessions.LoadOrStore(session.ID, session); loaded { + return fmt.Errorf("add session %q: %w", session.ID, ErrAlreadyExists) + } return nil } @@ -166,6 +173,17 @@ func (s *InMemorySessionStore) GetSession(_ context.Context, id string) (*Sessio return session, nil } +func (s *InMemorySessionStore) GetSessionByOrigin(_ context.Context, id, origin string) (*Session, error) { + if id == "" { + return nil, ErrEmptyID + } + session, exists := s.sessions.Load(id) + if !exists || session.Origin != origin { + return nil, ErrNotFound + } + return session, nil +} + func (s *InMemorySessionStore) GetSessions(_ context.Context) ([]*Session, error) { sessions := make([]*Session, 0, s.sessions.Length()) s.sessions.Range(func(key string, value *Session) bool { @@ -227,6 +245,7 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session session.mu.RLock() newSession := &Session{ ID: session.ID, + Origin: session.Origin, Title: session.Title, Evals: session.Evals, CreatedAt: session.CreatedAt, @@ -250,9 +269,13 @@ func (s *InMemorySessionStore) UpdateSession(_ context.Context, session *Session } session.mu.RUnlock() - // Preserve existing messages if session already exists + // Preserve existing messages and reject origin changes if session already exists. if existing, exists := s.sessions.Load(session.ID); exists { existing.mu.RLock() + if existing.Origin != newSession.Origin { + existing.mu.RUnlock() + return fmt.Errorf("update session %q: %w", session.ID, ErrOriginMismatch) + } newSession.Messages = make([]Item, len(existing.Messages)) copy(newSession.Messages, existing.Messages) existing.mu.RUnlock() @@ -383,7 +406,7 @@ type SQLiteSessionStore struct { // sessionSelectColumns is the canonical SELECT list for the sessions table. // The column order matches what scanSession expects; all read paths use this // constant so that adding a column requires updating exactly one place. -const sessionSelectColumns = `id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes` +const sessionSelectColumns = `id, origin, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes` // sessionPersistedFields holds the encoded form of a Session's JSON-bearing // columns plus the SQL representation of parent_id (nil for the empty @@ -557,11 +580,11 @@ func (s *SQLiteSessionStore) AddSession(ctx context.Context, session *Session) e _, err = tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, + id, origin, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title, + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + session.ID, session.Origin, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title, session.Cost, session.SendUserMessage, session.MaxIterations, session.WorkingDir, session.CreatedAt.Format(time.RFC3339), fields.PermissionsJSON, fields.AgentModelOverridesJSON, fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON, fields.AttributesJSON) @@ -602,7 +625,7 @@ func scanSession(scanner interface { ) err := scanner.Scan( - &sess.ID, &sess.ToolsApproved, &safetyPolicy, &sess.InputTokens, &sess.OutputTokens, + &sess.ID, &sess.Origin, &sess.ToolsApproved, &safetyPolicy, &sess.InputTokens, &sess.OutputTokens, &sess.Title, &sess.Cost, &sess.SendUserMessage, &sess.MaxIterations, &workingDir, &createdAtStr, &sess.Starred, &permissionsJSON, &agentModelOverridesJSON, &customModelsUsedJSON, &thinking, &parentID, &instructionContextJSON, &attributesJSON, @@ -787,6 +810,13 @@ func (s *SQLiteSessionStore) loadSessionItems(ctx context.Context, q querier, se return items, nil } +func (s *SQLiteSessionStore) GetSessionByOrigin(ctx context.Context, id, origin string) (*Session, error) { + if id == "" { + return nil, ErrEmptyID + } + return s.loadSessionByOrigin(ctx, s.db, id, origin) +} + // loadSession retrieves a session by ID using the supplied querier. func (s *SQLiteSessionStore) loadSession(ctx context.Context, q querier, id string) (*Session, error) { row := q.QueryRowContext(ctx, @@ -808,6 +838,26 @@ func (s *SQLiteSessionStore) loadSession(ctx context.Context, q querier, id stri return sess, nil } +func (s *SQLiteSessionStore) loadSessionByOrigin(ctx context.Context, q querier, id, origin string) (*Session, error) { + row := q.QueryRowContext(ctx, + "SELECT "+sessionSelectColumns+" FROM sessions WHERE id = ? AND origin = ?", id, origin) + + sess, err := scanSession(row) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrNotFound + } + return nil, err + } + + sess.Messages, err = s.loadSessionItems(ctx, q, id) + if err != nil { + return nil, fmt.Errorf("loading session items: %w", err) + } + + return sess, nil +} + // GetSessions retrieves all root sessions (excludes sub-sessions) func (s *SQLiteSessionStore) GetSessions(ctx context.Context) ([]*Session, error) { rows, err := s.db.QueryContext(ctx, @@ -922,6 +972,7 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session session.mu.RLock() snapshot := &Session{ ID: session.ID, + Origin: session.Origin, Title: session.Title, CreatedAt: session.CreatedAt, ToolsApproved: session.ToolsApproved, @@ -956,13 +1007,13 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session defer func() { _ = tx.Rollback() }() // Use INSERT OR REPLACE for upsert behavior - creates if not exists, updates if exists - _, err = tx.ExecContext(ctx, + result, err := tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, + id, origin, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, tools_approved = excluded.tools_approved, @@ -980,14 +1031,22 @@ func (s *SQLiteSessionStore) UpdateSession(ctx context.Context, session *Session thinking = excluded.thinking, parent_id = excluded.parent_id, instruction_context = excluded.instruction_context, - attributes = excluded.attributes`, - snapshot.ID, snapshot.ToolsApproved, string(snapshot.SafetyPolicy), snapshot.InputTokens, snapshot.OutputTokens, + attributes = excluded.attributes + WHERE sessions.origin = excluded.origin`, + snapshot.ID, snapshot.Origin, snapshot.ToolsApproved, string(snapshot.SafetyPolicy), snapshot.InputTokens, snapshot.OutputTokens, snapshot.Title, snapshot.Cost, snapshot.SendUserMessage, snapshot.MaxIterations, snapshot.WorkingDir, snapshot.CreatedAt.Format(time.RFC3339), snapshot.Starred, fields.PermissionsJSON, fields.AgentModelOverridesJSON, fields.CustomModelsUsedJSON, false, fields.ParentID, fields.InstructionContextJSON, fields.AttributesJSON) if err != nil { return err } + rowsAffected, err := result.RowsAffected() + if err != nil { + return err + } + if rowsAffected == 0 { + return fmt.Errorf("update session %q: %w", session.ID, ErrOriginMismatch) + } // Note: Messages are NOT persisted here. They are persisted via events // (UserMessageEvent, MessageAddedEvent, etc.) to avoid duplication. @@ -1129,12 +1188,12 @@ func (s *SQLiteSessionStore) addSessionTx(ctx context.Context, tx *sql.Tx, sessi _, err = tx.ExecContext(ctx, `INSERT INTO sessions ( - id, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, + id, origin, tools_approved, safety_policy, input_tokens, output_tokens, title, cost, send_user_message, max_iterations, working_dir, created_at, starred, permissions, agent_model_overrides, custom_models_used, thinking, parent_id, instruction_context, attributes ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - session.ID, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + session.ID, session.Origin, session.ToolsApproved, string(session.SafetyPolicy), session.InputTokens, session.OutputTokens, session.Title, session.Cost, session.SendUserMessage, session.MaxIterations, session.WorkingDir, session.CreatedAt.Format(time.RFC3339), session.Starred, fields.PermissionsJSON, fields.AgentModelOverridesJSON, fields.CustomModelsUsedJSON, false, diff --git a/pkg/session/store_memory_test.go b/pkg/session/store_memory_test.go index b39e1af40c..60f866e7e9 100644 --- a/pkg/session/store_memory_test.go +++ b/pkg/session/store_memory_test.go @@ -76,6 +76,94 @@ func TestNewSQLiteSessionStoreFromDB_RoundTripWithMessages(t *testing.T) { assert.Equal(t, "world", got.Messages[1].Message.Message.Content) } +// TestSessionOriginCannotChange verifies both stores reject an upsert that +// attempts to reclassify a persisted session. +func TestSessionOriginCannotChange(t *testing.T) { + t.Parallel() + ctx := t.Context() + + for name, newStore := range map[string]func() Store{ + "in memory": NewInMemorySessionStore, + "sqlite": func() Store { return openMemoryStore(t) }, + } { + t.Run(name, func(t *testing.T) { + store := newStore() + existing := New(WithID("origin-immutable"), WithOrigin("run"), WithTitle("original")) + require.NoError(t, store.AddSession(ctx, existing)) + + replacement := New(WithID(existing.ID), WithOrigin("a2a"), WithTitle("replacement")) + require.ErrorIs(t, store.UpdateSession(ctx, replacement), ErrOriginMismatch) + + got, err := store.GetSession(ctx, existing.ID) + require.NoError(t, err) + require.Equal(t, "run", got.Origin) + require.Equal(t, "original", got.Title) + }) + } +} + +// TestAddSessionRejectsConflictingOrigin verifies that adding a duplicate ID cannot +// reclassify an existing session, including in the SQLite store used by A2A. +func TestAddSessionRejectsConflictingOrigin(t *testing.T) { + t.Parallel() + ctx := t.Context() + + for name, newStore := range map[string]func() Store{ + "in memory": NewInMemorySessionStore, + "sqlite": func() Store { return openMemoryStore(t) }, + } { + t.Run(name, func(t *testing.T) { + store := newStore() + existing := New(WithID("origin-add-conflict"), WithOrigin("run"), WithTitle("original")) + require.NoError(t, store.AddSession(ctx, existing)) + + err := store.AddSession(ctx, New(WithID(existing.ID), WithOrigin("a2a"), WithTitle("replacement"))) + require.Error(t, err) + if name == "in memory" { + require.ErrorIs(t, err, ErrAlreadyExists) + } + + got, err := store.GetSession(ctx, existing.ID) + require.NoError(t, err) + assert.Equal(t, "run", got.Origin) + assert.Equal(t, "original", got.Title) + }) + } +} + +func TestSessionOriginRoundTripAndLookup(t *testing.T) { + t.Parallel() + ctx := t.Context() + + for name, newStore := range map[string]func() Store{ + "in memory": NewInMemorySessionStore, + "sqlite": func() Store { return openMemoryStore(t) }, + } { + t.Run(name, func(t *testing.T) { + store := newStore() + runSession := New(WithID("run-session")) + a2aSession := New(WithID("a2a-session"), WithOrigin("a2a")) + require.NoError(t, store.AddSession(ctx, runSession)) + require.NoError(t, store.AddSession(ctx, a2aSession)) + + got, err := store.GetSessionByOrigin(ctx, a2aSession.ID, "a2a") + require.NoError(t, err) + assert.Equal(t, "a2a", got.Origin) + + _, err = store.GetSessionByOrigin(ctx, runSession.ID, "a2a") + require.ErrorIs(t, err, ErrNotFound) + _, err = store.GetSessionByOrigin(ctx, "", "a2a") + require.ErrorIs(t, err, ErrEmptyID) + + a2aSession.SetTitle("updated") + require.NoError(t, store.UpdateSession(ctx, a2aSession)) + got, err = store.GetSession(ctx, a2aSession.ID) + require.NoError(t, err) + assert.Equal(t, "a2a", got.Origin) + }) + } +} + // TestMigration23_LegacySummaryRowsReadAsZeroCost simulates a database // created before migration 023 (no cost column on session_items, summary // rows written without one): opening the store applies the migration and @@ -114,4 +202,5 @@ func TestMigration23_LegacySummaryRowsReadAsZeroCost(t *testing.T) { assert.Equal(t, "old summary", got.Messages[0].Summary) assert.Equal(t, 2, got.Messages[0].FirstKeptEntry) assert.Zero(t, got.Messages[0].Cost, "legacy summary rows must read as cost 0") + assert.Equal(t, "run", got.Origin, "pre-origin sessions must be classified as normal runs") }