Function TestServeHTTP_DelegatedAPIKey

coderd/aibridged/aibridged_test.go:184–306

When the request context carries a delegated API key ID (set by the in-process transport on behalf of a trusted caller like chatd), the handler must authenticate via the key_id field, skipping the header-based key extraction entirely. Validation succeeds or fails exactly as it would for a real API k

(t *testing.T)

// LLM credentials must still be forwarded upstream while the Coder governance
// token is stripped.
func TestServeHTTP_DelegatedAPIKey(t *testing.T) {
	t.Parallel()

	const testKeyID = "abcdef1234"

	tests := []struct {
		name          string
		reqHeaders    map[string]string
		applyMocks    func(t *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler)
		expectStatus  int
		expectHandled bool
		expectPresent map[string]string
		expectAbsent  []string
	}{
		{
			// Delegated + centralized: identity comes from the
			// api key ID on the context, in lieu of a session
			// token. No header credentials are sent and SessionKey
			// is empty downstream.
			name: "valid centralized",
			applyMocks: func(t *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler) {
				client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).DoAndReturn(
					func(_ context.Context, in *proto.IsAuthorizedRequest) (*proto.IsAuthorizedResponse, error) {
						assert.Equal(t, testKeyID, in.GetKeyId(), "handler must use KeyId for delegated requests")
						assert.Empty(t, in.GetKey(), "handler must not set Key for delegated requests")
						return &proto.IsAuthorizedResponse{
							OwnerId:  uuid.NewString(),
							ApiKeyId: testKeyID,
							Username: "u",
						}, nil
					})
				pool.EXPECT().Acquire(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
					func(_ context.Context, req aibridged.Request, _ aibridged.ClientFunc, _ aibridged.MCPProxyBuilder) (http.Handler, error) {
						assert.Empty(t, req.SessionKey,
							"delegated centralized request carries no session token")
						return mockH, nil
					})
			},
			expectStatus:  http.StatusOK,
			expectHandled: true,
			expectAbsent: []string{
				"Authorization",
				"X-Api-Key",
				agplaibridge.HeaderCoderToken,
			},
		},
		{
			name: "valid BYOK preserves user credentials",
			reqHeaders: map[string]string{
				// Marks BYOK; this header must be stripped before
				// forwarding upstream. Its value is what gets
				// surfaced downstream as the SessionKey because
				// ExtractAuthToken prefers HeaderCoderToken.
				agplaibridge.HeaderCoderToken: "coder-token-byok",
				// The user's own LLM credential; must be preserved.
				"Authorization": "Bearer sk-ant-oat01-user-token",
			},
			applyMocks: func(t *testing.T, client *mock.MockDRPCClient, pool *mock.MockPooler, mockH *mockHandler) {