coderd/aibridgedserver/aibridgedserver_test.go

Rama main · réplica de lectura del repositorio original.

package aibridgedserver_test

import (
	"bufio"
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"math"
	"net"
	"net/url"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/prometheus/client_golang/prometheus"
	promtest "github.com/prometheus/client_golang/prometheus/testutil"
	"github.com/sqlc-dev/pqtype"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"
	"golang.org/x/xerrors"
	protobufproto "google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/anypb"
	"google.golang.org/protobuf/types/known/structpb"
	"google.golang.org/protobuf/types/known/timestamppb"
	"storj.io/drpc"
	"storj.io/drpc/drpcerr"

	"cdr.dev/slog/v3"
	"cdr.dev/slog/v3/sloggers/slogjson"
	"cdr.dev/slog/v3/sloggers/slogtest"
	"github.com/coder/coder/v2/aibridge/recorder"
	"github.com/coder/coder/v2/coderd/aibridged"
	"github.com/coder/coder/v2/coderd/aibridged/proto"
	"github.com/coder/coder/v2/coderd/aibridgedserver"
	agplaiseats "github.com/coder/coder/v2/coderd/aiseats"
	"github.com/coder/coder/v2/coderd/apikey"
	"github.com/coder/coder/v2/coderd/coderdtest"
	"github.com/coder/coder/v2/coderd/coderdtest/promhelp"
	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/database/dbauthz"
	"github.com/coder/coder/v2/coderd/database/dbfake"
	"github.com/coder/coder/v2/coderd/database/dbgen"
	"github.com/coder/coder/v2/coderd/database/dbmock"
	"github.com/coder/coder/v2/coderd/database/dbtestutil"
	"github.com/coder/coder/v2/coderd/database/dbtime"
	"github.com/coder/coder/v2/coderd/database/pubsub"
	"github.com/coder/coder/v2/coderd/externalauth"
	codermcp "github.com/coder/coder/v2/coderd/mcp"
	"github.com/coder/coder/v2/coderd/notifications"
	"github.com/coder/coder/v2/coderd/notifications/notificationstest"
	"github.com/coder/coder/v2/coderd/provisionerdserver"
	coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
	"github.com/coder/coder/v2/coderd/rbac"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/coder/v2/cryptorand"
	"github.com/coder/coder/v2/testutil"
	"github.com/coder/quartz"
	"github.com/coder/serpent"
)

var requiredExperiments = []codersdk.Experiment{
	codersdk.ExperimentMCPServerHTTP,
}

// TestAuthorization validates the authorization logic.
// No other tests are explicitly defined in this package because aibridgedserver is
// tested via integration tests in the aibridged package (see aibridged/aibridged_integration_test.go).
func TestAuthorization(t *testing.T) {
	t.Parallel()

	// A valid request without a workspace resolves to the Everyone group of
	// this organization.
	everyoneOrgID := uuid.New()

	cases := []struct {
		name string
		// Key will be set to the same key passed to mocksFn if unset.
		key string
		// mocksFn is called with a valid API key and user. If the test needs
		// invalid values, it should just mutate them directly.
		mocksFn     func(db *dbmock.MockStore, apiKey database.APIKey, user database.User)
		expectedErr error
	}{
		{
			name:        "invalid key format",
			key:         "foo",
			expectedErr: aibridgedserver.ErrInvalidKey,
		},
		{
			name:        "unknown key",
			expectedErr: aibridgedserver.ErrUnknownKey,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(database.APIKey{}, sql.ErrNoRows)
			},
		},
		{
			name:        "expired",
			expectedErr: aibridgedserver.ErrExpired,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				apiKey.ExpiresAt = dbtime.Now().Add(-time.Hour)
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
			},
		},
		{
			name:        "invalid key secret",
			expectedErr: aibridgedserver.ErrInvalidKey,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				apiKey.HashedSecret = []byte("differentsecret")
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
			},
		},
		{
			name:        "unknown user",
			expectedErr: aibridgedserver.ErrUnknownUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(database.User{}, sql.ErrNoRows)
			},
		},
		{
			name:        "deleted user",
			expectedErr: aibridgedserver.ErrDeletedUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Deleted = true
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			name:        "suspended user",
			expectedErr: aibridgedserver.ErrInactiveUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Status = database.UserStatusSuspended
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			name:        "dormant user",
			expectedErr: aibridgedserver.ErrInactiveUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Status = database.UserStatusDormant
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			name:        "system user",
			expectedErr: aibridgedserver.ErrSystemUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.IsSystem = true
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			name: "valid",
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
				expectSiteWideEveryoneBudgetFallback(db, user.ID, everyoneOrgID)
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			logger := testutil.Logger(t)

			// Make a fake user and an API key for the mock calls.
			now := dbtime.Now()
			user := database.User{
				ID:         uuid.New(),
				Email:      "test@coder.com",
				Username:   "test",
				Name:       "Test User",
				CreatedAt:  now,
				UpdatedAt:  now,
				RBACRoles:  []string{},
				LoginType:  database.LoginTypePassword,
				Status:     database.UserStatusActive,
				LastSeenAt: now,
			}

			keyID, _ := cryptorand.String(10)
			keySecret, keySecretHashed, _ := apikey.GenerateSecret(22)
			token := fmt.Sprintf("%s-%s", keyID, keySecret)
			apiKey := database.APIKey{
				ID:              keyID,
				LifetimeSeconds: 86400, // default in db
				HashedSecret:    keySecretHashed,
				IPAddress: pqtype.Inet{
					IPNet: net.IPNet{
						IP:   net.IPv4(127, 0, 0, 1),
						Mask: net.IPv4Mask(255, 255, 255, 255),
					},
					Valid: true,
				},
				UserID:    user.ID,
				LastUsed:  now,
				ExpiresAt: now.Add(time.Hour),
				CreatedAt: now,
				UpdatedAt: now,
				LoginType: database.LoginTypePassword,
				Scopes:    []database.APIKeyScope{database.ApiKeyScopeCoderAll},
				TokenName: "",
			}
			if tc.key == "" {
				tc.key = token
			}

			// Define any case-specific mocks.
			if tc.mocksFn != nil {
				tc.mocksFn(db, apiKey, user)
			}

			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)
			require.NotNil(t, srv)

			resp, err := srv.IsAuthorized(t.Context(), &proto.IsAuthorizedRequest{Key: tc.key})
			if tc.expectedErr != nil {
				require.Error(t, err)
				require.ErrorIs(t, err, tc.expectedErr)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
			} else {
				expected := proto.IsAuthorizedResponse{
					OwnerId:            user.ID.String(),
					ApiKeyId:           keyID,
					Username:           user.Username,
					Email:              user.Email,
					OrganizationId:     everyoneOrgID.String(),
					OrganizationSource: string(database.AibridgeInterceptionOrganizationSourceEffectiveGroup),
				}
				require.NoError(t, err)
				require.Equal(t, &expected, resp)
			}
		})
	}
}

// When IsAuthorizedRequest carries KeyId instead of Key, the server skips
// the secret check and validates only that the key exists, is unexpired, and
// belongs to an active, non-deleted, non-system user. This is the path used by
// in-process delegated callers (e.g., chatd) that hold only the key ID.
func TestAuthorization_Delegated(t *testing.T) {
	t.Parallel()

	// A valid delegated request resolves to the Everyone group of this
	// organization, as delegated requests carry no workspace.
	everyoneOrgID := uuid.New()

	cases := []struct {
		name        string
		mocksFn     func(db *dbmock.MockStore, apiKey database.APIKey, user database.User)
		bothFields  bool
		expectedErr error
	}{
		{
			name: "valid",
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
				expectSiteWideEveryoneBudgetFallback(db, user.ID, everyoneOrgID)
			},
		},
		{
			name:        "unknown key",
			expectedErr: aibridgedserver.ErrUnknownKey,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, _ database.User) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(database.APIKey{}, sql.ErrNoRows)
			},
		},
		{
			name:        "expired",
			expectedErr: aibridgedserver.ErrExpired,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, _ database.User) {
				apiKey.ExpiresAt = dbtime.Now().Add(-time.Hour)
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
			},
		},
		{
			// Sending both Key and KeyId is an API misuse and must be
			// rejected to avoid ambiguity about which path was taken.
			name:        "both fields set",
			bothFields:  true,
			expectedErr: aibridgedserver.ErrAmbiguousAuth,
		},
		{
			// A bogus secret has no effect on the delegated path because
			// the secret is never checked. This is the load-bearing
			// security property: trust is established out-of-band, not in
			// this RPC.
			name: "secret hash mismatch is ignored",
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				apiKey.HashedSecret = []byte("not-the-real-hash")
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
				expectSiteWideEveryoneBudgetFallback(db, user.ID, everyoneOrgID)
			},
		},
		{
			// The delegated path must still reject keys whose owner has
			// been deleted; trust at the transport boundary does not
			// extend to bypassing user-status checks.
			name:        "deleted user",
			expectedErr: aibridgedserver.ErrDeletedUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Deleted = true
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			// The delegated path must reject inactive users; transport
			// trust does not override account suspension.
			name:        "suspended user",
			expectedErr: aibridgedserver.ErrInactiveUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Status = database.UserStatusSuspended
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			// Dormant users are inactive unless they are explicitly
			// reactivated through the HTTP middleware path.
			name:        "dormant user",
			expectedErr: aibridgedserver.ErrInactiveUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.Status = database.UserStatusDormant
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
		{
			// Likewise, a system user must never be authenticated through
			// the delegated path.
			name:        "system user",
			expectedErr: aibridgedserver.ErrSystemUser,
			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
				user.IsSystem = true
				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			logger := testutil.Logger(t)

			now := dbtime.Now()
			user := database.User{
				ID:         uuid.New(),
				Email:      "test@coder.com",
				Username:   "test",
				Name:       "Test User",
				CreatedAt:  now,
				UpdatedAt:  now,
				RBACRoles:  []string{},
				LoginType:  database.LoginTypePassword,
				Status:     database.UserStatusActive,
				LastSeenAt: now,
			}
			keyID, _ := cryptorand.String(10)
			_, keySecretHashed, _ := apikey.GenerateSecret(22)
			apiKey := database.APIKey{
				ID:              keyID,
				LifetimeSeconds: 86400,
				HashedSecret:    keySecretHashed,
				UserID:          user.ID,
				LastUsed:        now,
				ExpiresAt:       now.Add(time.Hour),
				CreatedAt:       now,
				UpdatedAt:       now,
				LoginType:       database.LoginTypePassword,
				Scopes:          []database.APIKeyScope{database.ApiKeyScopeCoderAll},
			}

			if tc.mocksFn != nil {
				tc.mocksFn(db, apiKey, user)
			}

			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)
			require.NotNil(t, srv)

			req := &proto.IsAuthorizedRequest{KeyId: keyID}
			if tc.bothFields {
				req.Key = "anything-anything"
			}

			resp, err := srv.IsAuthorized(t.Context(), req)
			if tc.expectedErr != nil {
				require.Error(t, err)
				require.ErrorIs(t, err, tc.expectedErr)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
				return
			}
			require.NoError(t, err)
			require.Equal(t, &proto.IsAuthorizedResponse{
				OwnerId:            user.ID.String(),
				ApiKeyId:           keyID,
				Username:           user.Username,
				Email:              user.Email,
				OrganizationId:     everyoneOrgID.String(),
				OrganizationSource: string(database.AibridgeInterceptionOrganizationSourceEffectiveGroup),
			}, resp)
		})
	}
}

func TestAuthorization_WorkspaceAttribution(t *testing.T) {
	t.Parallel()

	// workspaceAttribution is lookup-free: it parses the workspace UUID from the
	// strict server-minted token name and validates the embedded owner against
	// key.UserID without touching the database. Only an attributed workspace
	// is looked up, to resolve its organization.
	tests := []struct {
		name      string
		configure func(key *database.APIKey, user database.User, workspaceID uuid.UUID)
		wantErr   error
		wantWsID  bool // expect resp.GetWorkspaceId() == workspaceID.String()
	}{
		{
			name: "valid",
			configure: func(key *database.APIKey, user database.User, workspaceID uuid.UUID) {
				key.TokenName = fmt.Sprintf("%s_%s_session_token", user.ID, workspaceID)
			},
			wantWsID: true,
		},
		{
			// LoginType == Token must never yield workspace attribution even when
			// the token name matches the strict pattern.
			name: "personal token spoof",
			configure: func(key *database.APIKey, user database.User, workspaceID uuid.UUID) {
				key.LoginType = database.LoginTypeToken
				key.TokenName = fmt.Sprintf("%s_%s_session_token", user.ID, workspaceID)
			},
		},
		{
			// A token name that matches the suffix but not the strict UUID-pair
			// prefix must not be attributed to any workspace.
			name: "chatd lookalike",
			configure: func(key *database.APIKey, user database.User, _ uuid.UUID) {
				key.TokenName = fmt.Sprintf("chatd_%s_session_token", user.ID)
			},
		},
		{
			// An extra word between the UUIDs does not match the strict pattern.
			name: "oauth lookalike",
			configure: func(key *database.APIKey, user database.User, workspaceID uuid.UUID) {
				key.TokenName = fmt.Sprintf("%s_%s_oauth_session_token", user.ID, workspaceID)
			},
		},
		{
			// The embedded owner UUID in the token name differs from key.UserID.
			// Fail closed: return ErrWorkspaceAttribution.
			name: "embedded owner mismatch",
			configure: func(key *database.APIKey, _ database.User, workspaceID uuid.UUID) {
				key.TokenName = fmt.Sprintf("%s_%s_session_token", uuid.New(), workspaceID)
			},
			wantErr: aibridgedserver.ErrWorkspaceAttribution,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			now := dbtime.Now()
			user := database.User{ID: uuid.New(), Username: "test", Status: database.UserStatusActive, LoginType: database.LoginTypePassword}
			workspaceID := uuid.New()
			keyID, err := cryptorand.String(10)
			require.NoError(t, err)
			secret, hashedSecret, err := apikey.GenerateSecret(22)
			require.NoError(t, err)
			key := database.APIKey{ID: keyID, UserID: user.ID, HashedSecret: hashedSecret, ExpiresAt: now.Add(time.Hour), LoginType: database.LoginTypePassword}
			tt.configure(&key, user, workspaceID)

			db.EXPECT().GetAPIKeyByID(gomock.Any(), key.ID).Return(key, nil)
			db.EXPECT().GetUserByID(gomock.Any(), user.ID).Return(user, nil)
			// Attributed requests resolve the workspace's organization.
			// Unattributed requests make no workspace lookups and resolve the
			// organization of the site-wide effective group instead.
			orgID := uuid.New()
			if tt.wantWsID {
				db.EXPECT().GetWorkspaceByID(gomock.Any(), workspaceID).Return(database.Workspace{ID: workspaceID, OrganizationID: orgID}, nil)
				db.EXPECT().GetOrganizationsByUserID(gomock.Any(), gomock.Any()).Return([]database.Organization{{ID: orgID}}, nil)
			} else if tt.wantErr == nil {
				expectSiteWideEveryoneBudgetFallback(db, user.ID, orgID)
			}

			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store: db, AISeatTracker: agplaiseats.Noop{}, AccessURL: "/",
				GatewayCfg: codersdk.AIBridgeConfig{}, Experiments: requiredExperiments,
				Logger: testutil.Logger(t), Clock: quartz.NewReal(),
			})
			require.NoError(t, err)

			resp, err := srv.IsAuthorized(t.Context(), &proto.IsAuthorizedRequest{Key: key.ID + "-" + secret})
			if tt.wantErr != nil {
				require.ErrorIs(t, err, tt.wantErr)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
				return
			}
			require.NoError(t, err)
			require.Equal(t, orgID.String(), resp.GetOrganizationId())
			if tt.wantWsID {
				require.Equal(t, workspaceID.String(), resp.GetWorkspaceId())
				require.Equal(t, string(database.AibridgeInterceptionOrganizationSourceWorkspace), resp.GetOrganizationSource())
			} else {
				require.Empty(t, resp.GetWorkspaceId())
				require.Equal(t, string(database.AibridgeInterceptionOrganizationSourceEffectiveGroup), resp.GetOrganizationSource())
			}
		})
	}
}

// TestAuthorization_InternalErrors verifies which IsAuthorized errors are
// marked as internal. A failure of coderd, such as a database error, carries
// proto.ErrCodeInternal so the gateway answers 500. A rejected request carries
// no code, so the gateway answers 403.
func TestAuthorization_InternalErrors(t *testing.T) {
	t.Parallel()

	errDB := xerrors.New("connection refused")

	// identity is the request's API key and user, and the workspace or
	// organization it names.
	type identity struct {
		key         database.APIKey
		user        database.User
		workspaceID uuid.UUID
		orgID       uuid.UUID
	}
	// identityFound sets up a key and user that pass every check before
	// organization resolution.
	identityFound := func(db *dbmock.MockStore, id identity) {
		db.EXPECT().GetAPIKeyByID(gomock.Any(), id.key.ID).Return(id.key, nil)
		db.EXPECT().GetUserByID(gomock.Any(), id.user.ID).Return(id.user, nil)
	}

	tests := []struct {
		name string
		// workspace makes the key a workspace session token.
		workspace bool
		// delegated authenticates by key ID, naming an organization, as chatd
		// does.
		delegated    bool
		configure    func(db *dbmock.MockStore, id identity)
		wantErr      error
		wantInternal bool
	}{
		{
			name: "key lookup fails",
			configure: func(db *dbmock.MockStore, id identity) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), id.key.ID).Return(database.APIKey{}, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name: "user lookup fails",
			configure: func(db *dbmock.MockStore, id identity) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), id.key.ID).Return(id.key, nil)
				db.EXPECT().GetUserByID(gomock.Any(), id.user.ID).Return(database.User{}, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name:      "workspace lookup fails",
			workspace: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetWorkspaceByID(gomock.Any(), id.workspaceID).Return(database.Workspace{}, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name:      "workspace organization lookup fails",
			workspace: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetWorkspaceByID(gomock.Any(), id.workspaceID).Return(database.Workspace{ID: id.workspaceID, OrganizationID: id.orgID}, nil)
				db.EXPECT().GetOrganizationsByUserID(gomock.Any(), gomock.Any()).Return(nil, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name:      "delegated organization lookup fails",
			delegated: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetOrganizationsByUserID(gomock.Any(), gomock.Any()).Return(nil, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name: "effective group lookup fails",
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(id.user.ID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, errDB)
			},
			wantErr:      errDB,
			wantInternal: true,
		},
		{
			name: "unknown key",
			configure: func(db *dbmock.MockStore, id identity) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), id.key.ID).Return(database.APIKey{}, sql.ErrNoRows)
			},
			wantErr: aibridgedserver.ErrUnknownKey,
		},
		{
			name: "unknown user",
			configure: func(db *dbmock.MockStore, id identity) {
				db.EXPECT().GetAPIKeyByID(gomock.Any(), id.key.ID).Return(id.key, nil)
				db.EXPECT().GetUserByID(gomock.Any(), id.user.ID).Return(database.User{}, sql.ErrNoRows)
			},
			wantErr: aibridgedserver.ErrUnknownUser,
		},
		{
			name:      "unknown workspace",
			workspace: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetWorkspaceByID(gomock.Any(), id.workspaceID).Return(database.Workspace{}, sql.ErrNoRows)
			},
			wantErr: aibridgedserver.ErrUnknownWorkspace,
		},
		{
			name:      "user not a member of workspace organization",
			workspace: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetWorkspaceByID(gomock.Any(), id.workspaceID).Return(database.Workspace{ID: id.workspaceID, OrganizationID: id.orgID}, nil)
				db.EXPECT().GetOrganizationsByUserID(gomock.Any(), gomock.Any()).Return([]database.Organization{{ID: uuid.New()}}, nil)
			},
			wantErr: aibridgedserver.ErrUnknownOrganizationMember,
		},
		{
			name:      "user not a member of delegated organization",
			delegated: true,
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetOrganizationsByUserID(gomock.Any(), gomock.Any()).Return([]database.Organization{{ID: uuid.New()}}, nil)
			},
			wantErr: aibridgedserver.ErrUnknownOrganizationMember,
		},
		{
			name: "no organization",
			configure: func(db *dbmock.MockStore, id identity) {
				identityFound(db, id)
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(id.user.ID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(id.user.ID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
				db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(id.user.ID)).Return(uuid.Nil, sql.ErrNoRows)
			},
			wantErr: aibridgedserver.ErrNoOrganization,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			id := identity{
				user:        database.User{ID: uuid.New(), Username: "test", Status: database.UserStatusActive, LoginType: database.LoginTypePassword},
				workspaceID: uuid.New(),
				orgID:       uuid.New(),
			}
			keyID, err := cryptorand.String(10)
			require.NoError(t, err)
			secret, hashedSecret, err := apikey.GenerateSecret(22)
			require.NoError(t, err)
			id.key = database.APIKey{ID: keyID, UserID: id.user.ID, HashedSecret: hashedSecret, ExpiresAt: dbtime.Now().Add(time.Hour), LoginType: database.LoginTypePassword}
			if tt.workspace {
				id.key.TokenName = fmt.Sprintf("%s_%s_session_token", id.user.ID, id.workspaceID)
			}
			tt.configure(db, id)

			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store: db, AISeatTracker: agplaiseats.Noop{}, AccessURL: "/",
				GatewayCfg: codersdk.AIBridgeConfig{}, Experiments: requiredExperiments,
				Logger: testutil.Logger(t), Clock: quartz.NewReal(),
			})
			require.NoError(t, err)

			req := &proto.IsAuthorizedRequest{Key: id.key.ID + "-" + secret}
			if tt.delegated {
				req = &proto.IsAuthorizedRequest{KeyId: id.key.ID, DelegatedOrganizationId: id.orgID.String()}
			}
			resp, err := srv.IsAuthorized(t.Context(), req)
			require.Nil(t, resp)
			require.ErrorIs(t, err, tt.wantErr)
			if tt.wantInternal {
				require.Equal(t, proto.ErrCodeInternal, drpcerr.Code(err))
			} else {
				require.Zero(t, drpcerr.Code(err))
			}
		})
	}
}

