coder · johnstcn · Sep 18, 2026 · Sep 11, 2026 · Sep 11, 2026 · Sep 11, 2026
diff --git a/aibridge/recorder/drpc_recorder.go b/aibridge/recorder/drpc_recorder.go
@@ -10,6 +10,7 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	agplaibridge "github.com/coder/coder/v2/coderd/aibridge"
 	"github.com/coder/coder/v2/coderd/aibridged/proto"
 )
 
@@ -22,7 +23,7 @@ type DRPCRecorder struct {
 }
 
 func (t *DRPCRecorder) RecordInterception(ctx context.Context, req *InterceptionRecord) error {
-	_, err := t.client.RecordInterception(ctx, &proto.RecordInterceptionRequest{
+	in := &proto.RecordInterceptionRequest{
 		Id:                          req.ID,
 		ApiKeyId:                    t.apiKeyID,
 		InitiatorId:                 req.InitiatorID,
@@ -39,7 +40,11 @@ func (t *DRPCRecorder) RecordInterception(ctx context.Context, req *Interception
 		CredentialHint:              req.CredentialHint,
 		AgentFirewallSessionId:      req.AgentFirewallSessionID,
 		AgentFirewallSequenceNumber: req.AgentFirewallSequenceNumber,
-	})
+	}
+	if attr, ok := agplaibridge.AttributionFromContext(ctx); ok {
+		in.WorkspaceId = attr.WorkspaceID.String()
+	}
+	_, err := t.client.RecordInterception(ctx, in)
 	return err
 }
 

diff --git a/coderd/aibridge/factory.go b/coderd/aibridge/factory.go
@@ -3,8 +3,33 @@ package aibridge
 import (
 	"context"
 	"net/http"
+
+	"github.com/google/uuid"
 )
 
+// Attribution carries contextual per-request attribution data for a request.
+type Attribution struct {
+	// WorkspaceID is the workspace bound to the chat, or uuid.Nil when no
+	// workspace is bound.
+	WorkspaceID uuid.UUID
+}
+
+type attributionCtxKey struct{}
+
+// WithAttribution returns a copy of ctx carrying the trusted Attribution.
+// Attribution MUST ONLY set by authentication or delegated in-process code,
+// and NEVER by client-provided HTTP headers.
+func WithAttribution(ctx context.Context, attr Attribution) context.Context {
+	return context.WithValue(ctx, attributionCtxKey{}, attr)
+}
+
+// AttributionFromContext returns the Attribution attached by [WithAttribution]
+// and whether a non-zero WorkspaceID was present.
+func AttributionFromContext(ctx context.Context) (Attribution, bool) {
+	attr, ok := ctx.Value(attributionCtxKey{}).(Attribution)
+	return attr, ok && attr.WorkspaceID != uuid.Nil
+}
+
 // Source identifies the call site that asked aibridge for a transport. It is
 // attached to the request context so downstream handlers and logs can attribute
 // traffic without changing behavior based on the value.
@@ -29,7 +54,10 @@ func SourceFromContext(ctx context.Context) Source {
 	return src
 }
 
-type delegatedAPIKeyIDCtxKey struct{}
+type (
+	deletedAPIKeyIDCtxKey      struct{}
+	delegatedAttributionCtxKey struct{}
+)
 
 // WithDelegatedAPIKeyID returns a copy of ctx carrying an API key ID on whose
 // behalf the request is being made. The in-process aibridge transport requires
@@ -40,16 +68,31 @@ type delegatedAPIKeyIDCtxKey struct{}
 // has not expired, and belongs to a non-deleted, non-system user. It does not
 // verify the key secret, because the caller never has it.
 func WithDelegatedAPIKeyID(ctx context.Context, id string) context.Context {
-	return context.WithValue(ctx, delegatedAPIKeyIDCtxKey{}, id)
+	/*
+		return WithDelegatedRequest(ctx, id, Attribution{})
+	*/
+	return context.WithValue(ctx, deletedAPIKeyIDCtxKey{}, id)
+}
+
+func WithDelegatedAttribution(ctx context.Context, attr Attribution) context.Context {
+	return context.WithValue(ctx, delegatedAttributionCtxKey{}, attr)
 }
 
 // DelegatedAPIKeyIDFromContext returns the API key ID attached by
 // [WithDelegatedAPIKeyID] and whether a non-empty value was set.
 func DelegatedAPIKeyIDFromContext(ctx context.Context) (string, bool) {
-	id, ok := ctx.Value(delegatedAPIKeyIDCtxKey{}).(string)
+	id, ok := ctx.Value(deletedAPIKeyIDCtxKey{}).(string)
 	return id, ok && id != ""
 }
 
+// DelegatedAttributionFromContext returns the trusted attribution attached to
+// the delegated request. Its boolean reports whether delegated authentication,
+// rather than attribution itself, was present.
+func DelegatedAttributionFromContext(ctx context.Context) (Attribution, bool) {
+	attr, ok := ctx.Value(delegatedAttributionCtxKey{}).(Attribution)
+	return attr, ok
+}
+
 // TransportFactory returns an [http.RoundTripper] that dispatches an aibridge
 // request in-process for a given provider instance name.
 //

diff --git a/coderd/aibridge/factory_test.go b/coderd/aibridge/factory_test.go
@@ -0,0 +1,64 @@
+package aibridge_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/coderd/aibridge"
+)
+
+func TestAttributionContext(t *testing.T) {
+	t.Parallel()
+
+	wsID := uuid.New()
+
+	t.Run("WithWorkspace", func(t *testing.T) {
+		t.Parallel()
+
+		attr := aibridge.Attribution{WorkspaceID: wsID}
+		ctx := aibridge.WithAttribution(context.Background(), attr)
+		got, ok := aibridge.AttributionFromContext(ctx)
+		require.True(t, ok)
+		require.Equal(t, attr, got)
+	})
+
+	t.Run("NoAttribution", func(t *testing.T) {
+		t.Parallel()
+
+		_, ok := aibridge.AttributionFromContext(context.Background())
+		require.False(t, ok)
+	})
+
+	t.Run("ZeroWorkspaceID", func(t *testing.T) {
+		t.Parallel()
+
+		// An Attribution with a zero WorkspaceID is not considered set.
+		attr := aibridge.Attribution{WorkspaceID: uuid.Nil}
+		ctx := aibridge.WithAttribution(context.Background(), attr)
+		_, ok := aibridge.AttributionFromContext(ctx)
+		require.False(t, ok, "Attribution with zero WorkspaceID must not be considered set")
+	})
+
+	t.Run("OverwrittenByInnerContext", func(t *testing.T) {
+		t.Parallel()
+
+		// A child context can shadow the parent's Attribution without
+		// mutating it, ensuring per-request immutability.
+		parent := aibridge.Attribution{WorkspaceID: uuid.New()}
+		child := aibridge.Attribution{WorkspaceID: wsID}
+
+		pCtx := aibridge.WithAttribution(context.Background(), parent)
+		cCtx := aibridge.WithAttribution(pCtx, child)
+
+		gotParent, ok := aibridge.AttributionFromContext(pCtx)
+		require.True(t, ok)
+		require.Equal(t, parent, gotParent)
+
+		gotChild, ok := aibridge.AttributionFromContext(cCtx)
+		require.True(t, ok)
+		require.Equal(t, child, gotChild)
+	})
+}
diff --git a/coderd/aibridged/http.go b/coderd/aibridged/http.go
@@ -145,6 +145,23 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 	}
 	logger = logger.With(slog.F("user_id", id))
 
