feat(config): session correlation header injection configuration

cli/cli.go

@@ -37,7 +37,22 @@ func NewCommand(version string) *serpent.Command {
   # Use allowlist from config file with additional CLI allow rules
   boundary --allow "domain=example.com" -- curl https://example.com
 
-  # Block everything by default (implicit)`
+  # Block everything by default (implicit)
+
+  # Enable session correlation inside a Coder workspace (inject target auto-derived from CODER_AGENT_URL)
+  boundary --enable-session-correlation \
+    --allow "domain=dev.coder.com" -- python train.py
+
+  # Enable session correlation with an explicit inject target (e.g. outside a workspace or custom deployment)
+  boundary --enable-session-correlation \
+    --session-id-inject-target "domain=mydeployment.coder.com path=/api/v2/aibridge/*" \
+    --allow "domain=mydeployment.coder.com" -- python train.py
+
+  # Enable session correlation with multiple inject targets (e.g. staging + prod)
+  boundary --enable-session-correlation \
+    --session-id-inject-target "domain=staging.coder.com path=/api/v2/aibridge/*" \
+    --session-id-inject-target "domain=prod.coder.com path=/api/v2/aibridge/*" \
+    --allow "domain=staging.coder.com" --allow "domain=prod.coder.com" -- python train.py`
 
 	return cmd
 }
@@ -169,6 +184,43 @@ func BaseCommand(version string) *serpent.Command {
 				Value:       &showVersion,
 				YAML:        "", // CLI only
 			},
+			// Session correlation header injection options.
+			{
+				Flag:        "enable-session-correlation",
+				Env:         "BOUNDARY_SESSION_CORRELATION_ENABLED",
+				Description: "Enable session correlation header injection. When no inject targets are configured, the target is auto-derived from CODER_AGENT_URL (set automatically inside Coder workspaces). Disable for deployments without Coder AI Gateway in front.",
+				Value:       &cliConfig.SessionCorrelationEnabled,
+				YAML:        "session_correlation_enabled",
+			},
+			{
+				Flag:        "session-id-inject-target",
+				Env:         "BOUNDARY_SESSION_ID_INJECT_TARGET",
+				Description: `Inject target (repeatable via flag; env accepts one value). Requests matching these targets receive session correlation headers. For multiple targets use the YAML config. Format: "domain=<host> [path=<glob>]".`,
+				Value:       &cliConfig.InjectSessionIDTarget,
+				YAML:        "", // CLI only, YAML uses session_id_inject_targets.
+			},
+			{
+				Flag:        "", // No CLI flag, YAML only.
+				Description: "Inject targets from config file (YAML only).",
+				Value:       &cliConfig.InjectSessionIDTargets,
+				YAML:        "session_id_inject_targets",
+			},
+			{
+				Flag:        "session-id-header-name",
+				Env:         "BOUNDARY_SESSION_ID_HEADER_NAME",
+				Description: "HTTP header name for the boundary session ID.",
+				Default:     config.DefaultSessionIDHeaderName,
+				Value:       &cliConfig.SessionIDHeaderName,
+				YAML:        "session_id_header_name",
+			},
+			{
+				Flag:        "sequence-number-header-name",
+				Env:         "BOUNDARY_SEQUENCE_NUMBER_HEADER_NAME",
+				Description: "HTTP header name for the boundary sequence number.",
+				Default:     config.DefaultSequenceNumberHeaderName,
+				Value:       &cliConfig.SequenceNumberHeaderName,
+				YAML:        "sequence_number_header_name",
+			},
 		},
 		Handler: func(inv *serpent.Invocation) error {
 			// Handle --version flag early

config/config.go

@@ -2,6 +2,7 @@ package config
 
 import (
 	"fmt"
+	"os"
 	"strings"
 
 	"github.com/coder/serpent"
@@ -70,6 +71,13 @@ type CliConfig struct {
 	NoUserNamespace    serpent.Bool           `yaml:"no_user_namespace"`
 	DisableAuditLogs   serpent.Bool           `yaml:"disable_audit_logs"`
 	LogProxySocketPath serpent.String         `yaml:"log_proxy_socket_path"`
+
+	// Session correlation header injection.
+	SessionCorrelationEnabled serpent.Bool        `yaml:"session_correlation_enabled"`
+	InjectSessionIDTarget     AllowStringsArray   `yaml:"-"`                         // From CLI flags only
+	InjectSessionIDTargets    serpent.StringArray `yaml:"session_id_inject_targets"` // From config file
+	SessionIDHeaderName       serpent.String      `yaml:"session_id_header_name"`
+	SequenceNumberHeaderName  serpent.String      `yaml:"sequence_number_header_name"`
 }
 
 type AppConfig struct {
@@ -87,6 +95,10 @@ type AppConfig struct {
 	DisableAuditLogs   bool
 	LogProxySocketPath string
 
+	// SessionCorrelation controls header injection for AI Bridge
+	// correlation. See SessionCorrelationConfig for details.
+	SessionCorrelation SessionCorrelationConfig
+
 	// SessionID is a UUIDv4 generated at process startup. It groups
 	// all audit events produced by this boundary invocation into a
 	// single session. Set by Run, not by configuration.
@@ -108,6 +120,12 @@ func NewAppConfigFromCliConfig(cfg CliConfig, targetCMD []string) (AppConfig, er
 
 	userInfo := GetUserInfo()
 
+	// Build session correlation config from CLI and YAML sources.
+	sc, err := buildSessionCorrelation(cfg, os.Environ())
+	if err != nil {
+		return AppConfig{}, fmt.Errorf("session correlation config: %w", err)
+	}
+
 	return AppConfig{
 		AllowRules:         allAllowStrings,
 		LogLevel:           cfg.LogLevel.Value(),
@@ -122,5 +140,54 @@ func NewAppConfigFromCliConfig(cfg CliConfig, targetCMD []string) (AppConfig, er
 		UserInfo:           userInfo,
 		DisableAuditLogs:   cfg.DisableAuditLogs.Value(),
 		LogProxySocketPath: cfg.LogProxySocketPath.Value(),
+		SessionCorrelation: sc,
 	}, nil
 }
+
+// buildSessionCorrelation merges CLI and YAML inject target sources,
+// parses each target string, applies header name defaults, and
+// validates the resulting configuration. environ is passed explicitly
+// (rather than reading os.Environ inside) so that callers and tests
+// can supply a controlled environment.
+func buildSessionCorrelation(cfg CliConfig, environ []string) (SessionCorrelationConfig, error) {
+	// Merge YAML targets with CLI targets.
+	rawTargets := append(cfg.InjectSessionIDTargets.Value(), cfg.InjectSessionIDTarget.Value()...)
+
+	var targets []InjectTarget
+	for _, raw := range rawTargets {
+		t, err := ParseInjectTarget(raw)
+		if err != nil {
+			return SessionCorrelationConfig{}, err
+		}
+		targets = append(targets, t)
+	}
+
+	if len(targets) == 0 && cfg.SessionCorrelationEnabled.Value() {
+		if t := DefaultInjectTargetFromEnv(environ); t != nil {
+			targets = []InjectTarget{*t}
+		}
+	}
+
+	// Apply defaults for header names.
+	sessionIDHeader := cfg.SessionIDHeaderName.Value()
+	if sessionIDHeader == "" {
+		sessionIDHeader = DefaultSessionIDHeaderName
+	}
+	seqHeader := cfg.SequenceNumberHeaderName.Value()
+	if seqHeader == "" {
+		seqHeader = DefaultSequenceNumberHeaderName
+	}
+
+	sc := SessionCorrelationConfig{
+		Enabled:                  cfg.SessionCorrelationEnabled.Value(),
+		InjectTargets:            targets,
+		SessionIDHeaderName:      sessionIDHeader,
+		SequenceNumberHeaderName: seqHeader,
+	}
+
+	if err := ValidateSessionCorrelation(sc); err != nil {
+		return SessionCorrelationConfig{}, err
+	}
+
+	return sc, nil
+}

config/session_correlation.go

@@ -0,0 +1,147 @@
+package config
+
+import (
+	"fmt"
+	"net/url"
+	"strings"
+)
+
+// Default header names and paths for session correlation.
+const (
+	DefaultSessionIDHeaderName      = "X-Coder-Agent-Firewall-Session-Id"
+	DefaultSequenceNumberHeaderName = "X-Coder-Agent-Firewall-Sequence-Number"
+
+	// DefaultAIBridgePath is the path glob used when auto-deriving an inject
+	// target from CODER_AGENT_URL.
+	DefaultAIBridgePath = "/api/v2/aibridge/*"
+
+	// CoderAgentURLEnv is the environment variable set by the Coder workspace
+	// agent that points to the control plane. Boundary uses it to derive a
+	// default inject target when none is explicitly configured.
+	CoderAgentURLEnv = "CODER_AGENT_URL"
+)
+
+// InjectTarget represents a parsed target for session correlation header
+// injection. Requests matching the domain (and optional path glob) will
+// receive the session ID and sequence number headers.
+type InjectTarget struct {
+	Domain string
+	Path   string
+}
+
+// SessionCorrelationConfig holds configuration for session correlation
+// header injection. When enabled, boundary injects its session ID and
+// sequence number as custom headers on matching outbound requests so
+// that an upstream AI Bridge can correlate the request back to the
+// boundary audit event stream.
+type SessionCorrelationConfig struct {
+	// Enabled controls whether session correlation headers are injected.
+	// Deployments without AI Bridge in front should set this to false.
+	Enabled bool
+
+	// InjectTargets is the list of domain/path patterns that should
+	// receive session correlation headers.
+	InjectTargets []InjectTarget
+
+	// SessionIDHeaderName is the HTTP header name used to carry the
+	// boundary session ID. Defaults to DefaultSessionIDHeaderName.
+	SessionIDHeaderName string
+
+	// SequenceNumberHeaderName is the HTTP header name used to carry
+	// the boundary sequence number. Defaults to
+	// DefaultSequenceNumberHeaderName.
+	SequenceNumberHeaderName string
+}
+
+// ParseInjectTarget parses a string of the form "domain=... path=..."
+// into an InjectTarget. The domain key is required; path is optional.
+func ParseInjectTarget(raw string) (InjectTarget, error) {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return InjectTarget{}, fmt.Errorf("inject target must not be empty")
+	}
+
+	var target InjectTarget
+	for _, part := range strings.Fields(raw) {
+		key, value, ok := strings.Cut(part, "=")
+		if !ok {
+			return InjectTarget{}, fmt.Errorf(
+				"inject target: malformed key-value pair %q, expected key=value", part,
+			)
+		}
+		switch key {
+		case "domain":
+			if value == "" {
+				return InjectTarget{}, fmt.Errorf("inject target: domain must not be empty")
+			}
+			target.Domain = value
+		case "path":
+			target.Path = value
+		default:
+			return InjectTarget{}, fmt.Errorf("inject target: unknown key %q", key)
+		}
+	}
+
+	if target.Domain == "" {
+		return InjectTarget{}, fmt.Errorf("inject target: domain is required")
+	}
+
+	return target, nil
+}
+
+// DefaultInjectTargetFromEnv derives an InjectTarget from the CODER_AGENT_URL
+// variable in the provided environment slice. It returns nil if the variable is
+// absent, empty, or not a valid URL with a host. The derived target uses
+// DefaultAIBridgePath as the path glob so that all AI Bridge traffic on the
+// control-plane host is matched.
+//
+// The environ parameter is accepted rather than reading os.Environ directly so
+// that callers (and tests) can supply an arbitrary environment.
+func DefaultInjectTargetFromEnv(environ []string) *InjectTarget {
+	var raw string
+	for _, e := range environ {
+		k, v, ok := strings.Cut(e, "=")
+		if ok && k == CoderAgentURLEnv {
+			raw = v
+			break
+		}
+	}
+	if raw == "" {
+		return nil
+	}
+
+	u, err := url.Parse(raw)
+	if err != nil || u.Host == "" {
+		return nil
+	}
+
+	return &InjectTarget{
+		Domain: u.Hostname(),
+		Path:   DefaultAIBridgePath,
+	}
+}
+
+// ValidateSessionCorrelation checks that the session correlation config
+// is internally consistent. It returns an error describing the first
+// problem found, or nil if the config is valid.
+func ValidateSessionCorrelation(cfg SessionCorrelationConfig) error {
+	if !cfg.Enabled {
+		return nil
+	}
+
+	if len(cfg.InjectTargets) == 0 {
+		return fmt.Errorf(
+			"session correlation is enabled but no inject targets are configured",
+		)
+	}
+
+	if cfg.SessionIDHeaderName == "" {
+		return fmt.Errorf("session-id-header-name must not be empty when session correlation is enabled")
+	}
+
+	if cfg.SequenceNumberHeaderName == "" {
+		return fmt.Errorf("sequence-number-header-name must not be empty when session correlation is enabled")
+	}
+
+	return nil
+}