// TestAuthorization_WorkspaceOrgResolution verifies that requests made with a
// workspace session token resolve to the workspace's organization, and are
// rejected when the workspace or the user's membership in its organization no
// longer exists. It runs against a real database behind dbauthz so the AI
// Gateway subject's permissions are exercised.
func TestAuthorization_WorkspaceOrgResolution(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name string
		// configure mutates state after the workspace and its session token
		// exist. It returns the token to authorize with.
		configure     func(t *testing.T, db database.Store, user database.User, org database.Organization, workspace database.WorkspaceTable) string
		wantErr       error
		wantWorkspace bool
		wantSource    database.AIBridgeInterceptionOrganizationSource
	}{
		{
			name: "workspace session token resolves workspace org",
			configure: func(t *testing.T, db database.Store, user database.User, _ database.Organization, workspace database.WorkspaceTable) string {
				_, token := dbgen.APIKey(t, db, database.APIKey{
					UserID:    user.ID,
					TokenName: provisionerdserver.WorkspaceSessionTokenName(user.ID, workspace.ID),
				})
				return token
			},
			wantWorkspace: true,
			wantSource:    database.AibridgeInterceptionOrganizationSourceWorkspace,
		},
		{
			// Tokens that are not workspace session tokens fall back to the
			// organization of the site-wide effective group: here the Everyone
			// group of the user's only organization.
			name: "non-workspace token resolves effective group org",
			configure: func(t *testing.T, db database.Store, user database.User, _ database.Organization, _ database.WorkspaceTable) string {
				_, token := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID})
				return token
			},
			wantSource: database.AibridgeInterceptionOrganizationSourceEffectiveGroup,
		},
		{
			name: "missing workspace",
			configure: func(t *testing.T, db database.Store, user database.User, _ database.Organization, _ database.WorkspaceTable) string {
				_, token := dbgen.APIKey(t, db, database.APIKey{
					UserID:    user.ID,
					TokenName: provisionerdserver.WorkspaceSessionTokenName(user.ID, uuid.New()),
				})
				return token
			},
			wantErr: aibridgedserver.ErrUnknownWorkspace,
		},
		{
			name: "deleted workspace",
			configure: func(t *testing.T, db database.Store, user database.User, _ database.Organization, workspace database.WorkspaceTable) string {
				_, token := dbgen.APIKey(t, db, database.APIKey{
					UserID:    user.ID,
					TokenName: provisionerdserver.WorkspaceSessionTokenName(user.ID, workspace.ID),
				})
				require.NoError(t, db.UpdateWorkspaceDeletedByID(t.Context(), database.UpdateWorkspaceDeletedByIDParams{
					ID:      workspace.ID,
					Deleted: true,
				}))
				return token
			},
			wantErr: aibridgedserver.ErrUnknownWorkspace,
		},
		{
			name: "user no longer a member of workspace org",
			configure: func(t *testing.T, db database.Store, user database.User, org database.Organization, workspace database.WorkspaceTable) string {
				_, token := dbgen.APIKey(t, db, database.APIKey{
					UserID:    user.ID,
					TokenName: provisionerdserver.WorkspaceSessionTokenName(user.ID, workspace.ID),
				})
				require.NoError(t, db.DeleteOrganizationMember(t.Context(), database.DeleteOrganizationMemberParams{
					OrganizationID: org.ID,
					UserID:         user.ID,
				}))
				return token
			},
			wantErr: aibridgedserver.ErrUnknownOrganizationMember,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctx := testutil.Context(t, testutil.WaitLong)
			logger := testutil.Logger(t)

			rawDB, _ := dbtestutil.NewDB(t)
			authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

			org := dbgen.Organization(t, rawDB, database.Organization{})
			user := dbgen.User(t, rawDB, database.User{})
			dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
			workspace := dbfake.WorkspaceBuild(t, rawDB, database.WorkspaceTable{
				OrganizationID: org.ID,
				OwnerID:        user.ID,
			}).Do().Workspace
			token := tt.configure(t, rawDB, user, org, workspace)

			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         authzDB,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			resp, err := srv.IsAuthorized(ctx, &proto.IsAuthorizedRequest{Key: token})
			if tt.wantErr != nil {
				require.ErrorIs(t, err, tt.wantErr)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
				return
			}
			require.NoError(t, err)
			require.Equal(t, user.ID.String(), resp.GetOwnerId())
			require.Equal(t, org.ID.String(), resp.GetOrganizationId())
			require.Equal(t, string(tt.wantSource), resp.GetOrganizationSource())
			if tt.wantWorkspace {
				require.Equal(t, workspace.ID.String(), resp.GetWorkspaceId())
			} else {
				require.Empty(t, resp.GetWorkspaceId())
			}
		})
	}
}

// TestAuthorization_DelegatedOrgResolution verifies that a delegated chatd
// request resolves to the chat's organization when the user is still a live
// member of it, and that the delegated organization is rejected anywhere else.
func TestAuthorization_DelegatedOrgResolution(t *testing.T) {
	t.Parallel()

	tests := []struct {
		name string
		// configure mutates state after the user joins org. It returns the
		// request to authorize with.
		configure func(t *testing.T, db database.Store, org database.Organization, key database.APIKey, token string) *proto.IsAuthorizedRequest
		wantErr   error
	}{
		{
			name: "delegated request resolves chat org",
			configure: func(_ *testing.T, _ database.Store, org database.Organization, key database.APIKey, _ string) *proto.IsAuthorizedRequest {
				return &proto.IsAuthorizedRequest{KeyId: key.ID, DelegatedOrganizationId: org.ID.String()}
			},
		},
		{
			name: "user no longer a member of chat org",
			configure: func(t *testing.T, db database.Store, org database.Organization, key database.APIKey, _ string) *proto.IsAuthorizedRequest {
				require.NoError(t, db.DeleteOrganizationMember(t.Context(), database.DeleteOrganizationMemberParams{
					OrganizationID: org.ID,
					UserID:         key.UserID,
				}))
				return &proto.IsAuthorizedRequest{KeyId: key.ID, DelegatedOrganizationId: org.ID.String()}
			},
			wantErr: aibridgedserver.ErrUnknownOrganizationMember,
		},
		{
			name: "deleted chat org",
			configure: func(t *testing.T, db database.Store, org database.Organization, key database.APIKey, _ string) *proto.IsAuthorizedRequest {
				require.NoError(t, db.UpdateOrganizationDeletedByID(t.Context(), database.UpdateOrganizationDeletedByIDParams{
					ID:        org.ID,
					UpdatedAt: dbtime.Now(),
				}))
				return &proto.IsAuthorizedRequest{KeyId: key.ID, DelegatedOrganizationId: org.ID.String()}
			},
			wantErr: aibridgedserver.ErrUnknownOrganizationMember,
		},
		{
			// Only in-process delegated callers can set the delegated
			// organization.
			name: "delegated org with key",
			configure: func(_ *testing.T, _ database.Store, org database.Organization, _ database.APIKey, token string) *proto.IsAuthorizedRequest {
				return &proto.IsAuthorizedRequest{Key: token, DelegatedOrganizationId: org.ID.String()}
			},
			wantErr: aibridgedserver.ErrAmbiguousAuth,
		},
		{
			name: "malformed delegated org",
			configure: func(_ *testing.T, _ database.Store, _ database.Organization, key database.APIKey, _ string) *proto.IsAuthorizedRequest {
				return &proto.IsAuthorizedRequest{KeyId: key.ID, DelegatedOrganizationId: "not-a-uuid"}
			},
			wantErr: aibridgedserver.ErrInvalidDelegatedOrganization,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctx := testutil.Context(t, testutil.WaitLong)
			logger := testutil.Logger(t)

			rawDB, _ := dbtestutil.NewDB(t)
			authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

			org := dbgen.Organization(t, rawDB, database.Organization{})
			user := dbgen.User(t, rawDB, database.User{})
			dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
			key, token := dbgen.APIKey(t, rawDB, database.APIKey{UserID: user.ID})
			req := tt.configure(t, rawDB, org, key, token)

			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         authzDB,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			resp, err := srv.IsAuthorized(ctx, req)
			if tt.wantErr != nil {
				require.ErrorIs(t, err, tt.wantErr)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
				require.Nil(t, resp)
				return
			}
			require.NoError(t, err)
			require.Equal(t, user.ID.String(), resp.GetOwnerId())
			require.Equal(t, org.ID.String(), resp.GetOrganizationId())
			require.Equal(t, string(database.AibridgeInterceptionOrganizationSourceChatd), resp.GetOrganizationSource())
		})
	}
}

// TestAuthorization_EffectiveGroupOrgResolution verifies that a request
// without a workspace resolves to the organization of the user's site-wide
// effective group, on both the direct and the delegated path. It runs against
// a real database behind dbauthz so the AI Gateway subject's permissions are
// exercised.
func TestAuthorization_EffectiveGroupOrgResolution(t *testing.T) {
	t.Parallel()

	// fixture is the state each case starts from: a user who is a member of
	// orgs A and C, with a budgeted group in A. Org B exists, but the user is
	// not a member.
	type fixture struct {
		user             database.User
		orgA, orgB, orgC database.Organization
		groupA           database.Group
	}

	// removeAllMemberships removes the user from every organization they
	// belong to, so no organization can resolve.
	removeAllMemberships := func(ctx context.Context, t *testing.T, db database.Store, f fixture) uuid.UUID {
		for _, org := range []database.Organization{f.orgA, f.orgC} {
			require.NoError(t, db.DeleteOrganizationMember(ctx, database.DeleteOrganizationMemberParams{
				OrganizationID: org.ID,
				UserID:         f.user.ID,
			}))
		}
		return uuid.Nil
	}

	tests := []struct {
		name      string
		delegated bool
		// configure adjusts the fixture and returns the organization the
		// request should resolve to, or uuid.Nil when the request must be
		// rejected with ErrNoOrganization.
		configure func(ctx context.Context, t *testing.T, db database.Store, f fixture) uuid.UUID
	}{
		{
			name: "budgeted group",
			configure: func(_ context.Context, _ *testing.T, _ database.Store, f fixture) uuid.UUID {
				return f.orgA.ID
			},
		},
		{
			// An override wins over a group budget, even a higher one.
			name: "override",
			configure: func(ctx context.Context, t *testing.T, db database.Store, f fixture) uuid.UUID {
				groupC := dbgen.Group(t, db, database.Group{OrganizationID: f.orgC.ID})
				dbgen.GroupMember(t, db, database.GroupMemberTable{UserID: f.user.ID, GroupID: groupC.ID})
				_, err := db.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
					UserID:           f.user.ID,
					GroupID:          groupC.ID,
					SpendLimitMicros: 1,
				})
				require.NoError(t, err)
				return f.orgC.ID
			},
		},
		{
			// With no budget anywhere, the Everyone group of the default org
			// is preferred over earlier memberships.
			name: "no budgets uses default org",
			configure: func(ctx context.Context, t *testing.T, db database.Store, f fixture) uuid.UUID {
				_, err := db.DeleteGroupAIBudget(ctx, f.groupA.ID)
				require.NoError(t, err)
				defaultOrg, err := db.GetDefaultOrganization(ctx)
				require.NoError(t, err)
				dbgen.OrganizationMember(t, db, database.OrganizationMember{OrganizationID: defaultOrg.ID, UserID: f.user.ID})
				return defaultOrg.ID
			},
		},
		{
			// A user in no live organization resolves no organization, so the
			// request is rejected rather than falling back to site-wide budget
			// and cost.
			name:      "no membership",
			configure: removeAllMemberships,
		},
		{
			name:      "delegated no membership",
			delegated: true,
			configure: removeAllMemberships,
		},
		{
			// Delegated requests carry no workspace, so they take the same
			// site-wide step.
			name:      "delegated",
			delegated: true,
			configure: func(_ context.Context, _ *testing.T, _ database.Store, f fixture) uuid.UUID {
				return f.orgA.ID
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctx := testutil.Context(t, testutil.WaitLong)
			logger := testutil.Logger(t)

			rawDB, _ := dbtestutil.NewDB(t)
			authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

			f := fixture{
				user: dbgen.User(t, rawDB, database.User{}),
				orgA: dbgen.Organization(t, rawDB, database.Organization{}),
				orgB: dbgen.Organization(t, rawDB, database.Organization{}),
				orgC: dbgen.Organization(t, rawDB, database.Organization{}),
			}
			for _, org := range []database.Organization{f.orgA, f.orgC} {
				dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: f.user.ID})
			}
			f.groupA = dbgen.Group(t, rawDB, database.Group{OrganizationID: f.orgA.ID})
			dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: f.user.ID, GroupID: f.groupA.ID})
			_, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
				GroupID:          f.groupA.ID,
				SpendLimitMicros: 1_000_000,
			})
			require.NoError(t, err)
			// A budget in org B must never count, as the user is not a member.
			groupB := dbgen.Group(t, rawDB, database.Group{OrganizationID: f.orgB.ID})
			_, err = rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
				GroupID:          groupB.ID,
				SpendLimitMicros: 1_000_000_000,
			})
			require.NoError(t, err)

			wantOrgID := tt.configure(ctx, t, rawDB, f)

			key, token := dbgen.APIKey(t, rawDB, database.APIKey{UserID: f.user.ID})
			req := &proto.IsAuthorizedRequest{Key: token}
			if tt.delegated {
				req = &proto.IsAuthorizedRequest{KeyId: key.ID}
			}

			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         authzDB,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			resp, err := srv.IsAuthorized(ctx, req)
			if wantOrgID == uuid.Nil {
				require.ErrorIs(t, err, aibridgedserver.ErrNoOrganization)
				require.Zero(t, drpcerr.Code(err), "a rejection must not be marked internal")
				require.Nil(t, resp)
				return
			}
			require.NoError(t, err)
			require.Equal(t, f.user.ID.String(), resp.GetOwnerId())
			require.Empty(t, resp.GetWorkspaceId())
			require.Equal(t, wantOrgID.String(), resp.GetOrganizationId())
			require.Equal(t, string(database.AibridgeInterceptionOrganizationSourceEffectiveGroup), resp.GetOrganizationSource())
		})
	}
}

func TestIsBudgetExceeded(t *testing.T) {
	t.Parallel()

	cases := []struct {
		name              string
		userIDStr         string
		setupMocks        func(db *dbmock.MockStore, userID uuid.UUID) (resp *proto.IsBudgetExceededResponse, blockedGroupID uuid.UUID)
		wantErrContains   string
		wantMetricOutcome string
	}{
		{
			// Invalid UUID short-circuits before any store call.
			name:              "invalid user_id",
			userIDStr:         "not-a-uuid",
			wantErrContains:   "invalid user_id",
			wantMetricOutcome: "error",
		},
		{
			// No override and no group budget resolves: pass-through.
			name: "no budget configured returns not exceeded",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         false,
					SpendLimitMicros: nil,
				}, uuid.Nil
			},
			wantMetricOutcome: "allowed",
		},
		{
			// Group budget resolves, spend below limit (spend 500 < limit 1000): pass-through.
			name: "under limit returns not exceeded",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				groupID := uuid.New()
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{SpendMicros: 500}, nil)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         false,
					SpendLimitMicros: new(int64(1_000)),
				}, uuid.Nil
			},
			wantMetricOutcome: "allowed",
		},
		{
			// Group budget resolves, spend at limit (spend 1000 == limit 1000): blocked.
			name: "at limit returns exceeded",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				groupID := uuid.New()
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{SpendMicros: 1_000}, nil)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         true,
					SpendLimitMicros: new(int64(1_000)),
				}, groupID
			},
			wantMetricOutcome: "blocked",
		},
		{
			// Limit of 0 is a valid "block-all" setting, distinct from
			// "no budget configured": blocked.
			name: "zero limit blocks all requests",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				groupID := uuid.New()
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{SpendMicros: 0}, nil)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         true,
					SpendLimitMicros: new(int64(0)),
				}, groupID
			},
			wantMetricOutcome: "blocked",
		},
		{
			// Group budget resolves, spend above limit (spend 1500 > limit 1000): blocked.
			name: "over limit returns exceeded",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				groupID := uuid.New()
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{SpendMicros: 1_500}, nil)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         true,
					SpendLimitMicros: new(int64(1_000)),
				}, groupID
			},
			wantMetricOutcome: "blocked",
		},
		{
			// User override wins, group lookup skipped, spend aggregated against
			// the override's group (spend 600 > limit 500): blocked.
			name: "user override wins over group budget",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				overrideGroupID := uuid.New()
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{
						UserID:           userID,
						GroupID:          overrideGroupID,
						SpendLimitMicros: 500,
					}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Cond(func(p database.GetUserAISpendSinceParams) bool {
					return assert.Equal(t, overrideGroupID, p.EffectiveGroupID, "spend aggregated against override group")
				})).Return(database.GetUserAISpendSinceRow{SpendMicros: 600}, nil)
				return &proto.IsBudgetExceededResponse{
					Exceeded:         true,
					SpendLimitMicros: new(int64(500)),
				}, overrideGroupID
			},
			wantMetricOutcome: "blocked",
		},
		{
			// Unexpected error from budget override lookup propagates.
			name: "budget resolution error propagates",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrConnDone)
				return nil, uuid.Nil
			},
			wantErrContains:   "resolve effective AI budget",
			wantMetricOutcome: "error",
		},
		{
			// Error from spend aggregation propagates (fail-closed).
			name: "spend aggregation error propagates",
			setupMocks: func(db *dbmock.MockStore, userID uuid.UUID) (*proto.IsBudgetExceededResponse, uuid.UUID) {
				db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
					Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
				db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
					Return(database.GetHighestGroupAIBudgetByUserRow{GroupID: uuid.New(), SpendLimitMicros: 1_000}, nil)
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{}, sql.ErrConnDone)
				return nil, uuid.Nil
			},
			wantErrContains:   "get user AI spend",
			wantMetricOutcome: "error",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			logger := testutil.Logger(t)

			userID := uuid.New()
			userIDStr := tc.userIDStr
			if userIDStr == "" {
				userIDStr = userID.String()
			}

			var wantResp *proto.IsBudgetExceededResponse
			var blockedGroupID uuid.UUID
			if tc.setupMocks != nil {
				wantResp, blockedGroupID = tc.setupMocks(db, userID)
			}

			reg := prometheus.NewRegistry()
			metrics := aibridgedserver.NewMetrics(reg)
			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
				Metrics:       metrics,
			})
			require.NoError(t, err)

			req := &proto.IsBudgetExceededRequest{UserId: userIDStr}
			resp, err := srv.IsBudgetExceeded(t.Context(), req)

			// The enforcement duration is always observed once, labeled by the
			// outcome, even when the check errors.
			require.Equal(t, 1, promtest.CollectAndCount(metrics.EnforcementDuration))
			require.EqualValues(t, 1, promhelp.HistogramValue(t, reg,
				"cost_control_enforcement_duration_seconds",
				prometheus.Labels{"outcome": tc.wantMetricOutcome}).GetSampleCount())
			wantBlocked := 0
			if tc.wantMetricOutcome == "blocked" {
				wantBlocked = 1
			}
			require.Equal(t, wantBlocked, promtest.CollectAndCount(metrics.BlockedRequests))
			if wantBlocked == 1 {
				require.Equal(t, 1, promhelp.CounterValue(t, reg,
					"cost_control_blocked_requests_total",
					prometheus.Labels{"group_id": blockedGroupID.String()}))
			}

			if tc.wantErrContains != "" {
				require.Error(t, err)
				require.Nil(t, resp)
				assert.ErrorContains(t, err, tc.wantErrContains)
				return
			}
			require.NoError(t, err)
			require.NotNil(t, resp)
			require.Equal(t, wantResp.GetExceeded(), resp.GetExceeded(), "exceeded")
			require.Equal(t, wantResp.SpendLimitMicros, resp.SpendLimitMicros, "spend_limit_micros")
		})
	}
}

// TestIsBudgetExceeded_Enforcement exercises real-DB scenarios that drive
// enforcement decisions.
func TestIsBudgetExceeded_Enforcement(t *testing.T) {
	t.Parallel()

	const groupLimitMicros = 1_000_000

	// setup provisions a user in an organization with a single budgeted group.
	setup := func(t *testing.T, clock quartz.Clock) (context.Context, database.Store, *aibridgedserver.Server, database.User, database.Group) {
		t.Helper()

		ctx := testutil.Context(t, testutil.WaitLong)
		logger := testutil.Logger(t)

		rawDB, _ := dbtestutil.NewDB(t)
		authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

		org := dbgen.Organization(t, rawDB, database.Organization{})
		user := dbgen.User(t, rawDB, database.User{})
		dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
		group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID})
		dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID})

		_, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
			GroupID:          group.ID,
			SpendLimitMicros: groupLimitMicros,
		})
		require.NoError(t, err)

		srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
			Store:         authzDB,
			AISeatTracker: agplaiseats.Noop{},
			AccessURL:     "/",
			GatewayCfg:    codersdk.AIBridgeConfig{},
			Experiments:   requiredExperiments,
			Logger:        logger,
			Clock:         clock,
		})
		require.NoError(t, err)

		return ctx, rawDB, srv, user, group
	}

	t.Run("period boundary excludes prior period spend", func(t *testing.T) {
		t.Parallel()

		// Use fixed dates to keep the test deterministic.
		clock := quartz.NewMock(t)
		ctx, rawDB, srv, user, group := setup(t, clock)

		prevMonth := time.Date(2026, time.January, 15, 0, 0, 0, 0, time.UTC)
		nextMonth := time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC)

		// Set now to 2026-01-15.
		clock.Set(prevMonth)

		// User spend on 2026-01-15 exceeds the group limit.
		_, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
			UserID:           user.ID,
			EffectiveGroupID: group.ID,
			Day:              prevMonth,
			CostMicros:       1_500_000,
		})
		require.NoError(t, err)

		// Current period is January: includes the 2026-01-15 spend, user exceeded.
		prevMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
			UserId: user.ID.String(),
		})
		require.NoError(t, err)
		require.True(t, prevMonthResp.GetExceeded())
		require.Equal(t, int64(groupLimitMicros), prevMonthResp.GetSpendLimitMicros())

		// Advance clock to 2026-02-05: excludes the 2026-01-15 spend, user not exceeded.
		clock.Set(nextMonth)
		nextMonthResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
			UserId: user.ID.String(),
		})
		require.NoError(t, err)
		require.False(t, nextMonthResp.GetExceeded())
		require.Equal(t, int64(groupLimitMicros), nextMonthResp.GetSpendLimitMicros())
	})

	t.Run("new user override unblocks user", func(t *testing.T) {
		t.Parallel()

		// Use fixed dates to keep the test deterministic.
		clock := quartz.NewMock(t)
		ctx, rawDB, srv, user, group := setup(t, clock)

		now := time.Date(2026, time.March, 15, 0, 0, 0, 0, time.UTC)

		// Set now to 2026-03-15.
		clock.Set(now)

		// User spend on 2026-03-15 exceeds the group limit.
		_, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
			UserID:           user.ID,
			EffectiveGroupID: group.ID,
			Day:              now,
			CostMicros:       1_500_000,
		})
		require.NoError(t, err)

		// User's spend exceeds the group limit.
		beforeResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
			UserId: user.ID.String(),
		})
		require.NoError(t, err)
		require.True(t, beforeResp.GetExceeded())
		require.Equal(t, int64(groupLimitMicros), beforeResp.GetSpendLimitMicros())

		// Add user override with a higher limit on the same group. The override
		// wins, so the user's spend is now under the effective limit.
		const overrideLimitMicros = 2_000_000
		_, err = rawDB.UpsertUserAIBudgetOverride(ctx, database.UpsertUserAIBudgetOverrideParams{
			UserID:           user.ID,
			GroupID:          group.ID,
			SpendLimitMicros: overrideLimitMicros,
		})
		require.NoError(t, err)

		afterResp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
			UserId: user.ID.String(),
		})
		require.NoError(t, err)
		require.False(t, afterResp.GetExceeded())
		require.Equal(t, int64(overrideLimitMicros), afterResp.GetSpendLimitMicros())
	})

	t.Run("unbudgeted member is not blocked", func(t *testing.T) {
		t.Parallel()

		clock := quartz.NewMock(t)
		clock.Set(time.Date(2026, time.March, 15, 0, 0, 0, 0, time.UTC))

		ctx := testutil.Context(t, testutil.WaitLong)
		logger := testutil.Logger(t)
		rawDB, _ := dbtestutil.NewDB(t)
		authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

		// An org member with no group budget and no override: spend is
		// unlimited, so enforcement never blocks them.
		org := dbgen.Organization(t, rawDB, database.Organization{})
		user := dbgen.User(t, rawDB, database.User{})
		dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})

		// Record spend attributed to the Everyone group. Without a configured
		// limit it must not cause a block.
		_, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
			UserID:           user.ID,
			EffectiveGroupID: org.ID,
			Day:              clock.Now(),
			CostMicros:       1_000_000_000, // $1,000 USD
		})
		require.NoError(t, err)

		srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
			Store:         authzDB,
			AISeatTracker: agplaiseats.Noop{},
			AccessURL:     "/",
			GatewayCfg:    codersdk.AIBridgeConfig{},
			Experiments:   requiredExperiments,
			Logger:        logger,
			Clock:         clock,
		})
		require.NoError(t, err)

		resp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{UserId: user.ID.String()})
		require.NoError(t, err)
		require.False(t, resp.GetExceeded())
		require.Nil(t, resp.SpendLimitMicros)
	})

	t.Run("organization scopes the budget", func(t *testing.T) {
		t.Parallel()

		clock := quartz.NewMock(t)
		clock.Set(time.Date(2026, time.March, 15, 0, 0, 0, 0, time.UTC))

		// The user is over budget in org A's group and under budget in org
		// B's group. A's group has the higher limit, so it also wins
		// site-wide.
		ctx, rawDB, srv, user, groupA := setup(t, clock)
		_, err := rawDB.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
			UserID:           user.ID,
			EffectiveGroupID: groupA.ID,
			Day:              clock.Now(),
			CostMicros:       1_500_000,
		})
		require.NoError(t, err)

		const groupBLimitMicros = groupLimitMicros / 2
		orgB := dbgen.Organization(t, rawDB, database.Organization{})
		dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: orgB.ID, UserID: user.ID})
		groupB := dbgen.Group(t, rawDB, database.Group{OrganizationID: orgB.ID})
		dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: groupB.ID})
		_, err = rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
			GroupID:          groupB.ID,
			SpendLimitMicros: groupBLimitMicros,
		})
		require.NoError(t, err)

		for _, tc := range []struct {
			name         string
			orgID        string
			wantExceeded bool
			wantLimit    int64
		}{
			{name: "org B", orgID: orgB.ID.String(), wantExceeded: false, wantLimit: groupBLimitMicros},
			{name: "org A", orgID: groupA.OrganizationID.String(), wantExceeded: true, wantLimit: groupLimitMicros},
			{name: "site-wide", orgID: "", wantExceeded: true, wantLimit: groupLimitMicros},
		} {
			resp, err := srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
				UserId:         user.ID.String(),
				OrganizationId: tc.orgID,
			})
			require.NoError(t, err, tc.name)
			require.Equal(t, tc.wantExceeded, resp.GetExceeded(), tc.name)
			require.Equal(t, tc.wantLimit, resp.GetSpendLimitMicros(), tc.name)
		}

		_, err = srv.IsBudgetExceeded(ctx, &proto.IsBudgetExceededRequest{
			UserId:         user.ID.String(),
			OrganizationId: "not-a-uuid",
		})
		require.ErrorContains(t, err, "invalid organization_id")
	})
}

