Function TestServeHTTP_ActorHeaders

coderd/aibridged/aibridged_test.go:609–702

TestServeHTTP_ActorHeaders validates that actor headers are correctly forwarded to upstream AI providers when SendActorHeaders is enabled in the provider configuration. These headers allow upstream providers to identify the user making the request for tracking and auditing purposes.

(t *testing.T)

// These headers allow upstream providers to identify the user making the request for
// tracking and auditing purposes.
func TestServeHTTP_ActorHeaders(t *testing.T) {
	t.Parallel()

	testUsername := "testuser"
	testUserID := uuid.New()

	cases := []struct {
		path string
	}{
		// Not a complete set of paths; we're not testing the specific APIs - just the provider configs.
		{
			path: "/openai/v1/chat/completions",
		},
		{
			path: "/anthropic/v1/messages",
		},
	}

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

			// Setup mock upstream AI server that captures headers.
			var receivedHeaders http.Header
			upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				receivedHeaders = r.Header.Clone()
				w.WriteHeader(http.StatusTeapot)
				_, _ = w.Write([]byte(`i am a teapot`))
			}))
			t.Cleanup(upstreamSrv.Close)

			// Setup with SendActorHeaders enabled.
			logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
			ctrl := gomock.NewController(t)
			client := mock.NewMockDRPCClient(ctrl)

			// Create providers with SendActorHeaders=true.
			providers := []aibridge.Provider{
				aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{
					BaseURL:          upstreamSrv.URL,
					SendActorHeaders: true,
				}),
				aibridge.NewAnthropicProvider(aibridge.AnthropicConfig{
					BaseURL:          upstreamSrv.URL,
					SendActorHeaders: true,
				}, nil),
			}

			pool, err := aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions, providers, logger, nil, testTracer)
			require.NoError(t, err)
			conn := &mockDRPCConn{}
			client.EXPECT().DRPCConn().AnyTimes().Return(conn)

			// Return authorization response with user ID and username.
			client.EXPECT().IsAuthorized(gomock.Any(), gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{
				OwnerId:  testUserID.String(),
				Username: testUsername,
			}, nil)