+	// Direct and delegated carriers are mutually exclusive and request-scoped.
+	// Delegated requests carry attribution stamped by the in-process caller
+	// (chatd); direct requests carry workspace attribution parsed from the
+	// server-minted token name by IsAuthorized.
+	var attribution agplaibridge.Attribution
+	if delegated {
+		attribution, _ = agplaibridge.DelegatedAttributionFromContext(ctx)
+	} else {
+		attribution, err = attributionFromAuthorization(resp)
+		if err != nil {
+			logger.Warn(ctx, "invalid authorization attribution", slog.Error(err))
+			http.Error(rw, ErrUnauthorized.Error(), http.StatusForbidden)
+			return
+		}
+	}
+	ctx = agplaibridge.WithAttribution(ctx, attribution)
+
 	budgetResp, err := client.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
 		UserId: id.String(),
 	})
@@ -183,3 +200,17 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 
 	handler.ServeHTTP(rw, r)
 }
+
+// attributionFromAuthorization extracts the workspace ID from the
+// IsAuthorizedResponse. The proto carries only workspace_id for attribution;
+// organization and workspace name are not returned.
+func attributionFromAuthorization(resp *proto.IsAuthorizedResponse) (agplaibridge.Attribution, error) {
+	if resp.GetWorkspaceId() == "" {
+		return agplaibridge.Attribution{}, nil
+	}
+	workspaceID, err := uuid.Parse(resp.GetWorkspaceId())
+	if err != nil {
+		return agplaibridge.Attribution{}, xerrors.Errorf("parse workspace ID: %w", err)
+	}
+	return agplaibridge.Attribution{WorkspaceID: workspaceID}, nil
+}