func TestGetMCPServerConfigs(t *testing.T) {
	t.Parallel()

	externalAuthCfgs := []*externalauth.Config{
		{
			ID:     "1",
			MCPURL: "1.com/mcp",
		},
		{
			ID: "2", // Will not be eligible for inclusion since MCPURL is not defined.
		},
	}

	cases := []struct {
		name                     string
		disableCoderMCPInjection bool
		experiments              codersdk.Experiments
		oauth2ProviderEnabled    bool
		externalAuthConfigs      []*externalauth.Config
		expectCoderMCP           bool
		expectedExternalMCP      bool
	}{
		{
			name:        "MCP experiment off, OAuth2 provider off",
			experiments: codersdk.Experiments{},
		},
		{
			name:        "MCP experiment on, OAuth2 provider off",
			experiments: requiredExperiments,
		},
		{
			name:                  "OAuth2 provider on, MCP experiment off",
			experiments:           codersdk.Experiments{},
			oauth2ProviderEnabled: true,
		},
		{
			name:                  "only internal MCP",
			experiments:           requiredExperiments,
			oauth2ProviderEnabled: true,
			expectCoderMCP:        true,
		},
		{
			name:                "only external MCP",
			externalAuthConfigs: externalAuthCfgs,
			expectedExternalMCP: true,
		},
		{
			name:                  "both internal & external MCP",
			experiments:           requiredExperiments,
			oauth2ProviderEnabled: true,
			externalAuthConfigs:   externalAuthCfgs,
			expectCoderMCP:        true,
			expectedExternalMCP:   true,
		},
		{
			name:                     "both internal & external MCP, but coder MCP tools not injected",
			disableCoderMCPInjection: true,
			experiments:              requiredExperiments,
			oauth2ProviderEnabled:    true,
			externalAuthConfigs:      externalAuthCfgs,
			expectCoderMCP:           false,
			expectedExternalMCP:      true,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			logger := testutil.Logger(t)

			accessURL := "https://my-cool-deployment.com"
			srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     accessURL,
				GatewayCfg: codersdk.AIBridgeConfig{
					InjectCoderMCPTools: serpent.Bool(!tc.disableCoderMCPInjection),
				},
				ExternalAuthConfigs:   tc.externalAuthConfigs,
				Experiments:           tc.experiments,
				OAuth2ProviderEnabled: tc.oauth2ProviderEnabled,
				Logger:                logger,
				Clock:                 quartz.NewReal(),
			})
			require.NoError(t, err)
			require.NotNil(t, srv)

			resp, err := srv.GetMCPServerConfigs(t.Context(), &proto.GetMCPServerConfigsRequest{})
			require.NoError(t, err)
			require.NotNil(t, resp)

			if tc.expectCoderMCP {
				coderConfig := resp.CoderMcpConfig
				require.NotNil(t, coderConfig)
				require.Equal(t, aibridged.InternalMCPServerID, coderConfig.GetId())
				expectedURL, err := url.JoinPath(accessURL, codermcp.MCPEndpoint)
				require.NoError(t, err)
				require.Equal(t, expectedURL, coderConfig.GetUrl())
				require.Empty(t, coderConfig.GetToolAllowRegex())
				require.Empty(t, coderConfig.GetToolDenyRegex())
			} else {
				require.Empty(t, resp.GetCoderMcpConfig())
			}

			if tc.expectedExternalMCP {
				require.Len(t, resp.GetExternalAuthMcpConfigs(), 1)
			} else {
				require.Empty(t, resp.GetExternalAuthMcpConfigs())
			}
		})
	}
}

func TestGetMCPServerAccessTokensBatch(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	db := dbmock.NewMockStore(ctrl)
	logger := testutil.Logger(t)

	// Given: 2 external auth configured with MCP and 1 without.
	srv, err := aibridgedserver.NewServer(t.Context(), aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		ExternalAuthConfigs: []*externalauth.Config{
			{
				ID:     "1",
				MCPURL: "1.com/mcp",
			},
			{
				ID:     "2",
				MCPURL: "2.com/mcp",
			},
			{
				ID: "3",
			},
		},
		Experiments: requiredExperiments,
		Logger:      logger,
		Clock:       quartz.NewReal(),
	})
	require.NoError(t, err)
	require.NotNil(t, srv)

	// When: requesting all external auth links, return all.
	db.EXPECT().GetExternalAuthLinksByUserID(gomock.Any(), gomock.Any()).MinTimes(1).DoAndReturn(func(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error) {
		return []database.ExternalAuthLink{
			{
				UserID:           userID,
				ProviderID:       "1",
				OAuthAccessToken: "1-token",
			},
			{
				UserID:           userID,
				ProviderID:       "2",
				OAuthAccessToken: "2-token",
				OAuthExpiry:      dbtime.Now().Add(-time.Minute), // This token is expired and should not be returned.
			},
			{
				UserID:           userID,
				ProviderID:       "3",
				OAuthAccessToken: "3-token",
			},
		}, nil
	})

	// When: accessing the MCP server access tokens, only the 2 with MCP configured should be returned, and the 1 without should
	// not fail the request but rather have an error returned specifically for that server.
	resp, err := srv.GetMCPServerAccessTokensBatch(t.Context(), &proto.GetMCPServerAccessTokensBatchRequest{
		UserId:             uuid.NewString(),
		McpServerConfigIds: []string{"1", "1", "2", "3"}, // Duplicates must be tolerated.
	})
	require.NoError(t, err)

	// Then: 2 MCP servers are eligible but only 1 will return a valid token as the other expired.
	require.Len(t, resp.GetAccessTokens(), 1)
	require.Equal(t, "1-token", resp.GetAccessTokens()["1"])
	require.Len(t, resp.GetErrors(), 2)
	require.Contains(t, resp.GetErrors()["2"], aibridgedserver.ErrExpiredOrInvalidOAuthToken.Error())
	require.Contains(t, resp.GetErrors()["3"], aibridgedserver.ErrNoMCPConfigFound.Error())
}

func TestRecordInterception(t *testing.T) {
	t.Parallel()

	var (
		metadataProto = map[string]*anypb.Any{
			"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
		}
		metadataJSON = `{"key":"value"}`
	)

	// organizationCase records an interception with the given raw organization
	// attribution and expects the given values in the insert. Invalid pairs
	// must be stored as NULL without failing the record.
	organizationCase := func(name, rawOrgID, rawSource string, wantOrgID uuid.NullUUID, wantSource database.NullAIBridgeInterceptionOrganizationSource) testRecordMethodCase[*proto.RecordInterceptionRequest] {
		return testRecordMethodCase[*proto.RecordInterceptionRequest]{
			name: name,
			request: &proto.RecordInterceptionRequest{
				Id:                 uuid.NewString(),
				ApiKeyId:           uuid.NewString(),
				InitiatorId:        uuid.NewString(),
				Provider:           "anthropic",
				Model:              "claude-4-opus",
				Metadata:           metadataProto,
				StartedAt:          timestamppb.Now(),
				OrganizationId:     rawOrgID,
				OrganizationSource: rawSource,
			},
			setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
				interceptionID, err := uuid.Parse(req.GetId())
				assert.NoError(t, err, "parse interception UUID")
				initiatorID, err := uuid.Parse(req.GetInitiatorId())
				assert.NoError(t, err, "parse interception initiator UUID")

				db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
					ID:                 interceptionID,
					APIKeyID:           sql.NullString{String: req.ApiKeyId, Valid: true},
					InitiatorID:        initiatorID,
					Provider:           req.GetProvider(),
					ProviderName:       req.GetProvider(),
					Model:              req.GetModel(),
					Metadata:           json.RawMessage(metadataJSON),
					StartedAt:          req.StartedAt.AsTime().UTC(),
					CredentialKind:     database.CredentialKindCentralized,
					OrganizationID:     wantOrgID,
					OrganizationSource: wantSource,
				}).Return(database.AIBridgeInterception{
					ID:                 interceptionID,
					InitiatorID:        initiatorID,
					OrganizationID:     wantOrgID,
					OrganizationSource: wantSource,
				}, nil)
			},
		}
	}

	orgID := uuid.New()
	noOrgID := uuid.NullUUID{}
	noSource := database.NullAIBridgeInterceptionOrganizationSource{}
	sourceWorkspace := string(database.AibridgeInterceptionOrganizationSourceWorkspace)
	var organizationCases []testRecordMethodCase[*proto.RecordInterceptionRequest]
	for _, source := range database.AllAIBridgeInterceptionOrganizationSourceValues() {
		organizationCases = append(organizationCases, organizationCase(
			"organization attribution stored/"+string(source), orgID.String(), string(source),
			uuid.NullUUID{UUID: orgID, Valid: true},
			database.NullAIBridgeInterceptionOrganizationSource{AIBridgeInterceptionOrganizationSource: source, Valid: true},
		))
	}
	organizationCases = append(organizationCases,
		organizationCase("absent organization treated as null", "", "", noOrgID, noSource),
		organizationCase("organization without source dropped", orgID.String(), "", noOrgID, noSource),
		organizationCase("source without organization dropped", "", sourceWorkspace, noOrgID, noSource),
		organizationCase("invalid organization ID dropped", "not-a-uuid", sourceWorkspace, noOrgID, noSource),
		organizationCase("unknown organization source dropped", orgID.String(), "not-a-source", noOrgID, noSource),
	)

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordInterceptionRequest) (*proto.RecordInterceptionResponse, error) {
			return srv.RecordInterception(ctx, req)
		},
		append([]testRecordMethodCase[*proto.RecordInterceptionRequest]{
			{
				name: "valid interception",
				request: &proto.RecordInterceptionRequest{
					Id:             uuid.NewString(),
					ApiKeyId:       uuid.NewString(),
					InitiatorId:    uuid.NewString(),
					Provider:       "anthropic",
					ProviderName:   "anthropic",
					Model:          "claude-4-opus",
					Metadata:       metadataProto,
					StartedAt:      timestamppb.Now(),
					CredentialKind: "byok",
					CredentialHint: "sk-a...efgh",
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:             interceptionID,
						APIKeyID:       sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:    initiatorID,
						Provider:       req.GetProvider(),
						ProviderName:   req.GetProviderName(),
						Model:          req.GetModel(),
						Metadata:       json.RawMessage(metadataJSON),
						StartedAt:      req.StartedAt.AsTime().UTC(),
						CredentialKind: database.CredentialKindByok,
						CredentialHint: "sk-a...efgh",
					}).Return(database.AIBridgeInterception{
						ID:             interceptionID,
						APIKeyID:       sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:    initiatorID,
						Provider:       req.GetProvider(),
						ProviderName:   req.GetProviderName(),
						Model:          req.GetModel(),
						StartedAt:      req.StartedAt.AsTime().UTC(),
						CredentialKind: database.CredentialKindByok,
						CredentialHint: "sk-a...efgh",
					}, nil)
				},
			},
			{
				name: "valid interception with client session ID",
				request: &proto.RecordInterceptionRequest{
					Id:              uuid.NewString(),
					ApiKeyId:        uuid.NewString(),
					InitiatorId:     uuid.NewString(),
					Provider:        "anthropic",
					Model:           "claude-4-opus",
					Metadata:        metadataProto,
					StartedAt:       timestamppb.Now(),
					ClientSessionId: new("session-abc-123"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:              interceptionID,
						APIKeyID:        sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:     initiatorID,
						Provider:        req.GetProvider(),
						ProviderName:    req.GetProvider(),
						Model:           req.GetModel(),
						Metadata:        json.RawMessage(metadataJSON),
						StartedAt:       req.StartedAt.AsTime().UTC(),
						ClientSessionID: sql.NullString{String: "session-abc-123", Valid: true},
						CredentialKind:  database.CredentialKindCentralized,
					}).Return(database.AIBridgeInterception{
						ID:              interceptionID,
						APIKeyID:        sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:     initiatorID,
						Provider:        req.GetProvider(),
						ProviderName:    req.GetProvider(),
						Model:           req.GetModel(),
						StartedAt:       req.StartedAt.AsTime().UTC(),
						ClientSessionID: sql.NullString{String: "session-abc-123", Valid: true},
					}, nil)
				},
			},
			{
				name: "empty client session ID treated as null",
				request: &proto.RecordInterceptionRequest{
					Id:              uuid.NewString(),
					ApiKeyId:        uuid.NewString(),
					InitiatorId:     uuid.NewString(),
					Provider:        "anthropic",
					Model:           "claude-4-opus",
					Metadata:        metadataProto,
					StartedAt:       timestamppb.Now(),
					ClientSessionId: new(""),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:              interceptionID,
						APIKeyID:        sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:     initiatorID,
						Provider:        req.GetProvider(),
						ProviderName:    req.GetProvider(),
						Model:           req.GetModel(),
						Metadata:        json.RawMessage(metadataJSON),
						StartedAt:       req.StartedAt.AsTime().UTC(),
						ClientSessionID: sql.NullString{},
						CredentialKind:  database.CredentialKindCentralized,
					}).Return(database.AIBridgeInterception{
						ID:           interceptionID,
						APIKeyID:     sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:  initiatorID,
						Provider:     req.GetProvider(),
						ProviderName: req.GetProvider(),
						Model:        req.GetModel(),
						StartedAt:    req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "valid interception with agent firewall correlation",
				request: &proto.RecordInterceptionRequest{
					Id:                          uuid.NewString(),
					ApiKeyId:                    uuid.NewString(),
					InitiatorId:                 uuid.NewString(),
					Provider:                    "anthropic",
					Model:                       "claude-4-opus",
					Metadata:                    metadataProto,
					StartedAt:                   timestamppb.Now(),
					AgentFirewallSessionId:      new(uuid.NewString()),
					AgentFirewallSequenceNumber: new(int32(42)),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")
					agentFirewallSessionID, err := uuid.Parse(req.GetAgentFirewallSessionId())
					assert.NoError(t, err, "parse agent firewall session UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:                          interceptionID,
						APIKeyID:                    sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:                 initiatorID,
						Provider:                    req.GetProvider(),
						ProviderName:                req.GetProvider(),
						Model:                       req.GetModel(),
						Metadata:                    json.RawMessage(metadataJSON),
						StartedAt:                   req.StartedAt.AsTime().UTC(),
						CredentialKind:              database.CredentialKindCentralized,
						AgentFirewallSessionID:      uuid.NullUUID{UUID: agentFirewallSessionID, Valid: true},
						AgentFirewallSequenceNumber: sql.NullInt32{Int32: 42, Valid: true},
					}).Return(database.AIBridgeInterception{
						ID:          interceptionID,
						InitiatorID: initiatorID,
						Provider:    req.GetProvider(),
						Model:       req.GetModel(),
						StartedAt:   req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "absent agent firewall fields treated as null",
				request: &proto.RecordInterceptionRequest{
					Id:          uuid.NewString(),
					ApiKeyId:    uuid.NewString(),
					InitiatorId: uuid.NewString(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					Metadata:    metadataProto,
					StartedAt:   timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:                          interceptionID,
						APIKeyID:                    sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:                 initiatorID,
						Provider:                    req.GetProvider(),
						ProviderName:                req.GetProvider(),
						Model:                       req.GetModel(),
						Metadata:                    json.RawMessage(metadataJSON),
						StartedAt:                   req.StartedAt.AsTime().UTC(),
						CredentialKind:              database.CredentialKindCentralized,
						AgentFirewallSessionID:      uuid.NullUUID{},
						AgentFirewallSequenceNumber: sql.NullInt32{},
					}).Return(database.AIBridgeInterception{
						ID:          interceptionID,
						InitiatorID: initiatorID,
						Provider:    req.GetProvider(),
						Model:       req.GetModel(),
						StartedAt:   req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "invalid agent firewall session ID treated as null",
				request: &proto.RecordInterceptionRequest{
					Id:                          uuid.NewString(),
					ApiKeyId:                    uuid.NewString(),
					InitiatorId:                 uuid.NewString(),
					Provider:                    "anthropic",
					Model:                       "claude-4-opus",
					Metadata:                    metadataProto,
					StartedAt:                   timestamppb.Now(),
					AgentFirewallSessionId:      new("not-a-uuid"),
					AgentFirewallSequenceNumber: new(int32(7)),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					// Malformed agent firewall session ID is stored as null
					// (and logged) rather than failing the interception.
					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:                          interceptionID,
						APIKeyID:                    sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:                 initiatorID,
						Provider:                    req.GetProvider(),
						ProviderName:                req.GetProvider(),
						Model:                       req.GetModel(),
						Metadata:                    json.RawMessage(metadataJSON),
						StartedAt:                   req.StartedAt.AsTime().UTC(),
						CredentialKind:              database.CredentialKindCentralized,
						AgentFirewallSessionID:      uuid.NullUUID{},
						AgentFirewallSequenceNumber: sql.NullInt32{Int32: 7, Valid: true},
					}).Return(database.AIBridgeInterception{
						ID:          interceptionID,
						InitiatorID: initiatorID,
						Provider:    req.GetProvider(),
						Model:       req.GetModel(),
						StartedAt:   req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "invalid interception ID",
				request: &proto.RecordInterceptionRequest{
					Id:          "not-a-uuid",
					InitiatorId: uuid.NewString(),
					ApiKeyId:    uuid.NewString(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					StartedAt:   timestamppb.Now(),
				},
				expectedErr: "invalid interception ID",
			},
			{
				name: "invalid initiator ID",
				request: &proto.RecordInterceptionRequest{
					Id:          uuid.NewString(),
					ApiKeyId:    uuid.NewString(),
					InitiatorId: "not-a-uuid",
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					StartedAt:   timestamppb.Now(),
				},
				expectedErr: "invalid initiator ID",
			},
			{
				name: "invalid interception no api key set",
				request: &proto.RecordInterceptionRequest{
					Id:          uuid.NewString(),
					InitiatorId: uuid.NewString(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					Metadata:    metadataProto,
					StartedAt:   timestamppb.Now(),
				},
				expectedErr: "empty API key ID",
			},
			{
				name: "provider name differs from provider type",
				request: &proto.RecordInterceptionRequest{
					Id:           uuid.NewString(),
					ApiKeyId:     uuid.NewString(),
					InitiatorId:  uuid.NewString(),
					Provider:     "copilot",
					ProviderName: "copilot-business",
					Model:        "gpt-4o",
					StartedAt:    timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:             interceptionID,
						APIKeyID:       sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:    initiatorID,
						Provider:       "copilot",
						ProviderName:   "copilot-business",
						Model:          req.GetModel(),
						Metadata:       json.RawMessage("{}"),
						StartedAt:      req.StartedAt.AsTime().UTC(),
						CredentialKind: database.CredentialKindCentralized,
					}).Return(database.AIBridgeInterception{
						ID:           interceptionID,
						InitiatorID:  initiatorID,
						Provider:     "copilot",
						ProviderName: "copilot-business",
						Model:        req.GetModel(),
						StartedAt:    req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "empty provider name defaults to provider",
				request: &proto.RecordInterceptionRequest{
					Id:          uuid.NewString(),
					ApiKeyId:    uuid.NewString(),
					InitiatorId: uuid.NewString(),
					Provider:    "copilot",
					Model:       "gpt-4o",
					StartedAt:   timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:             interceptionID,
						APIKeyID:       sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:    initiatorID,
						Provider:       "copilot",
						ProviderName:   "copilot",
						Model:          req.GetModel(),
						Metadata:       json.RawMessage("{}"),
						StartedAt:      req.StartedAt.AsTime().UTC(),
						CredentialKind: database.CredentialKindCentralized,
					}).Return(database.AIBridgeInterception{
						ID:           interceptionID,
						InitiatorID:  initiatorID,
						Provider:     "copilot",
						ProviderName: "copilot",
						Model:        req.GetModel(),
						StartedAt:    req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "whitespace provider name defaults to provider",
				request: &proto.RecordInterceptionRequest{
					Id:           uuid.NewString(),
					ApiKeyId:     uuid.NewString(),
					InitiatorId:  uuid.NewString(),
					Provider:     "copilot",
					ProviderName: "   ",
					Model:        "gpt-4o",
					StartedAt:    timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")
					initiatorID, err := uuid.Parse(req.GetInitiatorId())
					assert.NoError(t, err, "parse interception initiator UUID")

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{
						ID:             interceptionID,
						APIKeyID:       sql.NullString{String: req.ApiKeyId, Valid: true},
						InitiatorID:    initiatorID,
						Provider:       "copilot",
						ProviderName:   "copilot",
						Model:          req.GetModel(),
						Metadata:       json.RawMessage("{}"),
						StartedAt:      req.StartedAt.AsTime().UTC(),
						CredentialKind: database.CredentialKindCentralized,
					}).Return(database.AIBridgeInterception{
						ID:           interceptionID,
						InitiatorID:  initiatorID,
						Provider:     "copilot",
						ProviderName: "copilot",
						Model:        req.GetModel(),
						StartedAt:    req.StartedAt.AsTime().UTC(),
					}, nil)
				},
			},
			{
				name: "database error",
				request: &proto.RecordInterceptionRequest{
					Id:          uuid.NewString(),
					ApiKeyId:    uuid.NewString(),
					InitiatorId: uuid.NewString(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					StartedAt:   timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{}, sql.ErrConnDone)
				},
				expectedErr: "start interception",
			},
			{
				name: "ok with parent correlation",
				request: &proto.RecordInterceptionRequest{
					Id:                    uuid.UUID{3}.String(),
					ApiKeyId:              uuid.NewString(),
					InitiatorId:           uuid.NewString(),
					Provider:              "anthropic",
					Model:                 "claude-4-opus",
					StartedAt:             timestamppb.Now(),
					CorrelatingToolCallId: new("call_abc"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					selfID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse self UUID")
					parentID := uuid.UUID{4}
					rootID := uuid.UUID{5}

					db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(
						gomock.Any(),
						"call_abc",
					).Return(database.GetAIBridgeInterceptionLineageByToolCallIDRow{
						ThreadParentID: parentID,
						ThreadRootID:   rootID,
					}, nil)

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeInterceptionParams) bool {
						return assert.Equal(t, selfID, p.ID, "ID") &&
							assert.Equal(t, uuid.NullUUID{UUID: parentID, Valid: true}, p.ThreadParentInterceptionID, "thread parent interception ID") &&
							assert.Equal(t, uuid.NullUUID{UUID: rootID, Valid: true}, p.ThreadRootInterceptionID, "thread root interception ID")
					})).Return(database.AIBridgeInterception{
						ID: selfID,
					}, nil)
				},
			},
			{
				name: "no lineage",
				request: &proto.RecordInterceptionRequest{
					Id:                    uuid.UUID{3}.String(),
					ApiKeyId:              uuid.NewString(),
					InitiatorId:           uuid.NewString(),
					Provider:              "anthropic",
					Model:                 "claude-4-opus",
					StartedAt:             timestamppb.Now(),
					CorrelatingToolCallId: new("call_abc"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					selfID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse self UUID")

					db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(
						gomock.Any(),
						"call_abc",
					).Return(database.GetAIBridgeInterceptionLineageByToolCallIDRow{}, sql.ErrNoRows)

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeInterceptionParams) bool {
						return assert.Equal(t, selfID, p.ID, "ID") &&
							assert.Equal(t, uuid.NullUUID{}, p.ThreadParentInterceptionID, "thread parent interception ID") &&
							assert.Equal(t, uuid.NullUUID{}, p.ThreadRootInterceptionID, "thread root interception ID")
					})).Return(database.AIBridgeInterception{
						ID: selfID,
					}, nil)
				},
			},
			{
				name: "parent without root", // This should never happen since GetAIBridgeInterceptionLineageByToolCallID always returns both, but still...
				request: &proto.RecordInterceptionRequest{
					Id:                    uuid.UUID{3}.String(),
					ApiKeyId:              uuid.NewString(),
					InitiatorId:           uuid.NewString(),
					Provider:              "anthropic",
					Model:                 "claude-4-opus",
					StartedAt:             timestamppb.Now(),
					CorrelatingToolCallId: new("call_abc"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					selfID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse self UUID")
					parentID := uuid.UUID{4}

					db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(
						gomock.Any(),
						"call_abc",
					).Return(database.GetAIBridgeInterceptionLineageByToolCallIDRow{
						ThreadParentID: parentID,
					}, nil)

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeInterceptionParams) bool {
						return assert.Equal(t, selfID, p.ID, "ID") &&
							assert.Equal(t, uuid.NullUUID{UUID: parentID, Valid: true}, p.ThreadParentInterceptionID, "thread parent interception ID") &&
							assert.Equal(t, uuid.NullUUID{}, p.ThreadRootInterceptionID, "thread root interception ID not expected")
					})).Return(database.AIBridgeInterception{
						ID: selfID,
					}, nil)
				},
			},
			{
				name: "ok no parent found",
				request: &proto.RecordInterceptionRequest{
					Id:                    uuid.UUID{5}.String(),
					ApiKeyId:              uuid.NewString(),
					InitiatorId:           uuid.NewString(),
					Provider:              "anthropic",
					Model:                 "claude-4-opus",
					StartedAt:             timestamppb.Now(),
					CorrelatingToolCallId: new("call_orphan"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) {
					selfID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse self UUID")

					db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(
						gomock.Any(),
						"call_orphan",
					).Return(database.GetAIBridgeInterceptionLineageByToolCallIDRow{}, sql.ErrNoRows)

					db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeInterceptionParams) bool {
						return assert.Equal(t, selfID, p.ID, "ID") &&
							assert.Equal(t, uuid.NullUUID{}, p.ThreadParentInterceptionID, "thread parent interception ID") &&
							assert.Equal(t, uuid.NullUUID{}, p.ThreadRootInterceptionID, "thread root interception ID")
					})).Return(database.AIBridgeInterception{
						ID: selfID,
					}, nil)
				},
			},
		}, organizationCases...),
	)
}

