test(proxy): integration tests for session correlation audit and header agreement

audit/request.go

@@ -11,4 +11,11 @@ type Request struct {
 	Host    string
 	Allowed bool
 	Rule    string // The rule that matched (if any)
+
+	// SequenceNumber is a pre-allocated sequence number for this
+	// audit event. When non-nil the auditor must use this value
+	// instead of generating its own so that the audit log and
+	// any injected HTTP header carry the same number. When nil
+	// the auditor falls back to its internal SequenceCounter.
+	SequenceNumber *uint64
 }

audit/socket_auditor.go

@@ -80,10 +80,17 @@ func (s *SocketAuditor) AuditRequest(req Request) {
 		httpReq.MatchedRule = req.Rule
 	}
 
+	var seqNum uint64
+	if req.SequenceNumber != nil {
+		seqNum = *req.SequenceNumber
+	} else {
+		seqNum = s.seq.Next()
+	}
+
 	log := &agentproto.BoundaryLog{
 		Allowed:        req.Allowed,
 		Time:           timestamppb.Now(),
-		SequenceNumber: s.seq.Next(),
+		SequenceNumber: seqNum,
 		Resource:       &agentproto.BoundaryLog_HttpRequest_{HttpRequest: httpReq},
 	}
 

cli/cli.go

@@ -169,6 +169,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. Disable for deployments without AI Bridge in front.",
+				Value:       &cliConfig.SessionCorrelationEnabled,
+				YAML:        "session_correlation_enabled",
+			},
+			{
+				Flag:        "inject-session-id-on",
+				Env:         "BOUNDARY_INJECT_SESSION_ID_ON",
+				Description: `Inject target (repeatable). Requests matching these targets receive session correlation headers. Format: "domain=<host> [path=<glob>]".`,
+				Value:       &cliConfig.InjectSessionIDOn,
+				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.InjectSessionIDOnYAML,
+				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

@@ -69,6 +69,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"`
+	InjectSessionIDOn         AllowStringsArray      `yaml:"inject_session_id_on"`
+	InjectSessionIDOnYAML     serpent.StringArray     `yaml:"session_id_inject_targets"`
+	SessionIDHeaderName       serpent.String          `yaml:"session_id_header_name"`
+	SequenceNumberHeaderName  serpent.String          `yaml:"sequence_number_header_name"`
 }
 
 type AppConfig struct {
@@ -86,6 +93,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.
@@ -107,6 +118,12 @@ func NewAppConfigFromCliConfig(cfg CliConfig, targetCMD []string) (AppConfig, er
 
 	userInfo := GetUserInfo()
 
+	// Build session correlation config from CLI and YAML sources.
+	sc, err := buildSessionCorrelation(cfg)
+	if err != nil {
+		return AppConfig{}, fmt.Errorf("session correlation config: %w", err)
+	}
+
 	return AppConfig{
 		AllowRules:         allAllowStrings,
 		LogLevel:           cfg.LogLevel.Value(),
@@ -121,5 +138,46 @@ 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.
+func buildSessionCorrelation(cfg CliConfig) (SessionCorrelationConfig, error) {
+	// Merge YAML targets with CLI targets.
+	rawTargets := append(cfg.InjectSessionIDOnYAML.Value(), cfg.InjectSessionIDOn.Value()...)
+
+	var targets []InjectTarget
+	for _, raw := range rawTargets {
+		t, err := ParseInjectTarget(raw)
+		if err != nil {
+			return SessionCorrelationConfig{}, err
+		}
+		targets = append(targets, 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,105 @@
+package config
+
+import (
+	"fmt"
+	"strings"
+)
+
+// Default header names for session correlation.
+const (
+	DefaultSessionIDHeaderName      = "X-Coder-Agent-Firewall-Session-Id"
+	DefaultSequenceNumberHeaderName = "X-Coder-Agent-Firewall-Sequence-Number"
+)
+
+// 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
+}
+
+// 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
+}