func TestRecordInterceptionEnded(t *testing.T) {
	t.Parallel()

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse, error) {
			return srv.RecordInterceptionEnded(ctx, req)
		},
		[]testRecordMethodCase[*proto.RecordInterceptionEndedRequest]{
			{
				name: "ok",
				request: &proto.RecordInterceptionEndedRequest{
					Id:             uuid.UUID{1}.String(),
					EndedAt:        timestamppb.Now(),
					CredentialHint: "sk-a...efgh",
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")

					db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
						ID:             interceptionID,
						EndedAt:        req.EndedAt.AsTime(),
						CredentialHint: req.CredentialHint,
					}).Return(database.AIBridgeInterception{
						ID:             interceptionID,
						InitiatorID:    uuid.UUID{2},
						Provider:       "prov",
						Model:          "mod",
						StartedAt:      time.Now(),
						EndedAt:        sql.NullTime{Time: req.EndedAt.AsTime(), Valid: true},
						CredentialHint: req.CredentialHint,
					}, nil)
				},
			},
			{
				name: "ok_with_error",
				request: &proto.RecordInterceptionEndedRequest{
					Id:           uuid.UUID{1}.String(),
					EndedAt:      timestamppb.Now(),
					ErrorType:    protobufproto.String(string(database.AibridgeInterceptionErrorTypeRateLimited)),
					ErrorMessage: protobufproto.String("rate limited by upstream"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")

					db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
						ID:      interceptionID,
						EndedAt: req.EndedAt.AsTime(),
						ErrorType: database.NullAIBridgeInterceptionErrorType{
							AIBridgeInterceptionErrorType: database.AIBridgeInterceptionErrorType(req.GetErrorType()),
							Valid:                         true,
						},
						ErrorMessage: sql.NullString{String: req.GetErrorMessage(), Valid: true},
					}).Return(database.AIBridgeInterception{ID: interceptionID}, nil)
				},
			},
			{
				name: "invalid_error_type_is_unknown",
				request: &proto.RecordInterceptionEndedRequest{
					Id:        uuid.UUID{1}.String(),
					EndedAt:   timestamppb.Now(),
					ErrorType: protobufproto.String("not-a-real-type"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")

					// A non-empty but unrecognized error type is stored as
					// 'unknown' (not NULL), keeping the error columns consistent.
					db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
						ID:      interceptionID,
						EndedAt: req.EndedAt.AsTime(),
						ErrorType: database.NullAIBridgeInterceptionErrorType{
							AIBridgeInterceptionErrorType: database.AibridgeInterceptionErrorTypeUnknown,
							Valid:                         true,
						},
					}).Return(database.AIBridgeInterception{ID: interceptionID}, nil)
				},
			},
			{
				name: "message_without_error_type_stores_neither",
				request: &proto.RecordInterceptionEndedRequest{
					Id:           uuid.UUID{1}.String(),
					EndedAt:      timestamppb.Now(),
					ErrorMessage: protobufproto.String("orphan message with no type"),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
					interceptionID, err := uuid.Parse(req.GetId())
					assert.NoError(t, err, "parse interception UUID")

					// A message without a type is not a categorized error, so
					// both columns stay NULL to preserve the both-NULL == success
					// invariant rather than persisting a half-populated error.
					db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
						ID:           interceptionID,
						EndedAt:      req.EndedAt.AsTime(),
						ErrorType:    database.NullAIBridgeInterceptionErrorType{},
						ErrorMessage: sql.NullString{},
					}).Return(database.AIBridgeInterception{ID: interceptionID}, nil)
				},
			},
			{
				name: "bad_uuid_error",
				request: &proto.RecordInterceptionEndedRequest{
					Id: "this-is-not-uuid",
				},
				setupMocks:  func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {},
				expectedErr: "invalid interception ID",
			},
			{
				name: "database_error",
				request: &proto.RecordInterceptionEndedRequest{
					Id:      uuid.UUID{1}.String(),
					EndedAt: timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
					db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{}, sql.ErrConnDone)
				},
				expectedErr: "end interception: " + sql.ErrConnDone.Error(),
			},
		},
	)
}

func TestRecordTokenUsage(t *testing.T) {
	t.Parallel()

	var (
		metadataProto = map[string]*anypb.Any{
			"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
		}
		metadataJSON = `{"key":"value"}`
		// Use fixed dates to keep the test deterministic.
		now = time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)
	)

	// Budget resolution falls through to the Everyone group, for cases that vary
	// only provider resolution.
	expectBudgetLookups := func(db *dbmock.MockStore, intc database.AIBridgeInterception) {
		db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intc.ID).Return(intc, nil)
		db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(intc.InitiatorID)).
			Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
		db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(intc.InitiatorID)).
			Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
		db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(intc.InitiatorID)).Return(uuid.New(), nil)
	}

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) {
			return srv.RecordTokenUsage(ctx, req)
		},
		[]testRecordMethodCase[*proto.RecordTokenUsageRequest]{
			{
				// Budget resolves via group lookup, model is priced.
				name: "valid token usage with effective group and cost",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					CreatedAt:             timestamppb.New(now),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					groupID := uuid.New()
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
					price := &database.AIModelPrice{
						Provider:        intc.Provider,
						Model:           intc.Model,
						InputPrice:      sql.NullInt64{Int64: 3_000_000, Valid: true},
						OutputPrice:     sql.NullInt64{Int64: 6_000_000, Valid: true},
						CacheReadPrice:  sql.NullInt64{Int64: 300_000, Valid: true},
						CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true},
					}
					// No override
					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

					const wantCost int64 = 1555

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.Equal(t, price.InputPrice, p.InputPriceMicros, "input price") ||
							!assert.Equal(t, price.OutputPrice, p.OutputPriceMicros, "output price") ||
							!assert.Equal(t, price.CacheReadPrice, p.CacheReadPriceMicros, "cache read price") ||
							!assert.Equal(t, price.CacheWritePrice, p.CacheWritePriceMicros, "cache write price") ||
							!assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{
						UserID:           intc.InitiatorID,
						EffectiveGroupID: groupID,
						Day:              now.UTC().Truncate(24 * time.Hour),
						CostMicros:       wantCost,
					}).Return(database.AIUserDailySpend{}, nil)

					db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
						Return(database.GetUserAISpendSinceRow{SpendMicros: wantCost}, nil)
				},
				// A priced model does not increment unpriced_token_usage_records_total.
				assertMetrics: func(t *testing.T, reg *prometheus.Registry) {
					require.Nil(t, promhelp.MetricValue(t, reg, "cost_control_unpriced_token_usage_records_total",
						prometheus.Labels{"provider": "anthropic-eu", "provider_type": "anthropic", "model": "claude-sonnet-4-6", "provider_model": ""}))
				},
			},
			{
				// Budget resolves via user override, model is priced.
				name: "valid token usage with user override and cost",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.New(now),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					overrideGroupID := uuid.New()
					override := &database.GetUserAIBudgetOverrideWithOrganizationRow{
						UserID:           intc.InitiatorID,
						GroupID:          overrideGroupID,
						SpendLimitMicros: 1_500_000_000,
					}
					price := &database.AIModelPrice{
						Provider:   intc.Provider,
						Model:      intc.Model,
						InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true},
					}
					// No group
					expectTokenUsageCostLookups(db, intc, override, nil, nil, price)

					// input 300.
					const wantCost int64 = 300

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						// Override group wins.
						if !assert.Equal(t, uuid.NullUUID{UUID: overrideGroupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{
						UserID:           intc.InitiatorID,
						EffectiveGroupID: overrideGroupID,
						Day:              now.UTC().Truncate(24 * time.Hour),
						CostMicros:       wantCost,
					}).Return(database.AIUserDailySpend{}, nil)

					db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
						Return(database.GetUserAISpendSinceRow{SpendMicros: wantCost}, nil)
				},
			},
			{
				// No override or group budget, so attribution falls back to the
				// user's Everyone group.
				name: "valid token usage falls back to the Everyone group",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.New(now),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					everyoneID := uuid.New()
					price := &database.AIModelPrice{
						Provider:   intc.Provider,
						Model:      intc.Model,
						InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true},
					}
					expectTokenUsageCostLookups(db, intc, nil, nil, &everyoneID, price)

					// input 300.
					const wantCost int64 = 300

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						return assert.Equal(t, uuid.NullUUID{UUID: everyoneID, Valid: true}, p.EffectiveGroupID, "effective group ID") &&
							assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost")
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{
						UserID:           intc.InitiatorID,
						EffectiveGroupID: everyoneID,
						Day:              now.UTC().Truncate(24 * time.Hour),
						CostMicros:       wantCost,
					}).Return(database.AIUserDailySpend{}, nil)
				},
			},
			{
				// Model has no price row, so cost is NULL.
				name: "valid token usage with effective group and no price",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					CreatedAt:             timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					groupID := uuid.New()
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
					// Budget resolves to a group, but the model has no price row.
					// The resolved group must survive the price lookup's early
					// return on sql.ErrNoRows, while prices and cost stay NULL.
					expectTokenUsageCostLookups(db, intc, nil, group, nil, nil)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.False(t, p.InputPriceMicros.Valid, "input price null") ||
							!assert.False(t, p.OutputPriceMicros.Valid, "output price null") ||
							!assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") ||
							!assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") ||
							!assert.False(t, p.CostMicros.Valid, "cost null") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					// Spend update is skipped because cost is NULL.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
				// A missing price row increments unpriced_token_usage_records_total.
				assertMetrics: func(t *testing.T, reg *prometheus.Registry) {
					require.Equal(t, 1, promhelp.CounterValue(t, reg, "cost_control_unpriced_token_usage_records_total",
						prometheus.Labels{"provider": "anthropic-eu", "provider_type": "anthropic", "model": "claude-sonnet-4-6", "provider_model": ""}))
				},
			},
			{
				// Implausible counts are rejected before any DB work, so no row
				// is written. Persisting them would poison the organization
				// spend export, whose SUM cast raises rather than wraps.
				name:           "token usage above the allowed range is rejected",
				expectErrorLog: true,
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    math.MaxInt64,
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "reported token usage is out of range",
			},
			{
				name:           "negative token usage is rejected",
				expectErrorLog: true,
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    -1_000_000,
					OutputTokens:   2_000_000,
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "reported token usage is out of range",
			},
			{
				// Plausible token counts against a price row six orders of
				// magnitude too high. The record is written anyway, with prices
				// snapshotted and cost NULL.
				name:           "valid token usage with cost out of range",
				expectErrorLog: true,
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    1_000_000,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					groupID := uuid.New()
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
					// $20M per million tokens puts a 1M-token request well past
					// the per-interception cost bound.
					price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 20_000_000_000_000, Valid: true}}
					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						// Prices and tokens are populated even though cost is NULL.
						if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.True(t, p.InputPriceMicros.Valid, "input price populated") ||
							!assert.False(t, p.CostMicros.Valid, "cost null") ||
							!assert.Equal(t, int64(1_000_000), p.InputTokens, "input tokens recorded") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					// Spend update is skipped because cost is NULL.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
			},
			{
				// Price row exists with NULL columns, so cost is 0 (Valid).
				name: "valid token usage with effective group and NULL prices",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					CreatedAt:             timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					groupID := uuid.New()
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
					// The price row exists but every price column is NULL. Each
					// category is treated as zero for cost, so the columns are
					// recorded as NULL while cost is recorded as 0 (not NULL):
					// cost's NULL-ness tracks price row presence, not the price
					// values.
					price := &database.AIModelPrice{
						Provider:        intc.Provider,
						Model:           intc.Model,
						InputPrice:      sql.NullInt64{Valid: false},
						OutputPrice:     sql.NullInt64{Valid: false},
						CacheReadPrice:  sql.NullInt64{Valid: false},
						CacheWritePrice: sql.NullInt64{Valid: false},
					}
					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.False(t, p.InputPriceMicros.Valid, "input price null") ||
							!assert.False(t, p.OutputPriceMicros.Valid, "output price null") ||
							!assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") ||
							!assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") ||
							// Cost is recorded as 0 (Valid), not NULL, because the
							// price row exists.
							!assert.Equal(t, sql.NullInt64{Int64: 0, Valid: true}, p.CostMicros, "cost zero") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					// Spend update is skipped because cost is 0.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
			},
			{
				// Model is priced at zero, so cost is 0 (Valid).
				name: "valid token usage with effective group and zero prices",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					CreatedAt:             timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					groupID := uuid.New()
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000}
					// A model priced at zero is distinct from an unpriced model:
					// the price columns and cost are recorded as 0, not NULL.
					price := &database.AIModelPrice{
						Provider:        intc.Provider,
						Model:           intc.Model,
						InputPrice:      sql.NullInt64{Int64: 0, Valid: true},
						OutputPrice:     sql.NullInt64{Int64: 0, Valid: true},
						CacheReadPrice:  sql.NullInt64{Int64: 0, Valid: true},
						CacheWritePrice: sql.NullInt64{Int64: 0, Valid: true},
					}
					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						zero := sql.NullInt64{Int64: 0, Valid: true}
						if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.Equal(t, zero, p.InputPriceMicros, "input price zero") ||
							!assert.Equal(t, zero, p.OutputPriceMicros, "output price zero") ||
							!assert.Equal(t, zero, p.CacheReadPriceMicros, "cache read price zero") ||
							!assert.Equal(t, zero, p.CacheWritePriceMicros, "cache write price zero") ||
							// Cost is 0 but recorded (Valid), not NULL.
							!assert.Equal(t, zero, p.CostMicros, "cost zero") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					// Spend update is skipped because cost is 0.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
			},
			{
				// No budget and no price row: attribution falls back to Everyone and cost is NULL.
				name: "valid token usage with no budget and no price",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					Metadata:              metadataProto,
					CreatedAt:             timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					// No budget configured, so attribution falls back to the
					// Everyone group. The model has no price row, so cost and
					// prices stay NULL.
					intc := newTestInterception(interceptionID)
					everyoneID := uuid.New()
					expectTokenUsageCostLookups(db, intc, nil, nil, &everyoneID, nil)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						if !assert.NotEqual(t, uuid.Nil, p.ID, "ID") ||
							!assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") ||
							!assert.Equal(t, req.GetMsgId(), p.ProviderResponseID, "provider response ID") ||
							!assert.Equal(t, req.GetInputTokens(), p.InputTokens, "input tokens") ||
							!assert.Equal(t, req.GetOutputTokens(), p.OutputTokens, "output tokens") ||
							!assert.Equal(t, req.GetCacheReadInputTokens(), p.CacheReadInputTokens, "cache read input tokens") ||
							!assert.Equal(t, req.GetCacheWriteInputTokens(), p.CacheWriteInputTokens, "cache write input tokens") ||
							!assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") ||
							!assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") ||
							!assert.Equal(t, uuid.NullUUID{UUID: everyoneID, Valid: true}, p.EffectiveGroupID, "effective group ID") ||
							!assert.False(t, p.InputPriceMicros.Valid, "input price null") ||
							!assert.False(t, p.OutputPriceMicros.Valid, "output price null") ||
							!assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") ||
							!assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") ||
							!assert.False(t, p.CostMicros.Valid, "cost null") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{
						ID:                    uuid.New(),
						InterceptionID:        interceptionID,
						ProviderResponseID:    req.GetMsgId(),
						InputTokens:           req.GetInputTokens(),
						OutputTokens:          req.GetOutputTokens(),
						CacheReadInputTokens:  req.GetCacheReadInputTokens(),
						CacheWriteInputTokens: req.GetCacheWriteInputTokens(),
						Metadata: pqtype.NullRawMessage{
							RawMessage: json.RawMessage(metadataJSON),
							Valid:      true,
						},
						CreatedAt: req.GetCreatedAt().AsTime(),
					}, nil)

					// Spend update is skipped because cost is NULL.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
				// A missing price row increments unpriced_token_usage_records_total.
				assertMetrics: func(t *testing.T, reg *prometheus.Registry) {
					require.Equal(t, 1, promhelp.CounterValue(t, reg, "cost_control_unpriced_token_usage_records_total",
						prometheus.Labels{"provider": "anthropic-eu", "provider_type": "anthropic", "model": "claude-sonnet-4-6", "provider_model": ""}))
				},
			},
			{
				// A user with no organization has no effective group. Spend is
				// still recorded, but with a NULL group, and the daily spend
				// update is skipped.
				name: "valid token usage with no effective group",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId:        uuid.NewString(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					CreatedAt:             timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					price := &database.AIModelPrice{
						Provider:        intc.Provider,
						Model:           intc.Model,
						InputPrice:      sql.NullInt64{Int64: 3_000_000, Valid: true},
						OutputPrice:     sql.NullInt64{Int64: 6_000_000, Valid: true},
						CacheReadPrice:  sql.NullInt64{Int64: 300_000, Valid: true},
						CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true},
					}
					// Every resolution lookup misses, including the Everyone
					// fallback, so the group stays NULL while cost is computed.
					expectTokenUsageCostLookups(db, intc, nil, nil, nil, price)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)

					const wantCost int64 = 1555

					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						if !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") ||
							!assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") {
							return false
						}
						return true
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)

					// Spend update is skipped because the effective group is NULL.
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
			},
			{
				// An azure provider has the openai upstream wire format but bills at
				// its own rates.
				name: "openai wire format priced as azure",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					intc.Provider = "openai"
					intc.ProviderName = "azure-prod"
					intc.Model = "gpt-5-mini"
					expectBudgetLookups(db, intc)

					providerID := uuid.New()
					db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).
						Return(database.AIProvider{ID: providerID, Name: intc.ProviderName, Type: database.AIProviderTypeAzure}, nil)
					db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), database.GetAIModelPriceForProviderParams{
						ProviderType: string(database.AIProviderTypeAzure),
						Model:        intc.Model,
						ProviderID:   providerID,
					}).Return(database.GetAIModelPriceForProviderRow{
						Model:      intc.Model,
						InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true},
					}, nil)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)
					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						return assert.Equal(t, sql.NullInt64{Int64: 3_000_000, Valid: true}, p.InputPriceMicros, "input price") &&
							assert.Equal(t, sql.NullInt64{Int64: 300, Valid: true}, p.CostMicros, "cost")
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
						Return(database.AIUserDailySpend{}, nil)
				},
			},
			{
				// A bedrock provider has the anthropic upstream wire format but bills
				// at its own rates.
				name: "anthropic wire format priced as bedrock",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					intc.ProviderName = "bedrock-eu"
					expectBudgetLookups(db, intc)

					providerID := uuid.New()
					db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).
						Return(database.AIProvider{ID: providerID, Name: intc.ProviderName, Type: database.AIProviderTypeBedrock}, nil)
					db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), database.GetAIModelPriceForProviderParams{
						ProviderType: string(database.AIProviderTypeBedrock),
						Model:        intc.Model,
						ProviderID:   providerID,
					}).Return(database.GetAIModelPriceForProviderRow{
						Model:      intc.Model,
						InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true},
					}, nil)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)
					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						return assert.Equal(t, sql.NullInt64{Int64: 3_000_000, Valid: true}, p.InputPriceMicros, "input price") &&
							assert.Equal(t, sql.NullInt64{Int64: 300, Valid: true}, p.CostMicros, "cost")
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
						Return(database.AIUserDailySpend{}, nil)
				},
			},
			{
				name: "unresolved provider is unpriced",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					expectBudgetLookups(db, intc)

					db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).
						Return(database.AIProvider{}, sql.ErrNoRows)
					// Without a provider there is nothing to key the price on.
					db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), gomock.Any()).Times(0)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)
					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool {
						return assert.False(t, p.InputPriceMicros.Valid, "input price null") &&
							assert.False(t, p.CostMicros.Valid, "cost null")
					})).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0)
				},
				// The metric names the provider that failed to resolve.
				assertMetrics: func(t *testing.T, reg *prometheus.Registry) {
					require.Equal(t, 1, promhelp.CounterValue(t, reg, "cost_control_unpriced_token_usage_records_total",
						prometheus.Labels{"provider": "anthropic-eu", "provider_type": "unknown", "model": "claude-sonnet-4-6", "provider_model": ""}))
				},
			},
			{
				name: "invalid interception ID",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: "not-a-uuid",
					MsgId:          "msg_123",
					InputTokens:    100,
					OutputTokens:   200,
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "failed to parse interception_id",
			},
			{
				name: "interception lookup error",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					OutputTokens:   200,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					// An unexpected interception lookup error fails the record;
					// no token usage is inserted.
					db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), interceptionID).
						Return(database.AIBridgeInterception{}, sql.ErrConnDone)
				},
				expectedErr: "get interception",
			},
			{
				name: "provider lookup error",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					// An unexpected provider lookup error (not sql.ErrNoRows) fails
					// the record.
					intc := newTestInterception(interceptionID)
					expectBudgetLookups(db, intc)
					db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).
						Return(database.AIProvider{}, sql.ErrConnDone)
				},
				expectedErr: "get configured provider",
			},
			{
				name: "price lookup error",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					OutputTokens:   200,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					// An unexpected price lookup error (not sql.ErrNoRows) fails
					// the record.
					intc := newTestInterception(interceptionID)
					db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), interceptionID).Return(intc, nil)
					db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(intc.InitiatorID)).
						Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
					db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(intc.InitiatorID)).
						Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
					db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(intc.InitiatorID)).
						Return(uuid.New(), nil)
					db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).
						Return(database.AIProvider{Name: intc.ProviderName, Type: database.AIProviderTypeAnthropic}, nil)
					db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), gomock.Any()).
						Return(database.GetAIModelPriceForProviderRow{}, sql.ErrConnDone)
				},
				expectedErr: "resolve token usage cost",
			},
			{
				name: "insert token usage error",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					OutputTokens:   200,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					everyoneID := uuid.New()
					expectTokenUsageCostLookups(db, newTestInterception(interceptionID), nil, nil, &everyoneID, nil)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)
					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{}, sql.ErrConnDone)
				},
				expectedErr: "insert token usage",
			},
			{
				name: "increment user daily spend error",
				request: &proto.RecordTokenUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					InputTokens:    100,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					intc := newTestInterception(interceptionID)
					group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: uuid.New(), SpendLimitMicros: 1_000_000_000}
					price := &database.AIModelPrice{
						Provider:   intc.Provider,
						Model:      intc.Model,
						InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true},
					}
					expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

					db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
						func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
					)
					db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
						Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil)
					db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
						Return(database.AIUserDailySpend{}, sql.ErrConnDone)
				},
				expectedErr: "increment user daily spend",
			},
		},
	)
}

// TestRecordTokenUsageAuthorized exercises RecordTokenUsage end-to-end against a
// real database through the dbauthz layer as subjectAibridged. This catches missing
// RBAC grants on the aibridged subject and verifies the cost columns round-trip
// to storage along with the daily spend row.
func TestRecordTokenUsageAuthorized(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	rawDB, _ := dbtestutil.NewDB(t)
	authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

	// Seed prerequisites via the raw (unauthorized) store. The user belongs to a
	// group with a budget, so the effective group resolves to that group.
	org := dbgen.Organization(t, rawDB, database.Organization{})
	user := dbgen.User(t, rawDB, database.User{})
	dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
	group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID})
	dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID})

	_, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
		GroupID:          group.ID,
		SpendLimitMicros: 1_000_000_000,
	})
	require.NoError(t, err, "upsert group AI budget")

	const provider, model = "anthropic", "claude-sonnet-4-6"
	priceSeed, err := json.Marshal([]map[string]any{{
		"provider":          provider,
		"model":             model,
		"input_price":       3_000_000,
		"output_price":      6_000_000,
		"cache_read_price":  300_000,
		"cache_write_price": 4_000_000,
	}})
	require.NoError(t, err)
	require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")

	// The interception's provider name resolves to this provider, whose type keys
	// the price lookup.
	aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
		Name: "anthropic-eu",
		Type: database.AIProviderTypeAnthropic,
	})

	// The organization's price overrides the site-wide one, so reading it
	// needs the organization-scoped price permission.
	require.NoError(t, rawDB.UpsertOrganizationAIProviderModelConfig(ctx, database.UpsertOrganizationAIProviderModelConfigParams{
		OrganizationID:  org.ID,
		ProviderID:      aiProvider.ID,
		Model:           model,
		InputPrice:      sql.NullInt64{Int64: 5_000_000, Valid: true},
		OutputPrice:     sql.NullInt64{Int64: 10_000_000, Valid: true},
		CacheReadPrice:  sql.NullInt64{Int64: 500_000, Valid: true},
		CacheWritePrice: sql.NullInt64{Int64: 2_000_000, Valid: true},
	}), "set organization model price")

	intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{
		InitiatorID:    user.ID,
		Provider:       provider,
		ProviderName:   aiProvider.Name,
		Model:          model,
		OrganizationID: uuid.NullUUID{UUID: org.ID, Valid: true},
		OrganizationSource: database.NullAIBridgeInterceptionOrganizationSource{
			AIBridgeInterceptionOrganizationSource: database.AibridgeInterceptionOrganizationSourceWorkspace,
			Valid:                                  true,
		},
	}, nil)

	// Use fixed dates to keep the test deterministic.
	now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)

	// The server runs every store call as subjectAibridged via the authzDB.
	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         authzDB,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
		InterceptionId:        intc.ID.String(),
		MsgId:                 "msg_e2e",
		InputTokens:           100,
		OutputTokens:          200,
		CacheReadInputTokens:  50,
		CacheWriteInputTokens: 10,
		CreatedAt:             timestamppb.New(now),
	})
	require.NoError(t, err, "record token usage")

	// Read the persisted row back via the raw store and verify the snapshot.
	tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
	require.NoError(t, err)
	require.Len(t, tokenUsages, 1)
	tokenUsage := tokenUsages[0]

	require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, tokenUsage.EffectiveGroupID, "effective group")
	require.Equal(t, sql.NullInt64{Int64: 5_000_000, Valid: true}, tokenUsage.InputPriceMicros, "input price")
	require.Equal(t, sql.NullInt64{Int64: 10_000_000, Valid: true}, tokenUsage.OutputPriceMicros, "output price")
	require.Equal(t, sql.NullInt64{Int64: 500_000, Valid: true}, tokenUsage.CacheReadPriceMicros, "cache read price")
	require.Equal(t, sql.NullInt64{Int64: 2_000_000, Valid: true}, tokenUsage.CacheWritePriceMicros, "cache write price")
	const wantCost int64 = 2545 // 500 + 2000 + 25 + 20
	require.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, tokenUsage.CostMicros, "cost")

	// The daily spend row was incremented for (user, group, today) by the same cost.
	today := now.UTC().Truncate(24 * time.Hour)
	spend, err := rawDB.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{
		UserID:           user.ID,
		EffectiveGroupID: group.ID,
		PeriodStart:      today,
	})
	require.NoError(t, err, "get user AI spend since")
	require.Equal(t, user.ID, spend.UserID, "user ID")
	require.Equal(t, group.ID, spend.EffectiveGroupID, "effective group ID")
	require.True(t, today.Equal(spend.PeriodStart), "period start: want %s, got %s", today, spend.PeriodStart)
	require.Equal(t, wantCost, spend.SpendMicros, "spend micros")
}

// TestRecordTokenUsageUsesStoredOrganization verifies that cost resolves the
// effective group inside the organization stored on the interception, and
// never picks the site-wide winner from another organization. An interception
// without an organization keeps site-wide resolution.
func TestRecordTokenUsageUsesStoredOrganization(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	rawDB, _ := dbtestutil.NewDB(t)
	authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

	// The user has a budgeted group in org A and in org B. A's limit is higher,
	// so A's group wins site-wide.
	user := dbgen.User(t, rawDB, database.User{})
	budgetedGroup := func(limitMicros int64) database.Group {
		org := dbgen.Organization(t, rawDB, database.Organization{})
		dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
		group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID})
		dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID})
		_, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
			GroupID:          group.ID,
			SpendLimitMicros: limitMicros,
		})
		require.NoError(t, err)
		return group
	}
	groupA := budgetedGroup(1_000_000_000)
	groupB := budgetedGroup(1_000_000)

	const provider, model = "anthropic", "claude-sonnet-4-6"
	priceSeed, err := json.Marshal([]map[string]any{{
		"provider":     provider,
		"model":        model,
		"input_price":  3_000_000,
		"output_price": 6_000_000,
	}})
	require.NoError(t, err)
	require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}))
	aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
		Name: "anthropic-eu",
		Type: database.AIProviderTypeAnthropic,
	})

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         authzDB,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	for _, tc := range []struct {
		name      string
		orgID     uuid.NullUUID
		wantGroup database.Group
	}{
		{
			name:      "stored organization",
			orgID:     uuid.NullUUID{UUID: groupB.OrganizationID, Valid: true},
			wantGroup: groupB,
		},
		{
			// Old interceptions and old gateways store no organization.
			name:      "no stored organization",
			wantGroup: groupA,
		},
	} {
		intcParams := database.InsertAIBridgeInterceptionParams{
			InitiatorID:    user.ID,
			Provider:       provider,
			ProviderName:   aiProvider.Name,
			Model:          model,
			OrganizationID: tc.orgID,
		}
		if tc.orgID.Valid {
			intcParams.OrganizationSource = database.NullAIBridgeInterceptionOrganizationSource{
				AIBridgeInterceptionOrganizationSource: database.AibridgeInterceptionOrganizationSourceWorkspace,
				Valid:                                  true,
			}
		}
		intc := dbgen.AIBridgeInterception(t, rawDB, intcParams, nil)

		now := dbtime.Now()
		_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
			InterceptionId: intc.ID.String(),
			MsgId:          "msg_" + tc.name,
			InputTokens:    100,
			OutputTokens:   200,
			CreatedAt:      timestamppb.New(now),
		})
		require.NoError(t, err, tc.name)

		tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
		require.NoError(t, err, tc.name)
		require.Len(t, tokenUsages, 1, tc.name)
		require.Equal(t, uuid.NullUUID{UUID: tc.wantGroup.ID, Valid: true}, tokenUsages[0].EffectiveGroupID, tc.name)
		require.True(t, tokenUsages[0].CostMicros.Valid, tc.name)

		// Each case charges a different group, so the group's spend is
		// exactly this record's cost.
		spend, err := rawDB.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{
			UserID:           user.ID,
			EffectiveGroupID: tc.wantGroup.ID,
			PeriodStart:      dbtime.StartOfDay(now),
		})
		require.NoError(t, err, tc.name)
		require.Equal(t, tokenUsages[0].CostMicros.Int64, spend.SpendMicros, tc.name)
	}
}

// TestBudgetNotificationAuthorized exercises the budget threshold notification
// path end-to-end against a real database through the dbauthz layer as
// subjectAibridged. This catches missing RBAC grants on the aibridged subject
// and verifies the group and organization labels round-trip from storage.
func TestBudgetNotificationAuthorized(t *testing.T) {
	t.Parallel()

	ctx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	rawDB, _ := dbtestutil.NewDB(t)
	authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

	// Seed prerequisites via the raw (unauthorized) store. The user belongs to a
	// group with a budget, so the effective group resolves to that group.
	org := dbgen.Organization(t, rawDB, database.Organization{})
	user := dbgen.User(t, rawDB, database.User{})
	dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})
	group := dbgen.Group(t, rawDB, database.Group{OrganizationID: org.ID})
	dbgen.GroupMember(t, rawDB, database.GroupMemberTable{UserID: user.ID, GroupID: group.ID})

	// The interception below costs 1555 micros, which crosses 85% of this limit
	// (1530 micros) without reaching the limit itself.
	_, err := rawDB.UpsertGroupAIBudget(ctx, database.UpsertGroupAIBudgetParams{
		GroupID:          group.ID,
		SpendLimitMicros: 1_800,
	})
	require.NoError(t, err, "upsert group AI budget")

	const provider, model = "anthropic", "claude-sonnet-4-6"
	priceSeed, err := json.Marshal([]map[string]any{{
		"provider":          provider,
		"model":             model,
		"input_price":       3_000_000,
		"output_price":      6_000_000,
		"cache_read_price":  300_000,
		"cache_write_price": 4_000_000,
	}})
	require.NoError(t, err)
	require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")

	aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
		Name: "anthropic-eu",
		Type: database.AIProviderTypeAnthropic,
	})
	intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{
		InitiatorID:  user.ID,
		Provider:     provider,
		ProviderName: aiProvider.Name,
		Model:        model,
	}, nil)

	enq := &notificationstest.FakeEnqueuer{}
	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         authzDB,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Enqueuer:      enq,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
		InterceptionId:        intc.ID.String(),
		MsgId:                 "msg_budget_authz",
		InputTokens:           100,
		OutputTokens:          200,
		CacheReadInputTokens:  50,
		CacheWriteInputTokens: 10,
		CreatedAt:             timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)),
	})
	require.NoError(t, err, "record token usage")

	// Notification failures are logged rather than returned, so the enqueued
	// message is the only signal that both lookups were authorized.
	sent := enq.Sent(notificationstest.WithTemplateID(notifications.TemplateAIBudgetWarningUser))
	require.Len(t, sent, 1, "expected one budget warning notification")
	require.Equal(t, group.Name, sent[0].Labels["effective_group_name"])
	require.Equal(t, org.Name, sent[0].Labels["organization_name"])
}

// TestRecordTokenUsageModelPriceResolution covers which price an interception
// snapshots when a model carries a price from the embedded book, a site-wide
// price set through the API, a price of the interception's stored
// organization, or several of them.
func TestRecordTokenUsageModelPriceResolution(t *testing.T) {
	t.Parallel()

	const provider, model, providerModel = "anthropic", "requested-model", "provider-model"

	type modelPrice struct {
		model                                string
		input, output, cacheRead, cacheWrite int64
		cost                                 int64
	}

	var (
		requestedDefault = modelPrice{model: model, input: 3_000_000, output: 6_000_000, cacheRead: 300_000, cacheWrite: 4_000_000, cost: 1555}         // 300 + 1200 + 15 + 40
		requestedCustom  = modelPrice{model: model, input: 9_000_000, output: 12_000_000, cacheRead: 900_000, cacheWrite: 8_000_000, cost: 3425}        // 900 + 2400 + 45 + 80
		providerDefault  = modelPrice{model: providerModel, input: 2_000_000, output: 4_000_000, cacheRead: 200_000, cacheWrite: 1_000_000, cost: 1020} // 200 + 800 + 10 + 10
		providerCustom   = modelPrice{model: providerModel, input: 1_000_000, output: 2_000_000, cacheRead: 100_000, cacheWrite: 500_000, cost: 510}    // 100 + 400 + 5 + 5
		requestedOrg     = modelPrice{model: model, input: 5_000_000, output: 10_000_000, cacheRead: 500_000, cacheWrite: 2_000_000, cost: 2545}        // 500 + 2000 + 25 + 20
		providerOrg      = modelPrice{model: providerModel, input: 4_000_000, output: 7_000_000, cacheRead: 400_000, cacheWrite: 3_000_000, cost: 1850} // 400 + 1400 + 20 + 30
	)

	priceSeed := func(prices []modelPrice) json.RawMessage {
		rows := make([]map[string]any, 0, len(prices))
		for _, p := range prices {
			rows = append(rows, map[string]any{
				"provider":          provider,
				"model":             p.model,
				"input_price":       p.input,
				"output_price":      p.output,
				"cache_read_price":  p.cacheRead,
				"cache_write_price": p.cacheWrite,
			})
		}
		seed, err := json.Marshal(rows)
		require.NoError(t, err)
		return seed
	}

	tests := []struct {
		name          string
		providerModel string
		defaultPrices []modelPrice
		customPrices  []modelPrice
		// storedOrg stores the organization on the interception. Without it
		// the interception has a NULL organization, like old interceptions.
		storedOrg bool
		// orgPrices are live prices of the organization for the provider.
		orgPrices []modelPrice
		// deletedOrgPrices are organization prices that were soft-deleted.
		deletedOrgPrices []modelPrice
		want             *modelPrice
	}{
		{
			name:          "RequestedCustomOnly",
			providerModel: providerModel,
			customPrices:  []modelPrice{requestedCustom},
			want:          &requestedCustom,
		},
		{
			name:          "RequestedDefaultOnly",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			want:          &requestedDefault,
		},
		{
			name:          "ProviderCustomOnly",
			providerModel: providerModel,
			customPrices:  []modelPrice{providerCustom},
			want:          &providerCustom,
		},
		{
			name:          "ProviderDefaultOnly",
			providerModel: providerModel,
			defaultPrices: []modelPrice{providerDefault},
			want:          &providerDefault,
		},
		{
			name:          "RequestedCustomWinsOverRequestedDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			customPrices:  []modelPrice{requestedCustom},
			want:          &requestedCustom,
		},
		{
			name:          "RequestedCustomWinsOverProviderCustom",
			providerModel: providerModel,
			customPrices:  []modelPrice{requestedCustom, providerCustom},
			want:          &requestedCustom,
		},
		{
			name:          "RequestedCustomWinsOverProviderDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{providerDefault},
			customPrices:  []modelPrice{requestedCustom},
			want:          &requestedCustom,
		},
		{
			name:          "RequestedDefaultWinsOverProviderCustom",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			customPrices:  []modelPrice{providerCustom},
			want:          &requestedDefault,
		},
		{
			name:          "RequestedDefaultWinsOverProviderDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault, providerDefault},
			want:          &requestedDefault,
		},
		{
			name:          "ProviderCustomWinsOverProviderDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{providerDefault},
			customPrices:  []modelPrice{providerCustom},
			want:          &providerCustom,
		},
		{
			name:          "Unpriced",
			providerModel: providerModel,
		},
		{
			name:          "ProviderModelNotReported",
			defaultPrices: []modelPrice{requestedDefault},
			want:          &requestedDefault,
		},
		{
			name:          "NoStoredOrgIgnoresOrgPriceForCustom",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			customPrices:  []modelPrice{requestedCustom},
			orgPrices:     []modelPrice{requestedOrg},
			want:          &requestedCustom,
		},
		{
			name:          "NoStoredOrgIgnoresOrgPriceForDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			orgPrices:     []modelPrice{requestedOrg},
			want:          &requestedDefault,
		},
		{
			name:          "OrgPriceWinsOverCustomAndDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			customPrices:  []modelPrice{requestedCustom},
			storedOrg:     true,
			orgPrices:     []modelPrice{requestedOrg},
			want:          &requestedOrg,
		},
		{
			name:          "AbsentOrgPriceFallsBackToCustom",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			customPrices:  []modelPrice{requestedCustom},
			storedOrg:     true,
			want:          &requestedCustom,
		},
		{
			name:             "DeletedOrgPriceFallsBackToCustom",
			providerModel:    providerModel,
			defaultPrices:    []modelPrice{requestedDefault},
			customPrices:     []modelPrice{requestedCustom},
			storedOrg:        true,
			deletedOrgPrices: []modelPrice{requestedOrg},
			want:             &requestedCustom,
		},
		{
			name:          "NoOrgOrCustomPriceFallsBackToDefault",
			providerModel: providerModel,
			defaultPrices: []modelPrice{requestedDefault},
			storedOrg:     true,
			want:          &requestedDefault,
		},
		{
			name:          "StoredOrgUnpriced",
			providerModel: providerModel,
			storedOrg:     true,
		},
		{
			name:          "ProviderOrgWinsOverProviderCustom",
			providerModel: providerModel,
			customPrices:  []modelPrice{providerCustom},
			storedOrg:     true,
			orgPrices:     []modelPrice{providerOrg},
			want:          &providerOrg,
		},
		{
			name:          "RequestedCustomWinsOverProviderOrg",
			providerModel: providerModel,
			customPrices:  []modelPrice{requestedCustom},
			storedOrg:     true,
			orgPrices:     []modelPrice{providerOrg},
			want:          &requestedCustom,
		},
		{
			name:          "RequestedOrgWinsOverProviderOrg",
			providerModel: providerModel,
			storedOrg:     true,
			orgPrices:     []modelPrice{requestedOrg, providerOrg},
			want:          &requestedOrg,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			ctx := testutil.Context(t, testutil.WaitLong)
			logger := testutil.Logger(t)

			rawDB, _ := dbtestutil.NewDB(t)
			authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

			org := dbgen.Organization(t, rawDB, database.Organization{})
			user := dbgen.User(t, rawDB, database.User{})
			dbgen.OrganizationMember(t, rawDB, database.OrganizationMember{OrganizationID: org.ID, UserID: user.ID})

			if len(tt.defaultPrices) > 0 {
				require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
					Seed:   priceSeed(tt.defaultPrices),
					Source: database.AIModelPriceSourceDefault,
				}), "seed model prices")
			}
			if len(tt.customPrices) > 0 {
				require.NoError(t, rawDB.UpsertAIModelPrices(ctx, database.UpsertAIModelPricesParams{
					Seed:   priceSeed(tt.customPrices),
					Source: database.AIModelPriceSourceCustom,
				}), "set custom model price")
			}

			aiProvider := dbgen.AIProvider(t, rawDB, database.AIProvider{
				Name: "anthropic-eu",
				Type: database.AIProviderTypeAnthropic,
			})

			setOrgPrice := func(p modelPrice) {
				require.NoError(t, rawDB.UpsertOrganizationAIProviderModelConfig(ctx, database.UpsertOrganizationAIProviderModelConfigParams{
					OrganizationID:  org.ID,
					ProviderID:      aiProvider.ID,
					Model:           p.model,
					InputPrice:      sql.NullInt64{Int64: p.input, Valid: true},
					OutputPrice:     sql.NullInt64{Int64: p.output, Valid: true},
					CacheReadPrice:  sql.NullInt64{Int64: p.cacheRead, Valid: true},
					CacheWritePrice: sql.NullInt64{Int64: p.cacheWrite, Valid: true},
				}), "set organization model price")
			}
			for _, p := range tt.deletedOrgPrices {
				setOrgPrice(p)
				deleted, err := rawDB.DeleteOrganizationAIProviderModelConfig(ctx, database.DeleteOrganizationAIProviderModelConfigParams{
					OrganizationID: org.ID,
					ProviderID:     aiProvider.ID,
					Model:          p.model,
				})
				require.NoError(t, err, "delete organization model price")
				require.EqualValues(t, 1, deleted, "deleted organization model prices")
			}
			for _, p := range tt.orgPrices {
				setOrgPrice(p)
			}

			intcParams := database.InsertAIBridgeInterceptionParams{
				InitiatorID:  user.ID,
				Provider:     provider,
				ProviderName: aiProvider.Name,
				Model:        model,
			}
			if tt.storedOrg {
				intcParams.OrganizationID = uuid.NullUUID{UUID: org.ID, Valid: true}
				intcParams.OrganizationSource = database.NullAIBridgeInterceptionOrganizationSource{
					AIBridgeInterceptionOrganizationSource: database.AibridgeInterceptionOrganizationSourceWorkspace,
					Valid:                                  true,
				}
			}
			intc := dbgen.AIBridgeInterception(t, rawDB, intcParams, nil)

			reg := prometheus.NewRegistry()
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         authzDB,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
				Metrics:       aibridgedserver.NewMetrics(reg),
			})
			require.NoError(t, err)

			_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
				InterceptionId:        intc.ID.String(),
				MsgId:                 "msg_price_resolution",
				ProviderModel:         tt.providerModel,
				InputTokens:           100,
				OutputTokens:          200,
				CacheReadInputTokens:  50,
				CacheWriteInputTokens: 10,
				CreatedAt:             timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)),
			})
			require.NoError(t, err, "record token usage")

			var want database.AIBridgeTokenUsage
			if tt.want != nil {
				want = database.AIBridgeTokenUsage{
					PricedModel:           sql.NullString{String: tt.want.model, Valid: true},
					InputPriceMicros:      sql.NullInt64{Int64: tt.want.input, Valid: true},
					OutputPriceMicros:     sql.NullInt64{Int64: tt.want.output, Valid: true},
					CacheReadPriceMicros:  sql.NullInt64{Int64: tt.want.cacheRead, Valid: true},
					CacheWritePriceMicros: sql.NullInt64{Int64: tt.want.cacheWrite, Valid: true},
					CostMicros:            sql.NullInt64{Int64: tt.want.cost, Valid: true},
				}
			}

			tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
			require.NoError(t, err)
			require.Len(t, tokenUsages, 1)

			tokenUsage := tokenUsages[0]
			require.Equal(t, sql.NullString{String: tt.providerModel, Valid: tt.providerModel != ""}, tokenUsage.ProviderModel, "provider model")
			require.Equal(t, want.PricedModel, tokenUsage.PricedModel, "priced model")
			require.Equal(t, want.InputPriceMicros, tokenUsage.InputPriceMicros, "input price")
			require.Equal(t, want.OutputPriceMicros, tokenUsage.OutputPriceMicros, "output price")
			require.Equal(t, want.CacheReadPriceMicros, tokenUsage.CacheReadPriceMicros, "cache read price")
			require.Equal(t, want.CacheWritePriceMicros, tokenUsage.CacheWritePriceMicros, "cache write price")
			require.Equal(t, want.CostMicros, tokenUsage.CostMicros, "cost")

			if tt.want == nil {
				require.Equal(t, 1, promhelp.CounterValue(t, reg, "cost_control_unpriced_token_usage_records_total",
					prometheus.Labels{"provider": aiProvider.Name, "provider_type": provider, "model": model, "provider_model": tt.providerModel}))
			}
		})
	}
}

// TestRecordTokenUsageProviderResolution covers provider resolution against a real
// database through dbauthz, where the live-row filter and name reuse apply.
func TestRecordTokenUsageProviderResolution(t *testing.T) {
	t.Parallel()

	const claudeModel, gptModel = "claude-sonnet-4-6", "gpt-5-mini"
	const anthropicInputPrice, bedrockInputPrice, openaiInputPrice, azureInputPrice int64 = 2_000_000, 3_000_000, 4_000_000, 5_000_000

	setupCtx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	rawDB, _ := dbtestutil.NewDB(t)
	authzDB := dbauthz.New(rawDB, rbac.NewStrictAuthorizer(prometheus.NewRegistry()), logger, coderdtest.AccessControlStorePointer())

	user := dbgen.User(t, rawDB, database.User{})

	// Prices differ per provider type so the asserted cost identifies which type resolved.
	priceSeed, err := json.Marshal([]map[string]any{
		{"provider": string(database.AIProviderTypeAnthropic), "model": claudeModel, "input_price": anthropicInputPrice},
		{"provider": string(database.AIProviderTypeBedrock), "model": claudeModel, "input_price": bedrockInputPrice},
		{"provider": string(database.AIProviderTypeOpenai), "model": gptModel, "input_price": openaiInputPrice},
		{"provider": string(database.AIProviderTypeAzure), "model": gptModel, "input_price": azureInputPrice},
	})
	require.NoError(t, err)
	require.NoError(t, rawDB.UpsertAIModelPrices(setupCtx, database.UpsertAIModelPricesParams{Seed: priceSeed, Source: database.AIModelPriceSourceDefault}), "seed model prices")

	srv, err := aibridgedserver.NewServer(setupCtx, aibridgedserver.Options{
		Store:         authzDB,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	cases := []struct {
		name string
		// wireProvider is the upstream wire format recorded on the interception.
		wireProvider string
		// providerName is the provider instance name recorded on the interception.
		providerName string
		// providerType is the configured provider type of the live provider.
		providerType database.AIProviderType
		model        string
		// setupProvider creates and deletes the case's providers.
		setupProvider  func(t *testing.T, ctx context.Context, providerName string, providerType database.AIProviderType)
		wantInputPrice sql.NullInt64
		wantCost       sql.NullInt64
	}{
		{
			// The common configuration, where the configured provider type matches
			// the upstream wire format.
			name:         "provider named after its type",
			wireProvider: "anthropic",
			providerName: "anthropic",
			providerType: database.AIProviderTypeAnthropic,
			model:        claudeModel,
			// One live anthropic provider.
			setupProvider: func(t *testing.T, _ context.Context, providerName string, providerType database.AIProviderType) {
				dbgen.AIProvider(t, rawDB, database.AIProvider{Name: providerName, Type: providerType})
			},
			wantInputPrice: sql.NullInt64{Int64: anthropicInputPrice, Valid: true},
			// 100 input tokens at the anthropic input price: $0.0002.
			wantCost: sql.NullInt64{Int64: 200, Valid: true},
		},
		{
			name:         "priced by configured provider type",
			wireProvider: "anthropic",
			providerName: "bedrock-eu",
			providerType: database.AIProviderTypeBedrock,
			model:        claudeModel,
			// One live bedrock provider.
			setupProvider: func(t *testing.T, _ context.Context, providerName string, providerType database.AIProviderType) {
				dbgen.AIProvider(t, rawDB, database.AIProvider{Name: providerName, Type: providerType})
			},
			wantInputPrice: sql.NullInt64{Int64: bedrockInputPrice, Valid: true},
			// 100 input tokens at the bedrock input price: $0.0003.
			wantCost: sql.NullInt64{Int64: 300, Valid: true},
		},
		{
			name:         "deleted provider is unpriced",
			wireProvider: "anthropic",
			providerName: "bedrock-deleted",
			providerType: database.AIProviderTypeBedrock,
			model:        claudeModel,
			// One bedrock provider, deleted before the usage is recorded.
			setupProvider: func(t *testing.T, ctx context.Context, providerName string, providerType database.AIProviderType) {
				provider := dbgen.AIProvider(t, rawDB, database.AIProvider{Name: providerName, Type: providerType})
				require.NoError(t, rawDB.DeleteAIProviderByID(ctx, provider.ID), "delete provider")
			},
			wantInputPrice: sql.NullInt64{Valid: false},
			wantCost:       sql.NullInt64{Valid: false},
		},
		{
			// Names are unique only among live providers, so a deleted name can be
			// reused by a provider of a different configured provider type.
			name:         "reused name resolves to the live provider",
			wireProvider: "openai",
			providerName: "reused-name",
			providerType: database.AIProviderTypeAzure,
			model:        gptModel,
			// A deleted openai provider and a live azure provider sharing the name.
			setupProvider: func(t *testing.T, ctx context.Context, providerName string, providerType database.AIProviderType) {
				deleted := dbgen.AIProvider(t, rawDB, database.AIProvider{Name: providerName, Type: database.AIProviderTypeOpenai})
				require.NoError(t, rawDB.DeleteAIProviderByID(ctx, deleted.ID), "delete provider")
				dbgen.AIProvider(t, rawDB, database.AIProvider{Name: providerName, Type: providerType})
			},
			wantInputPrice: sql.NullInt64{Int64: azureInputPrice, Valid: true},
			// 100 input tokens at the azure input price: $0.0005.
			wantCost: sql.NullInt64{Int64: 500, Valid: true},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctx := testutil.Context(t, testutil.WaitLong)
			tc.setupProvider(t, ctx, tc.providerName, tc.providerType)

			intc := dbgen.AIBridgeInterception(t, rawDB, database.InsertAIBridgeInterceptionParams{
				InitiatorID:  user.ID,
				Provider:     tc.wireProvider,
				ProviderName: tc.providerName,
				Model:        tc.model,
			}, nil)

			_, err := srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
				InterceptionId: intc.ID.String(),
				MsgId:          "msg_e2e",
				InputTokens:    100,
				CreatedAt:      timestamppb.Now(),
			})
			require.NoError(t, err, "record token usage")

			tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID)
			require.NoError(t, err)
			require.Len(t, tokenUsages, 1)
			require.Equal(t, tc.wantInputPrice, tokenUsages[0].InputPriceMicros, "input price")
			require.Equal(t, tc.wantCost, tokenUsages[0].CostMicros, "cost")
		})
	}
}

// TestRecordTokenUsageBudgetNotifications verifies that recording token usage
// enqueues the right budget notifications: the warning template when spend
// crosses the warning threshold, the limit-reached template at 100%, both when
// a single interception crosses both, and nothing when no threshold is
// freshly crossed.
func TestRecordTokenUsageBudgetNotifications(t *testing.T) {
	t.Parallel()

	const (
		dollar          int64 = 1_000_000    // micros per USD dollar
		spendLimit      int64 = 100 * dollar // $100 limit (100% threshold)
		warnAt          int64 = 85 * dollar  // $85 (85% of the limit)
		inputPrice      int64 = dollar       // $1 per million tokens
		tokensPerDollar int64 = 1_000_000    // 1,000,000 tokens = $1 at the price above
	)
	now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)

	// Each case sets its cost via inputTokens (tokensPerDollar tokens = $1).
	// newSpend is the post-increment period total; the code derives the
	// pre-increment total as newSpend - cost and fires a threshold only when the
	// pre-increment total is below the threshold and the post-increment total is
	// at or above it.
	price := &database.AIModelPrice{
		InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true},
	}

	// Threshold percentage label expected for each template.
	wantThreshold := map[uuid.UUID]string{
		notifications.TemplateAIBudgetWarningUser:      "85",
		notifications.TemplateAIBudgetLimitReachedUser: "100",
	}

	testCases := []struct {
		name          string
		inputTokens   int64
		newSpend      int64 // post-increment period total
		wantTemplates []uuid.UUID
	}{
		{
			name: "crosses warning threshold",
			// pre = $84.50 (< $85), post = $85.50 (>= $85) -> warning.
			inputTokens:   tokensPerDollar,
			newSpend:      warnAt + dollar/2,
			wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetWarningUser},
		},
		{
			name: "crosses warning threshold exactly",
			// pre = $84 (< $85), post = $85 (>= $85) -> warning.
			inputTokens:   tokensPerDollar,
			newSpend:      warnAt,
			wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetWarningUser},
		},
		{
			name: "stays below warning threshold",
			// post = $84.50 (< $85) -> no crossing.
			inputTokens:   tokensPerDollar,
			newSpend:      warnAt - dollar/2,
			wantTemplates: nil,
		},
		{
			name: "already at warning threshold",
			// pre = $85 (not < $85), post = $86 -> no fresh crossing.
			inputTokens:   tokensPerDollar,
			newSpend:      warnAt + dollar,
			wantTemplates: nil,
		},
		{
			name: "already above warning threshold",
			// pre = $89 (>= $85, < $100), post = $90 -> no crossing.
			inputTokens:   tokensPerDollar,
			newSpend:      warnAt + 5*dollar,
			wantTemplates: nil,
		},
		{
			name: "crosses limit",
			// pre = $99.50 (>= $85, so no warning; < $100), post = $100.50 (>= $100) -> limit.
			inputTokens:   tokensPerDollar,
			newSpend:      spendLimit + dollar/2,
			wantTemplates: []uuid.UUID{notifications.TemplateAIBudgetLimitReachedUser},
		},
		{
			name: "crosses warning and limit in one interception",
			// pre = $80 (< $85), post = $100 (>= $85 and >= $100) -> warning + limit.
			inputTokens: 20 * tokensPerDollar,
			newSpend:    spendLimit,
			wantTemplates: []uuid.UUID{
				notifications.TemplateAIBudgetWarningUser,
				notifications.TemplateAIBudgetLimitReachedUser,
			},
		},
		{
			name: "already above limit",
			// pre = $109 (>= $100), post = $110 -> no fresh crossing.
			inputTokens:   tokensPerDollar,
			newSpend:      spendLimit + 10*dollar,
			wantTemplates: nil,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			enq := &notificationstest.FakeEnqueuer{}

			intc := newTestInterception(uuid.New())
			groupID := uuid.New()
			group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit}

			expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

			db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
				func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
			)
			db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
				Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
			db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
				Return(database.AIUserDailySpend{}, nil)
			db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
				Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil)
			// The group, organization, and user are resolved once per
			// interception that notifies, regardless of how many thresholds it
			// crosses.
			if len(tc.wantTemplates) > 0 {
				orgID := uuid.New()
				db.EXPECT().GetGroupByID(gomock.Any(), groupID).
					Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil)
				db.EXPECT().GetOrganizationByID(gomock.Any(), orgID).
					Return(database.Organization{ID: orgID, Name: "coder"}, nil)
				db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
					Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
				// No admins configured, so only the user is notified.
				db.EXPECT().GetUsers(gomock.Any(), gomock.Any()).Return(nil, nil)
			}

			ctx := testutil.Context(t, testutil.WaitLong)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Enqueuer:      enq,
				Logger:        testutil.Logger(t),
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
				InterceptionId: intc.ID.String(),
				MsgId:          "msg_123",
				InputTokens:    tc.inputTokens,
				CreatedAt:      timestamppb.New(now),
			})
			require.NoError(t, err)

			require.Len(t, enq.Sent(), len(tc.wantTemplates), "unexpected number of notifications")
			for _, tmpl := range tc.wantTemplates {
				sent := enq.Sent(notificationstest.WithTemplateID(tmpl))
				require.Len(t, sent, 1, "expected one notification for template %s", tmpl)
				require.Equal(t, intc.InitiatorID, sent[0].UserID)
				require.Equal(t, wantThreshold[tmpl], sent[0].Labels["threshold"])
				require.Equal(t, "$100.00", sent[0].Labels["limit"])
				require.Equal(t, "Engineering", sent[0].Labels["effective_group_name"])
				// The interception is recorded at 2026-06-25, so its budget
				// period runs June 1 - July 1, 2026.
				require.Equal(t, "June 1, 2026", sent[0].Labels["period_start"])
				require.Equal(t, "July 1, 2026", sent[0].Labels["period_end"])
			}
		})
	}
}

// TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary verifies that an
// interception created in one budget period but processed after the period has
// rolled over is still evaluated against the period it belongs to, so a genuine
// threshold crossing is detected rather than lost across the boundary.
func TestRecordTokenUsageBudgetNotificationAcrossPeriodBoundary(t *testing.T) {
	t.Parallel()

	const (
		dollar          int64 = 1_000_000    // micros per USD dollar
		spendLimit      int64 = 100 * dollar // $100 limit (100% threshold)
		warnAt          int64 = 85 * dollar  // $85 (85% of the limit)
		inputPrice      int64 = dollar       // $1 per million tokens
		tokensPerDollar int64 = 1_000_000    // 1,000,000 tokens = $1 at the price above
	)

	// The interception was created in the final second of January but is
	// processed just after the rollover into February. The spend is bucketed
	// into January, so detection must sum against January's period.
	createdAt := time.Date(2026, 1, 31, 23, 59, 59, 0, time.UTC)
	processedAt := time.Date(2026, 2, 1, 0, 0, 1, 0, time.UTC)
	januaryStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)

	ctrl := gomock.NewController(t)
	db := dbmock.NewMockStore(ctrl)
	enq := &notificationstest.FakeEnqueuer{}

	intc := newTestInterception(uuid.New())
	groupID := uuid.New()
	group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit}
	price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true}}

	expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

	db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
		func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
	)
	db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
		Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
	db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
		Return(database.AIUserDailySpend{}, nil)

	// The spend query must run against the period the interception belongs to
	// (January), not the period it was processed in (February).
	var gotPeriodStart time.Time
	db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
		DoAndReturn(func(_ context.Context, p database.GetUserAISpendSinceParams) (database.GetUserAISpendSinceRow, error) {
			gotPeriodStart = p.PeriodStart
			return database.GetUserAISpendSinceRow{SpendMicros: warnAt}, nil
		})
	orgID := uuid.New()
	db.EXPECT().GetGroupByID(gomock.Any(), groupID).
		Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil)
	db.EXPECT().GetOrganizationByID(gomock.Any(), orgID).
		Return(database.Organization{ID: orgID, Name: "coder"}, nil)
	db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
		Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
	// No admins configured, so only the user is notified.
	db.EXPECT().GetUsers(gomock.Any(), gomock.Any()).Return(nil, nil)

	clock := quartz.NewMock(t)
	clock.Set(processedAt)

	ctx := testutil.Context(t, testutil.WaitLong)
	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Enqueuer:      enq,
		Logger:        testutil.Logger(t),
		Clock:         clock,
	})
	require.NoError(t, err)

	_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
		InterceptionId: intc.ID.String(),
		MsgId:          "msg_boundary",
		InputTokens:    tokensPerDollar,
		CreatedAt:      timestamppb.New(createdAt),
	})
	require.NoError(t, err)

	require.Equal(t, januaryStart, gotPeriodStart,
		"spend must be summed against the period the interception belongs to, not the processing period")

	sent := enq.Sent(notificationstest.WithTemplateID(notifications.TemplateAIBudgetWarningUser))
	require.Len(t, sent, 1, "expected the crossing to be detected against the interception's period")
}

// TestRecordTokenUsageBudgetNotificationBestEffort verifies that a failure while
// detecting or sending a budget notification is swallowed: the token usage and
// spend are still recorded (RecordTokenUsage returns no error) and no
// notification is enqueued. This guards the best-effort contract, e.g. that a
// detection error is not propagated out of the transaction (which would roll
// back the committed spend).
func TestRecordTokenUsageBudgetNotificationBestEffort(t *testing.T) {
	t.Parallel()

	const (
		dollar          int64 = 1_000_000
		spendLimit      int64 = 100 * dollar
		warnAt          int64 = 85 * dollar
		inputPrice      int64 = dollar
		tokensPerDollar int64 = 1_000_000
	)
	now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)
	price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true}}

	testCases := []struct {
		name string
		// spendSinceErr fails detection (the read inside the transaction);
		// groupLookupErr fails the notification after a crossing is detected.
		spendSinceErr  error
		groupLookupErr error
	}{
		{name: "detection read fails", spendSinceErr: sql.ErrConnDone},
		{name: "group lookup fails", groupLookupErr: sql.ErrConnDone},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			enq := &notificationstest.FakeEnqueuer{}

			intc := newTestInterception(uuid.New())
			groupID := uuid.New()
			group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit}

			expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

			db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
				func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
			)
			// The token usage and spend are recorded before detection runs.
			db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
				Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
			db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
				Return(database.AIUserDailySpend{}, nil)

			switch {
			case tc.spendSinceErr != nil:
				// Detection fails; the group is never looked up.
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{}, tc.spendSinceErr)
			case tc.groupLookupErr != nil:
				// A crossing is detected ($84 -> $85), but resolving the group
				// for the notification fails.
				db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
					Return(database.GetUserAISpendSinceRow{SpendMicros: warnAt}, nil)
				db.EXPECT().GetGroupByID(gomock.Any(), groupID).
					Return(database.Group{}, tc.groupLookupErr)
			default:
				t.Fatal("test case must set spendSinceErr or groupLookupErr")
			}

			ctx := testutil.Context(t, testutil.WaitLong)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Enqueuer:      enq,
				// The detect/notify failure is logged; ignore it here since
				// triggering it is the point of the test.
				Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
				Clock:  quartz.NewReal(),
			})
			require.NoError(t, err)

			// The failure must not surface as an error from RecordTokenUsage.
			_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
				InterceptionId: intc.ID.String(),
				MsgId:          "msg_123",
				InputTokens:    tokensPerDollar,
				CreatedAt:      timestamppb.New(now),
			})
			require.NoError(t, err)
			require.Empty(t, enq.Sent(), "no notification should be enqueued when detection or lookup fails")
		})
	}
}

// TestRecordTokenUsageBudgetNotificationZeroLimit verifies that a zero spend
// limit (used to block a group entirely) produces no budget notification: there
// is no meaningful threshold to cross, and such users are already blocked by
// pre-request enforcement. The token usage is still recorded.
func TestRecordTokenUsageBudgetNotificationZeroLimit(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	db := dbmock.NewMockStore(ctrl)
	enq := &notificationstest.FakeEnqueuer{}

	intc := newTestInterception(uuid.New())
	groupID := uuid.New()
	// A zero limit blocks the group; there is no threshold to cross.
	group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 0}
	price := &database.AIModelPrice{InputPrice: sql.NullInt64{Int64: 1_000_000, Valid: true}}

	expectTokenUsageCostLookups(db, intc, nil, group, nil, price)

	db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
		func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
	)
	// The spend is still recorded; detection then short-circuits on the zero
	// limit without reading spend or looking up the group.
	db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
		Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
	db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
		Return(database.AIUserDailySpend{}, nil)

	ctx := testutil.Context(t, testutil.WaitLong)
	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Enqueuer:      enq,
		Logger:        testutil.Logger(t),
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
		InterceptionId: intc.ID.String(),
		MsgId:          "msg_123",
		InputTokens:    1_000_000, // $1 at $1 per million tokens
		CreatedAt:      timestamppb.New(time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)),
	})
	require.NoError(t, err)
	require.Empty(t, enq.Sent(), "a zero spend limit must not produce a notification")
}

// TestRecordTokenUsageBudgetAdminNotification verifies that crossing a budget
// threshold also notifies deployment owners and user admins, and that the
// affected user is not double-notified as an admin.
func TestRecordTokenUsageBudgetAdminNotification(t *testing.T) {
	t.Parallel()

	const (
		dollar          int64 = 1_000_000    // micros per USD dollar
		spendLimit      int64 = 100 * dollar // $100 limit (100% threshold)
		warnAt          int64 = 85 * dollar  // $85 (85% of the limit)
		inputPrice      int64 = dollar       // $1 per million tokens
		tokensPerDollar int64 = 1_000_000    // 1,000,000 tokens = $1 at the price above
	)
	now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC)

	cases := []struct {
		name              string
		useOverride       bool
		inputTokens       int64
		newSpend          int64
		wantThreshold     string
		wantLimitSource   string
		wantUserTemplate  uuid.UUID
		wantAdminTemplate uuid.UUID
	}{
		{
			// A $1 interception takes spend to warnAt ($85): pre ($84) < warnAt <= post.
			name:              "warning",
			inputTokens:       tokensPerDollar,
			newSpend:          warnAt,
			wantThreshold:     "85",
			wantLimitSource:   string(codersdk.AIBudgetLimitSourceGroup),
			wantUserTemplate:  notifications.TemplateAIBudgetWarningUser,
			wantAdminTemplate: notifications.TemplateAIBudgetWarningAdmin,
		},
		{
			// A $5 interception takes spend from $95 to the $100 limit, so the limit threshold is crossed.
			name:              "limit reached",
			inputTokens:       5 * tokensPerDollar,
			newSpend:          spendLimit,
			wantThreshold:     "100",
			wantLimitSource:   string(codersdk.AIBudgetLimitSourceGroup),
			wantUserTemplate:  notifications.TemplateAIBudgetLimitReachedUser,
			wantAdminTemplate: notifications.TemplateAIBudgetLimitReachedAdmin,
		},
		{
			// A per-user override supplies the limit, so limit_source is user_override.
			name:              "warning, per-user override",
			useOverride:       true,
			inputTokens:       tokensPerDollar,
			newSpend:          warnAt,
			wantThreshold:     "85",
			wantLimitSource:   string(codersdk.AIBudgetLimitSourceUserOverride),
			wantUserTemplate:  notifications.TemplateAIBudgetWarningUser,
			wantAdminTemplate: notifications.TemplateAIBudgetWarningAdmin,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			price := &database.AIModelPrice{
				InputPrice: sql.NullInt64{Int64: inputPrice, Valid: true},
			}

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			enq := &notificationstest.FakeEnqueuer{}

			intc := newTestInterception(uuid.New())
			groupID := uuid.New()

			admin := database.GetUsersRow{ID: uuid.New(), Username: "admin1"}
			// The affected user is also an admin; they must not receive the admin copy.
			selfAdmin := database.GetUsersRow{ID: intc.InitiatorID, Username: "bob"}

			// Whether the limit comes from a group budget or a per-user override, the
			// spend is attributed to the same effective group, so the group lookup and
			// notifications are identical apart from the limit_source label.
			if tc.useOverride {
				override := &database.GetUserAIBudgetOverrideWithOrganizationRow{GroupID: groupID, SpendLimitMicros: spendLimit}
				expectTokenUsageCostLookups(db, intc, override, nil, nil, price)
			} else {
				group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: spendLimit}
				expectTokenUsageCostLookups(db, intc, nil, group, nil, price)
			}
			db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
				func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
			)
			db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).
				Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: intc.ID}, nil)
			db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).
				Return(database.AIUserDailySpend{}, nil)
			db.EXPECT().GetUserAISpendSince(gomock.Any(), gomock.Any()).
				Return(database.GetUserAISpendSinceRow{SpendMicros: tc.newSpend}, nil)
			orgID := uuid.New()
			db.EXPECT().GetGroupByID(gomock.Any(), groupID).
				Return(database.Group{ID: groupID, Name: "Engineering", OrganizationID: orgID}, nil)
			db.EXPECT().GetOrganizationByID(gomock.Any(), orgID).
				Return(database.Organization{ID: orgID, Name: "coder"}, nil)
			db.EXPECT().GetUserByID(gomock.Any(), intc.InitiatorID).
				Return(database.User{ID: intc.InitiatorID, Username: "bob"}, nil)
			db.EXPECT().GetUsers(gomock.Any(), database.GetUsersParams{
				RbacRole: []string{codersdk.RoleOwner, codersdk.RoleUserAdmin},
			}).Return([]database.GetUsersRow{admin, selfAdmin}, nil)

			ctx := testutil.Context(t, testutil.WaitLong)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Enqueuer:      enq,
				Logger:        testutil.Logger(t),
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			_, err = srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
				InterceptionId: intc.ID.String(),
				MsgId:          "msg_123",
				InputTokens:    tc.inputTokens,
				CreatedAt:      timestamppb.New(now),
			})
			require.NoError(t, err)

			// The user who crossed the threshold gets the user-facing notification.
			userSent := enq.Sent(notificationstest.WithTemplateID(tc.wantUserTemplate))
			require.Len(t, userSent, 1)
			require.Equal(t, intc.InitiatorID, userSent[0].UserID)

			// The admin (but not the affected user, who is also an admin) gets the
			// admin notification naming the affected user.
			adminSent := enq.Sent(notificationstest.WithTemplateID(tc.wantAdminTemplate))
			require.Len(t, adminSent, 1)
			require.Equal(t, admin.ID, adminSent[0].UserID)
			require.Equal(t, "bob", adminSent[0].Labels["username"])
			require.Equal(t, tc.wantThreshold, adminSent[0].Labels["threshold"])
			require.Equal(t, "$100.00", adminSent[0].Labels["limit"])
			require.Equal(t, "Engineering", adminSent[0].Labels["effective_group_name"])
			require.Equal(t, "coder", adminSent[0].Labels["organization_name"])
			require.Equal(t, tc.wantLimitSource, adminSent[0].Labels["limit_source"])
		})
	}
}

// newTestInterception returns an interception with a fixed initiator, provider,
// and model for cost-attribution test setup. The provider name intentionally
// differs from the upstream wire format.
func newTestInterception(id uuid.UUID) database.AIBridgeInterception {
	return database.AIBridgeInterception{
		ID:           id,
		InitiatorID:  uuid.New(),
		Provider:     "anthropic",
		ProviderName: "anthropic-eu",
		Model:        "claude-sonnet-4-6",
	}
}

// siteWideOverrideParams, siteWideHighestParams and siteWideFallbackParams
// build site-wide budget lookup params: the user ID and a null organization.
func siteWideOverrideParams(userID uuid.UUID) database.GetUserAIBudgetOverrideWithOrganizationParams {
	return database.GetUserAIBudgetOverrideWithOrganizationParams{UserID: userID, OrganizationID: uuid.NullUUID{}}
}

func siteWideHighestParams(userID uuid.UUID) database.GetHighestGroupAIBudgetByUserParams {
	return database.GetHighestGroupAIBudgetByUserParams{UserID: userID, OrganizationID: uuid.NullUUID{}}
}

func siteWideFallbackParams(userID uuid.UUID) database.GetUserEveryoneFallbackGroupParams {
	return database.GetUserEveryoneFallbackGroupParams{UserID: userID, OrganizationID: uuid.NullUUID{}}
}

// expectSiteWideEveryoneBudgetFallback mocks the site-wide effective group lookups
// IsAuthorized makes for a request without a workspace: no override, no group
// budget, and the Everyone group of orgID. The Everyone group's ID equals its
// organization's ID.
func expectSiteWideEveryoneBudgetFallback(db *dbmock.MockStore, userID, orgID uuid.UUID) {
	db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(userID)).
		Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
	db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(userID)).
		Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
	db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(userID)).Return(orgID, nil)
}

// expectTokenUsageCostLookups mocks the store lookups made by resolveTokenUsageCost
// (budget resolution, provider resolution, and the price lookup). A nil override, group,
// everyoneGroupID, or price makes that lookup return sql.ErrNoRows. Budget resolution
// mirrors production code: a non-nil override wins and skips the group lookup, and the
// Everyone fallback is consulted only when both override and group are nil. The provider
// name resolves to a provider whose configured provider type equals the interception's
// upstream wire format.
func expectTokenUsageCostLookups(
	db *dbmock.MockStore,
	intc database.AIBridgeInterception,
	override *database.GetUserAIBudgetOverrideWithOrganizationRow,
	group *database.GetHighestGroupAIBudgetByUserRow,
	everyoneGroupID *uuid.UUID,
	price *database.AIModelPrice,
) {
	db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intc.ID).Return(intc, nil)

	if override != nil {
		db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(intc.InitiatorID)).Return(*override, nil)
	} else {
		db.EXPECT().GetUserAIBudgetOverrideWithOrganization(gomock.Any(), siteWideOverrideParams(intc.InitiatorID)).
			Return(database.GetUserAIBudgetOverrideWithOrganizationRow{}, sql.ErrNoRows)
		if group != nil {
			db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(intc.InitiatorID)).Return(*group, nil)
		} else {
			db.EXPECT().GetHighestGroupAIBudgetByUser(gomock.Any(), siteWideHighestParams(intc.InitiatorID)).
				Return(database.GetHighestGroupAIBudgetByUserRow{}, sql.ErrNoRows)
			if everyoneGroupID != nil {
				db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(intc.InitiatorID)).
					Return(*everyoneGroupID, nil)
			} else {
				db.EXPECT().GetUserEveryoneFallbackGroup(gomock.Any(), siteWideFallbackParams(intc.InitiatorID)).
					Return(uuid.Nil, sql.ErrNoRows)
			}
		}
	}

	providerID := uuid.New()
	db.EXPECT().GetAIProviderByName(gomock.Any(), intc.ProviderName).Return(database.AIProvider{
		ID:   providerID,
		Name: intc.ProviderName,
		Type: database.AIProviderType(intc.Provider),
	}, nil)

	if price != nil {
		db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), database.GetAIModelPriceForProviderParams{
			ProviderType:   intc.Provider,
			Model:          intc.Model,
			OrganizationID: intc.OrganizationID,
			ProviderID:     providerID,
		}).Return(database.GetAIModelPriceForProviderRow{
			Model:           intc.Model,
			InputPrice:      price.InputPrice,
			OutputPrice:     price.OutputPrice,
			CacheReadPrice:  price.CacheReadPrice,
			CacheWritePrice: price.CacheWritePrice,
		}, nil)
	} else {
		db.EXPECT().GetAIModelPriceForProvider(gomock.Any(), gomock.Any()).
			Return(database.GetAIModelPriceForProviderRow{}, sql.ErrNoRows)
	}
}

func TestRecordPromptUsage(t *testing.T) {
	t.Parallel()

	var (
		metadataProto = map[string]*anypb.Any{
			"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
		}
		metadataJSON = `{"key":"value"}`
	)

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) {
			return srv.RecordPromptUsage(ctx, req)
		},
		[]testRecordMethodCase[*proto.RecordPromptUsageRequest]{
			{
				name: "valid prompt usage",
				request: &proto.RecordPromptUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					Prompt:         "yo",
					Metadata:       metadataProto,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordPromptUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeUserPromptParams) bool {
						if !assert.NotEqual(t, uuid.Nil, p.ID, "ID") ||
							!assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") ||
							!assert.Equal(t, req.GetMsgId(), p.ProviderResponseID, "provider response ID") ||
							!assert.Equal(t, req.GetPrompt(), p.Prompt, "prompt") ||
							!assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") ||
							!assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") {
							return false
						}
						return true
					})).Return(database.AIBridgeUserPrompt{
						ID:                 uuid.New(),
						InterceptionID:     interceptionID,
						ProviderResponseID: req.GetMsgId(),
						Prompt:             req.GetPrompt(),
						Metadata: pqtype.NullRawMessage{
							RawMessage: json.RawMessage(metadataJSON),
							Valid:      true,
						},
						CreatedAt: req.GetCreatedAt().AsTime(),
					}, nil)
				},
			},
			{
				name: "invalid interception ID",
				request: &proto.RecordPromptUsageRequest{
					InterceptionId: "not-a-uuid",
					MsgId:          "msg_123",
					Prompt:         "yo",
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "failed to parse interception_id",
			},
			{
				name: "database error",
				request: &proto.RecordPromptUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					Prompt:         "yo",
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordPromptUsageRequest) {
					db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Any()).Return(database.AIBridgeUserPrompt{}, sql.ErrConnDone)
				},
				expectedErr: "insert user prompt",
			},
		},
	)
}

func TestRecordToolUsage(t *testing.T) {
	t.Parallel()

	var (
		metadataProto = map[string]*anypb.Any{
			"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_NumberValue{NumberValue: 123.45}}),
		}
		metadataJSON = `{"key":123.45}`
	)

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordToolUsageRequest) (*proto.RecordToolUsageResponse, error) {
			return srv.RecordToolUsage(ctx, req)
		},
		[]testRecordMethodCase[*proto.RecordToolUsageRequest]{
			{
				name: "valid tool usage with all fields",
				request: &proto.RecordToolUsageRequest{
					InterceptionId:  uuid.NewString(),
					MsgId:           "msg_123",
					ToolCallId:      "call_xyz",
					ItemId:          "fc_item_xyz",
					ServerUrl:       new("https://api.example.com"),
					Tool:            "read_file",
					Input:           `{"path": "/etc/hosts"}`,
					Injected:        false,
					InvocationError: new("permission denied"),
					Metadata:        metadataProto,
					CreatedAt:       timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordToolUsageRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					dbServerURL := sql.NullString{}
					if req.ServerUrl != nil {
						dbServerURL.String = *req.ServerUrl
						dbServerURL.Valid = true
					}

					dbInvocationError := sql.NullString{}
					if req.InvocationError != nil {
						dbInvocationError.String = *req.InvocationError
						dbInvocationError.Valid = true
					}

					db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeToolUsageParams) bool {
						if !assert.NotEqual(t, uuid.Nil, p.ID, "ID") ||
							!assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") ||
							!assert.Equal(t, req.GetMsgId(), p.ProviderResponseID, "provider response ID") ||
							!assert.Equal(t, sql.NullString{String: "call_xyz", Valid: true}, p.ProviderToolCallID, "provider tool call ID") ||
							!assert.Equal(t, sql.NullString{String: "fc_item_xyz", Valid: true}, p.ProviderItemID, "provider item ID") ||
							!assert.Equal(t, req.GetTool(), p.Tool, "tool") ||
							!assert.Equal(t, dbServerURL, p.ServerUrl, "server URL") ||
							!assert.Equal(t, req.GetInput(), p.Input, "input") ||
							!assert.Equal(t, req.GetInjected(), p.Injected, "injected") ||
							!assert.Equal(t, dbInvocationError, p.InvocationError, "invocation error") ||
							!assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") ||
							!assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") {
							return false
						}
						return true
					})).Return(database.AIBridgeToolUsage{
						ID:                 uuid.New(),
						InterceptionID:     interceptionID,
						ProviderResponseID: req.GetMsgId(),
						Tool:               req.GetTool(),
						ServerUrl:          dbServerURL,
						Input:              req.GetInput(),
						Injected:           req.GetInjected(),
						InvocationError:    dbInvocationError,
						Metadata: pqtype.NullRawMessage{
							RawMessage: json.RawMessage(metadataJSON),
							Valid:      true,
						},
						CreatedAt: req.GetCreatedAt().AsTime(),
					}, nil)
				},
			},
			{
				name: "invalid interception ID",
				request: &proto.RecordToolUsageRequest{
					InterceptionId: "not-a-uuid",
					MsgId:          "msg_123",
					Tool:           "read_file",
					Input:          `{"path": "/etc/hosts"}`,
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "failed to parse interception_id",
			},
			{
				name: "database error",
				request: &proto.RecordToolUsageRequest{
					InterceptionId: uuid.NewString(),
					MsgId:          "msg_123",
					Tool:           "read_file",
					Input:          `{"path": "/etc/hosts"}`,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordToolUsageRequest) {
					db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeToolUsage{}, sql.ErrConnDone)
				},
				expectedErr: "insert tool usage",
			},
		},
	)
}

func TestRecordModelThought(t *testing.T) {
	t.Parallel()

	var (
		metadataProto = map[string]*anypb.Any{
			"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
		}
		metadataJSON = `{"key":"value"}`
	)

	testRecordMethod(t,
		func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordModelThoughtRequest) (*proto.RecordModelThoughtResponse, error) {
			return srv.RecordModelThought(ctx, req)
		},
		[]testRecordMethodCase[*proto.RecordModelThoughtRequest]{
			{
				name: "valid model thought",
				request: &proto.RecordModelThoughtRequest{
					InterceptionId: uuid.NewString(),
					Content:        "I should list the files.",
					Metadata:       metadataProto,
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordModelThoughtRequest) {
					interceptionID, err := uuid.Parse(req.GetInterceptionId())
					assert.NoError(t, err, "parse interception UUID")

					db.EXPECT().InsertAIBridgeModelThought(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeModelThoughtParams) bool {
						if !assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") ||
							!assert.Equal(t, "I should list the files.", p.Content, "content") ||
							!assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") {
							return false
						}
						return true
					})).Return(database.AIBridgeModelThought{
						InterceptionID: interceptionID,
						Content:        "I should list the files.",
						Metadata: pqtype.NullRawMessage{
							RawMessage: json.RawMessage(metadataJSON),
							Valid:      true,
						},
					}, nil)
				},
			},
			{
				name: "invalid interception ID",
				request: &proto.RecordModelThoughtRequest{
					InterceptionId: "not-a-uuid",
					Content:        "thinking...",
					CreatedAt:      timestamppb.Now(),
				},
				expectedErr: "failed to parse interception_id",
			},
			{
				name: "database error",
				request: &proto.RecordModelThoughtRequest{
					InterceptionId: uuid.NewString(),
					Content:        "thinking...",
					CreatedAt:      timestamppb.Now(),
				},
				setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordModelThoughtRequest) {
					db.EXPECT().InsertAIBridgeModelThought(gomock.Any(), gomock.Any()).Return(database.AIBridgeModelThought{}, sql.ErrConnDone)
				},
				expectedErr: "insert model thought",
			},
		},
	)
}

type testRecordMethodCase[Req any] struct {
	name    string
	request Req
	// setupMocks is called with the mock store and the above request.
	setupMocks  func(t *testing.T, db *dbmock.MockStore, req Req)
	expectedErr string
	// assertMetrics, when set, is called after the method returns to assert
	// the metrics recorded on the server's registry.
	assertMetrics func(t *testing.T, reg *prometheus.Registry)
	// expectErrorLog tolerates ERROR-level logs, which slogtest otherwise
	// treats as a test failure. Set it only for cases whose expected behavior
	// includes logging an error, so every other case stays strict.
	expectErrorLog bool
}

// testRecordMethod is a helper that abstracts the common testing pattern for all Record* methods.
func testRecordMethod[Req any, Resp any](
	t *testing.T,
	callMethod func(srv *aibridgedserver.Server, ctx context.Context, req Req) (Resp, error),
	cases []testRecordMethodCase[Req],
) {
	t.Helper()

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			logger := testutil.Logger(t)
			if tc.expectErrorLog {
				logger = slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
			}

			if tc.setupMocks != nil {
				tc.setupMocks(t, db, tc.request)
			}

			ctx := testutil.Context(t, testutil.WaitLong)
			reg := prometheus.NewRegistry()
			metrics := aibridgedserver.NewMetrics(reg)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   requiredExperiments,
				Logger:        logger,
				Clock:         quartz.NewReal(),
				Metrics:       metrics,
			})
			require.NoError(t, err)

			resp, err := callMethod(srv, ctx, tc.request)
			if tc.expectedErr != "" {
				require.Error(t, err, "Expected error for test case: %s", tc.name)
				require.Contains(t, err.Error(), tc.expectedErr)
			} else {
				require.NoError(t, err, "Unexpected error for test case: %s", tc.name)
				require.NotNil(t, resp)
			}
			if tc.assertMetrics != nil {
				tc.assertMetrics(t, reg)
			}
		})
	}
}

// Helper functions.
func mustMarshalAny(t *testing.T, msg protobufproto.Message) *anypb.Any {
	t.Helper()
	v, err := anypb.New(msg)
	require.NoError(t, err)
	return v
}

// logLine represents a parsed JSON log entry.
type logLine struct {
	Msg    string         `json:"msg"`
	Level  string         `json:"level"`
	Fields map[string]any `json:"fields"`
}

// parseLogLines parses JSON log lines from a buffer.
func parseLogLines(buf *bytes.Buffer) []logLine {
	var lines []logLine
	scanner := bufio.NewScanner(buf)
	for scanner.Scan() {
		var line logLine
		if err := json.Unmarshal(scanner.Bytes(), &line); err == nil {
			lines = append(lines, line)
		}
	}
	return lines
}

// getLogLinesWithMessage returns all log lines with the given message.
func getLogLinesWithMessage(lines []logLine, msg string) []logLine {
	var result []logLine
	for _, line := range lines {
		if line.Msg == msg {
			result = append(result, line)
		}
	}
	return result
}

func TestStructuredLogging(t *testing.T) {
	t.Parallel()

	metadataProto := map[string]*anypb.Any{
		"key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}),
	}

	type testCase struct {
		name              string
		structuredLogging bool
		expectedErr       error
		setupMocks        func(db *dbmock.MockStore, interceptionID uuid.UUID)
		recordFn          func(srv *aibridgedserver.Server, ctx context.Context, interceptionID uuid.UUID) error
		expectedFields    map[string]any
	}

	interceptionID := uuid.UUID{1}
	initiatorID := uuid.UUID{2}
	threadParentID := uuid.UUID{3}
	threadRootID := uuid.UUID{4}

	toolCallID := "my-tool-call"
	sessionID := "some-session-id"
	organizationID := uuid.UUID{5}

	cases := []testCase{
		{
			name:              "RecordInterception_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(gomock.Any(), toolCallID).Return(database.GetAIBridgeInterceptionLineageByToolCallIDRow{
					ThreadParentID: threadParentID,
					ThreadRootID:   threadRootID,
				}, nil)

				db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
					ID:             intcID,
					InitiatorID:    initiatorID,
					ThreadParentID: uuid.NullUUID{UUID: threadParentID, Valid: true},
					ThreadRootID:   uuid.NullUUID{UUID: threadRootID, Valid: true},
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
					Id:                    intcID.String(),
					ApiKeyId:              "api-key-123",
					InitiatorId:           initiatorID.String(),
					Provider:              "anthropic",
					Model:                 "claude-4-opus",
					Metadata:              metadataProto,
					StartedAt:             timestamppb.Now(),
					CorrelatingToolCallId: new(toolCallID),
					ClientSessionId:       new(sessionID),
					OrganizationId:        organizationID.String(),
					OrganizationSource:    string(database.AibridgeInterceptionOrganizationSourceEffectiveGroup),
				})

				return err
			},
			expectedFields: map[string]any{
				"record_type":              "interception_start",
				"interception_id":          interceptionID.String(),
				"initiator_id":             initiatorID.String(),
				"provider":                 "anthropic",
				"model":                    "claude-4-opus",
				"correlating_tool_call_id": toolCallID,
				"thread_parent_id":         threadParentID.String(),
				"thread_root_id":           threadRootID.String(),
				"client_session_id":        sessionID,
				"organization_id":          organizationID.String(),
				"organization_source":      "effective_group",
			},
		},
		{
			name:              "RecordInterception_does_not_log_when_disabled",
			structuredLogging: false,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
					ID:          intcID,
					InitiatorID: initiatorID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
					Id:          intcID.String(),
					ApiKeyId:    "api-key-123",
					InitiatorId: initiatorID.String(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					StartedAt:   timestamppb.Now(),
				})
				return err
			},
			expectedFields: nil, // No log expected.
		},
		{
			name:              "RecordInterception_log_on_db_error",
			structuredLogging: true,
			expectedErr:       sql.ErrConnDone,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{}, sql.ErrConnDone)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
					Id:          intcID.String(),
					ApiKeyId:    "api-key-123",
					InitiatorId: initiatorID.String(),
					Provider:    "anthropic",
					Model:       "claude-4-opus",
					StartedAt:   timestamppb.Now(),
				})
				return err
			},
			// Even though the database call errored, we must still write the logs.
			expectedFields: map[string]any{
				"record_type":     "interception_start",
				"interception_id": interceptionID.String(),
				"initiator_id":    initiatorID.String(),
				"provider":        "anthropic",
				"model":           "claude-4-opus",
				// Organization fields are omitted when no organization was resolved.
				"organization_id":     nil,
				"organization_source": nil,
			},
		},
		{
			name:              "RecordInterceptionEnded_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{
					ID: intcID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordInterceptionEnded(ctx, &proto.RecordInterceptionEndedRequest{
					Id:      intcID.String(),
					EndedAt: timestamppb.Now(),
				})
				return err
			},
			expectedFields: map[string]any{
				"record_type":     "interception_end",
				"interception_id": interceptionID.String(),
			},
		},
		{
			name:              "RecordTokenUsage_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				everyoneID := uuid.New()
				expectTokenUsageCostLookups(db, newTestInterception(intcID), nil, nil, &everyoneID, nil)
				db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
					func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) },
				)
				db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{
					ID:             uuid.New(),
					InterceptionID: intcID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordTokenUsage(ctx, &proto.RecordTokenUsageRequest{
					InterceptionId:        intcID.String(),
					MsgId:                 "msg_123",
					InputTokens:           100,
					OutputTokens:          200,
					CacheReadInputTokens:  50,
					CacheWriteInputTokens: 10,
					Metadata:              metadataProto,
					CreatedAt:             timestamppb.Now(),
				})
				return err
			},
			expectedFields: map[string]any{
				"record_type":              "token_usage",
				"interception_id":          interceptionID.String(),
				"input_tokens":             float64(100), // JSON numbers are float64.
				"output_tokens":            float64(200),
				"cache_read_input_tokens":  float64(50),
				"cache_write_input_tokens": float64(10),
			},
		},
		{
			name:              "RecordPromptUsage_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Any()).Return(database.AIBridgeUserPrompt{
					ID:             uuid.New(),
					InterceptionID: intcID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{
					InterceptionId: intcID.String(),
					MsgId:          "msg_123",
					Prompt:         "Hello, Claude!",
					Metadata:       metadataProto,
					CreatedAt:      timestamppb.Now(),
				})
				return err
			},
			expectedFields: map[string]any{
				"record_type":     "prompt_usage",
				"interception_id": interceptionID.String(),
				"prompt":          "Hello, Claude!",
			},
		},
		{
			name:              "RecordToolUsage_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeToolUsage{
					ID:             uuid.New(),
					InterceptionID: intcID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
					InterceptionId:  intcID.String(),
					MsgId:           "msg_123",
					ServerUrl:       new("https://api.example.com"),
					Tool:            "read_file",
					Input:           `{"path": "/etc/hosts"}`,
					Injected:        true,
					InvocationError: new("permission denied"),
					Metadata:        metadataProto,
					CreatedAt:       timestamppb.Now(),
				})
				return err
			},
			expectedFields: map[string]any{
				"record_type":      "tool_usage",
				"interception_id":  interceptionID.String(),
				"tool":             "read_file",
				"input":            `{"path": "/etc/hosts"}`,
				"injected":         true,
				"invocation_error": "permission denied",
			},
		},
		{
			name:              "RecordModelThought_logs_when_enabled",
			structuredLogging: true,
			setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) {
				db.EXPECT().InsertAIBridgeModelThought(gomock.Any(), gomock.Any()).Return(database.AIBridgeModelThought{
					InterceptionID: intcID,
				}, nil)
			},
			recordFn: func(srv *aibridgedserver.Server, ctx context.Context, intcID uuid.UUID) error {
				_, err := srv.RecordModelThought(ctx, &proto.RecordModelThoughtRequest{
					InterceptionId: intcID.String(),
					Content:        "I need to list the files.",
					Metadata:       metadataProto,
					CreatedAt:      timestamppb.Now(),
				})
				return err
			},
			expectedFields: map[string]any{
				"record_type":     "model_thought",
				"interception_id": interceptionID.String(),
				"content":         "I need to list the files.",
			},
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			buf := &bytes.Buffer{}
			logger := slog.Make(slogjson.Sink(buf)).Leveled(slog.LevelDebug)

			tc.setupMocks(db, interceptionID)

			ctx := testutil.Context(t, testutil.WaitLong)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: agplaiseats.Noop{},
				AccessURL:     "/",
				GatewayCfg: codersdk.AIBridgeConfig{
					StructuredLogging: serpent.Bool(tc.structuredLogging),
				},
				Experiments: requiredExperiments,
				Logger:      logger,
				Clock:       quartz.NewReal(),
			})
			require.NoError(t, err)

			err = tc.recordFn(srv, ctx, interceptionID)
			if tc.expectedErr != nil {
				require.Error(t, err)
			} else {
				require.NoError(t, err)
			}

			lines := parseLogLines(buf)
			if tc.expectedFields == nil {
				// No log expected (disabled or error case).
				require.Empty(t, lines)
			} else {
				matchedLines := getLogLinesWithMessage(lines, recorder.InterceptionLogMarker)
				require.GreaterOrEqual(t, len(matchedLines), 1, "expected at least 1 log line(s) with message %q", recorder.InterceptionLogMarker)

				fields := matchedLines[0].Fields
				for key, expected := range tc.expectedFields {
					require.Equal(t, expected, fields[key], "field %q mismatch", key)
				}
			}
		})
	}
}

// TestInferredThreadsByToolCalls verifies that a chain of interceptions linked via
// tool call IDs correctly propagates thread_parent_id and thread_root_id.
//
// The chain is: A → B → C
//   - A is the root (no parent, no root)
//   - B correlates via a tool call recorded by A (parent=A, root=A)
//   - C correlates via a tool call recorded by B (parent=B, root=A)
func TestInferredThreadsByToolCalls(t *testing.T) {
	t.Parallel()
	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	user := dbgen.User(t, db, database.User{})

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	aID := uuid.New()
	bID := uuid.New()
	cID := uuid.New()

	// Record interception A (root of the chain, no correlation).
	_, err = srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
		Id:          aID.String(),
		ApiKeyId:    uuid.NewString(),
		InitiatorId: user.ID.String(),
		Provider:    "anthropic",
		Model:       "claude-4-opus",
		StartedAt:   timestamppb.Now(),
	})
	require.NoError(t, err)

	// No thread association yet.
	intcA, err := db.GetAIBridgeInterceptionByID(ctx, aID)
	require.NoError(t, err)
	require.Equal(t, uuid.NullUUID{}, intcA.ThreadParentID)
	require.Equal(t, uuid.NullUUID{}, intcA.ThreadRootID)

	// Record tool usage on A with a known tool call ID.
	_, err = srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
		InterceptionId: aID.String(),
		MsgId:          "resp_a",
		ToolCallId:     "call_a",
		Tool:           "bash",
		Input:          "{}",
		CreatedAt:      timestamppb.Now(),
	})
	require.NoError(t, err)

	// Record interception B correlating to A's tool call.
	_, err = srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
		Id:                    bID.String(),
		ApiKeyId:              uuid.NewString(),
		InitiatorId:           user.ID.String(),
		Provider:              "anthropic",
		Model:                 "claude-4-opus",
		StartedAt:             timestamppb.Now(),
		CorrelatingToolCallId: new("call_a"),
	})
	require.NoError(t, err)

	intcB, err := db.GetAIBridgeInterceptionByID(ctx, bID)
	require.NoError(t, err)
	require.Equal(t, uuid.NullUUID{UUID: aID, Valid: true}, intcB.ThreadParentID)
	require.Equal(t, uuid.NullUUID{UUID: aID, Valid: true}, intcB.ThreadRootID)

	// Record tool usage on B.
	_, err = srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
		InterceptionId: bID.String(),
		MsgId:          "resp_b",
		ToolCallId:     "call_b",
		Tool:           "bash",
		Input:          "{}",
		CreatedAt:      timestamppb.Now(),
	})
	require.NoError(t, err)

	// Record interception C correlating to B's tool call.
	_, err = srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
		Id:                    cID.String(),
		ApiKeyId:              uuid.NewString(),
		InitiatorId:           user.ID.String(),
		Provider:              "anthropic",
		Model:                 "claude-4-opus",
		StartedAt:             timestamppb.Now(),
		CorrelatingToolCallId: new("call_b"),
	})
	require.NoError(t, err)

	intcC, err := db.GetAIBridgeInterceptionByID(ctx, cID)
	require.NoError(t, err)
	require.Equal(t, uuid.NullUUID{UUID: bID, Valid: true}, intcC.ThreadParentID)
	require.Equal(t, uuid.NullUUID{UUID: aID, Valid: true}, intcC.ThreadRootID)
}

// TestRecordToolUsageProviderItemID exercises the RecordToolUsage RPC against a
// real database and confirms that provider_item_id is persisted in its own
// column for both shapes of Responses-API tool call. Agentic tools carry both
// an item id and a tool_call_id; hosted tools (e.g. web_search_call) carry only
// an item id. The hosted case is the important one: it proves the item id is
// stored even when tool_call_id is absent, so persistence is not gated on the
// tool_call_id being present, and the two ids are written to their own columns.
func TestRecordToolUsageProviderItemID(t *testing.T) {
	t.Parallel()
	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	logger := testutil.Logger(t)

	user := dbgen.User(t, db, database.User{})

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Experiments:   requiredExperiments,
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	intcID := uuid.New()
	_, err = srv.RecordInterception(ctx, &proto.RecordInterceptionRequest{
		Id:          intcID.String(),
		ApiKeyId:    uuid.NewString(),
		InitiatorId: user.ID.String(),
		Provider:    "openai",
		Model:       "gpt-5",
		StartedAt:   timestamppb.Now(),
	})
	require.NoError(t, err)

	// Agentic tool: both item_id and tool_call_id are present.
	_, err = srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
		InterceptionId: intcID.String(),
		MsgId:          "resp_1",
		ToolCallId:     "call_agentic",
		ItemId:         "fc_item_1",
		Tool:           "function_call",
		Input:          "{}",
		CreatedAt:      timestamppb.Now(),
	})
	require.NoError(t, err)

	// Hosted tool: only item_id is present, tool_call_id is empty.
	_, err = srv.RecordToolUsage(ctx, &proto.RecordToolUsageRequest{
		InterceptionId: intcID.String(),
		MsgId:          "resp_1",
		ItemId:         "ws_item_1",
		Tool:           "web_search_call",
		Input:          "{}",
		CreatedAt:      timestamppb.Now(),
	})
	require.NoError(t, err)

	usages, err := db.GetAIBridgeToolUsagesByInterceptionID(ctx, intcID)
	require.NoError(t, err)
	require.Len(t, usages, 2)

	byItemID := make(map[string]database.AIBridgeToolUsage, len(usages))
	for _, u := range usages {
		require.True(t, u.ProviderItemID.Valid, "item ID should be persisted for %q", u.Tool)
		byItemID[u.ProviderItemID.String] = u
	}

	// Agentic tool: item id and tool_call_id land in their own columns.
	agentic, ok := byItemID["fc_item_1"]
	require.True(t, ok, "agentic tool usage persisted by item ID")
	require.Equal(t, sql.NullString{String: "call_agentic", Valid: true}, agentic.ProviderToolCallID)

	// Hosted tool: item id is persisted even though the tool_call_id is empty.
	hosted, ok := byItemID["ws_item_1"]
	require.True(t, ok, "hosted tool usage persisted by item ID")
	require.Equal(t, sql.NullString{}, hosted.ProviderToolCallID, "hosted tool has no tool_call_id")
}

// TestGetAIProviders exercises the row-to-proto mapping over a real database:
// enabled providers carry their keys (and typed Bedrock settings), disabled
// providers are included but withhold keys and settings, Copilot (a keyless
// BYOK provider) round-trips with no keys, and an enabled provider whose
// settings blob cannot be decoded is skipped rather than failing the fetch.
func TestGetAIProviders(t *testing.T) {
	t.Parallel()

	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	// The skipped misconfigured provider is logged at Error level by design,
	// so error logs are expected here.
	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})

	// Enabled OpenAI with two keys.
	openai := dbgen.AIProvider(t, db, database.AIProvider{
		Type:    database.AIProviderTypeOpenai,
		Name:    "openai",
		Enabled: true,
		BaseUrl: "https://api.openai.com/",
	})
	dbgen.AIProviderKey(t, db, database.AIProviderKey{ProviderID: openai.ID, APIKey: "sk-openai-1"})
	dbgen.AIProviderKey(t, db, database.AIProviderKey{ProviderID: openai.ID, APIKey: "sk-openai-2"})

	// Enabled Bedrock with typed settings.
	bedrockSettings, err := json.Marshal(codersdk.AIProviderSettings{
		Bedrock: &codersdk.AIProviderBedrockSettings{
			Region:          "us-east-1",
			Model:           "anthropic.claude-3",
			SmallFastModel:  "anthropic.claude-haiku",
			AccessKey:       new("AKID"),
			AccessKeySecret: new("secret"),
			RoleARN:         "arn:aws:iam::123456789012:role/bedrock",
		},
	})
	require.NoError(t, err)
	dbgen.AIProvider(t, db, database.AIProvider{
		Type:     database.AIProviderTypeBedrock,
		Name:     "bedrock",
		Enabled:  true,
		BaseUrl:  "https://bedrock-runtime.us-east-1.amazonaws.com/",
		Settings: sql.NullString{String: string(bedrockSettings), Valid: true},
	})

	// Enabled Anthropic using Claude Platform for AWS in IAM mode. It has no
	// keys: it authenticates by signing.
	claudePlatformSettings, err := json.Marshal(codersdk.AIProviderSettings{
		ClaudePlatformAWS: &codersdk.AIProviderClaudePlatformAWSSettings{
			Region:      "us-west-2",
			WorkspaceID: "wrkspc_123",
		},
	})
	require.NoError(t, err)
	dbgen.AIProvider(t, db, database.AIProvider{
		Type:     database.AIProviderTypeAnthropic,
		Name:     "claude-platform",
		Enabled:  true,
		BaseUrl:  "https://aws-external-anthropic.us-west-2.api.aws/",
		Settings: sql.NullString{String: string(claudePlatformSettings), Valid: true},
	})

	// Enabled Copilot, which is keyless (BYOK per request).
	dbgen.AIProvider(t, db, database.AIProvider{
		Type:    database.AIProviderTypeCopilot,
		Name:    "copilot",
		Enabled: true,
		BaseUrl: "https://api.githubcopilot.com/",
	})

	// Disabled Anthropic with a key; the key must be withheld.
	disabled := dbgen.AIProvider(t, db, database.AIProvider{
		Type:    database.AIProviderTypeAnthropic,
		Name:    "anthropic-off",
		BaseUrl: "https://api.anthropic.com/",
	}, func(p *database.InsertAIProviderParams) {
		p.Enabled = false
	})
	dbgen.AIProviderKey(t, db, database.AIProviderKey{ProviderID: disabled.ID, APIKey: "sk-secret"})

	// Enabled provider with an undecodable settings blob; it must be skipped
	// so one corrupt row does not break provider config for every gateway.
	dbgen.AIProvider(t, db, database.AIProvider{
		Type:     database.AIProviderTypeBedrock,
		Name:     "broken-settings",
		Enabled:  true,
		BaseUrl:  "https://bedrock-runtime.us-east-1.amazonaws.com/",
		Settings: sql.NullString{String: "{not valid json", Valid: true},
	})

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	resp, err := srv.GetAIProviders(ctx, &proto.GetAIProvidersRequest{})
	require.NoError(t, err)

	byName := make(map[string]*proto.AIProvider, len(resp.GetProviders()))
	for _, p := range resp.GetProviders() {
		byName[p.GetName()] = p
	}
	require.Len(t, byName, 5)
	assert.NotContains(t, byName, "broken-settings", "provider with undecodable settings must be skipped")

	gotOpenAI := byName["openai"]
	require.NotNil(t, gotOpenAI)
	assert.True(t, gotOpenAI.GetEnabled())
	assert.Equal(t, string(database.AIProviderTypeOpenai), gotOpenAI.GetType())
	assert.Equal(t, "https://api.openai.com/", gotOpenAI.GetBaseUrl())
	assert.ElementsMatch(t, []string{"sk-openai-1", "sk-openai-2"}, gotOpenAI.GetKeys())
	assert.Nil(t, gotOpenAI.GetBedrock())

	gotBedrock := byName["bedrock"]
	require.NotNil(t, gotBedrock)
	assert.True(t, gotBedrock.GetEnabled())
	require.NotNil(t, gotBedrock.GetBedrock())
	assert.Equal(t, "us-east-1", gotBedrock.GetBedrock().GetRegion())
	assert.Equal(t, "anthropic.claude-3", gotBedrock.GetBedrock().GetModel())
	assert.Equal(t, "anthropic.claude-haiku", gotBedrock.GetBedrock().GetSmallFastModel())
	assert.Equal(t, "AKID", gotBedrock.GetBedrock().GetAccessKey())
	assert.Equal(t, "secret", gotBedrock.GetBedrock().GetAccessKeySecret())
	assert.Equal(t, "arn:aws:iam::123456789012:role/bedrock", gotBedrock.GetBedrock().GetRoleArn())

	gotClaudePlatform := byName["claude-platform"]
	require.NotNil(t, gotClaudePlatform)
	assert.True(t, gotClaudePlatform.GetEnabled())
	assert.Equal(t, string(database.AIProviderTypeAnthropic), gotClaudePlatform.GetType(),
		"claude platform is an authentication method on anthropic, not a provider type")
	assert.Nil(t, gotClaudePlatform.GetBedrock())
	require.NotNil(t, gotClaudePlatform.GetClaudePlatformAws())
	assert.Equal(t, "us-west-2", gotClaudePlatform.GetClaudePlatformAws().GetRegion())
	assert.Equal(t, "wrkspc_123", gotClaudePlatform.GetClaudePlatformAws().GetWorkspaceId())
	assert.Empty(t, gotClaudePlatform.GetKeys())

	gotCopilot := byName["copilot"]
	require.NotNil(t, gotCopilot)
	assert.True(t, gotCopilot.GetEnabled())
	assert.Empty(t, gotCopilot.GetKeys())

	gotDisabled := byName["anthropic-off"]
	require.NotNil(t, gotDisabled)
	assert.False(t, gotDisabled.GetEnabled())
	assert.Empty(t, gotDisabled.GetKeys(), "keys must be withheld for disabled providers")
	assert.Nil(t, gotDisabled.GetBedrock())
	assert.Nil(t, gotDisabled.GetClaudePlatformAws())
}

// TestWatchAIProviders asserts that the WatchAIProviders handler emits an
// initial signal on subscribe, one signal per AIProvidersChangedChannel publish,
// and returns cleanly when the stream context is canceled.
func TestWatchAIProviders(t *testing.T) {
	t.Parallel()

	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	logger := slogtest.Make(t, nil)
	// In-memory pubsub delivers Publish synchronously for deterministic signals.
	ps := pubsub.NewInMemory()

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		Pubsub:        ps,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	streamCtx, streamCancel := context.WithCancel(ctx)
	defer streamCancel()
	stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)}

	watchErr := make(chan error, 1)
	go func() {
		watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream)
	}()

	// The handler sends an initial signal immediately on subscribe. Draining it
	// before publishing guarantees the next publish is not coalesced into the
	// initial signal.
	testutil.TryReceive(ctx, t, stream.sent)

	require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil))
	testutil.TryReceive(ctx, t, stream.sent)

	require.NoError(t, ps.Publish(coderdpubsub.AIProvidersChangedChannel, nil))
	testutil.TryReceive(ctx, t, stream.sent)

	streamCancel()
	require.NoError(t, testutil.TryReceive(ctx, t, watchErr))
}

// TestWatchAIProvidersSignalsOnDeliveryError asserts that a dropped-message
// delivery error is forwarded as a change signal rather than failing the
// stream, so the gateway reconverges after a pubsub drop.
func TestWatchAIProvidersSignalsOnDeliveryError(t *testing.T) {
	t.Parallel()

	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	logger := slogtest.Make(t, nil)
	ps := &captureListenerPubsub{listenerC: make(chan pubsub.ListenerWithErr, 1)}

	srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
		Store:         db,
		Pubsub:        ps,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	streamCtx, streamCancel := context.WithCancel(ctx)
	defer streamCancel()
	stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)}

	watchErr := make(chan error, 1)
	go func() {
		watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream)
	}()

	// Capture the registered listener and drain the initial subscribe signal so
	// the delivery-error signal that follows is not coalesced into it.
	listener := testutil.TryReceive(ctx, t, ps.listenerC)
	testutil.TryReceive(ctx, t, stream.sent)

	// A delivery error must still produce a signal, exercising the pubsub-error
	// branch of the handler.
	listener(ctx, nil, pubsub.ErrDroppedMessages)
	testutil.TryReceive(ctx, t, stream.sent)

	streamCancel()
	require.NoError(t, testutil.TryReceive(ctx, t, watchErr))
}

// TestWatchAIProvidersStopsOnLifecycleCancel asserts the handler returns when
// the server lifecycle context is canceled even though the stream context
// remains open, so a stream that outlives the server does not leak a goroutine
// on shutdown.
func TestWatchAIProvidersStopsOnLifecycleCancel(t *testing.T) {
	t.Parallel()

	db, _ := dbtestutil.NewDB(t)
	ctx := testutil.Context(t, testutil.WaitLong)
	logger := slogtest.Make(t, nil)
	ps := pubsub.NewInMemory()

	// The lifecycle context is independent of the stream context so it can be
	// canceled while the stream stays open.
	lifecycleCtx, lifecycleCancel := context.WithCancel(ctx)
	defer lifecycleCancel()
	srv, err := aibridgedserver.NewServer(lifecycleCtx, aibridgedserver.Options{
		Store:         db,
		Pubsub:        ps,
		AISeatTracker: agplaiseats.Noop{},
		AccessURL:     "/",
		GatewayCfg:    codersdk.AIBridgeConfig{},
		Logger:        logger,
		Clock:         quartz.NewReal(),
	})
	require.NoError(t, err)

	streamCtx, streamCancel := context.WithCancel(ctx)
	defer streamCancel()
	stream := &fakeWatchProvidersStream{ctx: streamCtx, sent: make(chan struct{}, 16)}

	watchErr := make(chan error, 1)
	go func() {
		watchErr <- srv.WatchAIProviders(&proto.WatchAIProvidersRequest{}, stream)
	}()

	// Drain the initial subscribe signal to confirm the handler is running
	// before the lifecycle is canceled.
	testutil.TryReceive(ctx, t, stream.sent)

	// Canceling only the lifecycle context must stop the handler even though
	// the stream context is still open.
	lifecycleCancel()
	require.NoError(t, testutil.TryReceive(ctx, t, watchErr))
}

var _ pubsub.Pubsub = (*captureListenerPubsub)(nil)

// captureListenerPubsub captures the ListenerWithErr registered via
// SubscribeWithErr so a test can drive delivery (including errors) directly.
type captureListenerPubsub struct {
	listenerC chan pubsub.ListenerWithErr
}

func (*captureListenerPubsub) Subscribe(string, pubsub.Listener) (func(), error) {
	return nil, xerrors.New("Subscribe not implemented")
}

func (p *captureListenerPubsub) SubscribeWithErr(_ string, listener pubsub.ListenerWithErr) (func(), error) {
	p.listenerC <- listener
	return func() {}, nil
}

func (*captureListenerPubsub) Publish(string, []byte) error {
	return xerrors.New("Publish not implemented")
}

func (*captureListenerPubsub) Close() error { return nil }

// fakeWatchProvidersStream is a minimal proto.DRPCProviderConfigurator_WatchAIProvidersStream
// that records Send calls on a channel.
type fakeWatchProvidersStream struct {
	ctx  context.Context
	sent chan struct{}
}

func (s *fakeWatchProvidersStream) Send(*proto.WatchAIProvidersResponse) error {
	select {
	case s.sent <- struct{}{}:
		return nil
	case <-s.ctx.Done():
		return s.ctx.Err()
	}
}

func (s *fakeWatchProvidersStream) Context() context.Context                { return s.ctx }
func (*fakeWatchProvidersStream) MsgSend(drpc.Message, drpc.Encoding) error { return nil }
func (*fakeWatchProvidersStream) MsgRecv(drpc.Message, drpc.Encoding) error { return nil }
func (*fakeWatchProvidersStream) CloseSend() error                          { return nil }
func (*fakeWatchProvidersStream) Close() error                              { return nil }

// countingSeatTracker records the number of RecordUsage calls.
type countingSeatTracker struct {
	calls atomic.Int64
}

func (c *countingSeatTracker) RecordUsage(context.Context, uuid.UUID, agplaiseats.Reason) {
	c.calls.Add(1)
}

// TestRecordInterceptionAISeat verifies that bridge usage claims an AI
// Governance seat only when the seat exclusion experiment is disabled.
func TestRecordInterceptionAISeat(t *testing.T) {
	t.Parallel()

	newRequest := func() *proto.RecordInterceptionRequest {
		return &proto.RecordInterceptionRequest{
			Id:          uuid.NewString(),
			ApiKeyId:    uuid.NewString(),
			InitiatorId: uuid.NewString(),
			Provider:    "anthropic",
			Model:       "claude-4-opus",
			StartedAt:   timestamppb.Now(),
		}
	}

	cases := []struct {
		name          string
		experiments   []codersdk.Experiment
		expectedCalls int64
	}{
		{
			name:          "experiment off records a seat",
			experiments:   requiredExperiments,
			expectedCalls: 1,
		},
		{
			name:          "seat exclusion skips the seat",
			experiments:   append([]codersdk.Experiment{codersdk.ExperimentAIGatewaySeatExclusion}, requiredExperiments...),
			expectedCalls: 0,
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			db := dbmock.NewMockStore(ctrl)
			db.EXPECT().InsertAIBridgeInterception(gomock.Any(), gomock.Any()).
				Return(database.AIBridgeInterception{}, nil)

			tracker := &countingSeatTracker{}
			ctx := testutil.Context(t, testutil.WaitLong)
			srv, err := aibridgedserver.NewServer(ctx, aibridgedserver.Options{
				Store:         db,
				AISeatTracker: tracker,
				AccessURL:     "/",
				GatewayCfg:    codersdk.AIBridgeConfig{},
				Experiments:   tc.experiments,
				Logger:        testutil.Logger(t),
				Clock:         quartz.NewReal(),
			})
			require.NoError(t, err)

			_, err = srv.RecordInterception(ctx, newRequest())
			require.NoError(t, err)
			require.Equal(t, tc.expectedCalls, tracker.calls.Load())
		})
	}
}