feat: record user agent in AI Bridge interceptions

Author: pawbana

enterprise/aibridged/aibridged_integration_test.go

@@ -3,6 +3,7 @@ package aibridged_test
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -226,10 +227,13 @@ func TestIntegration(t *testing.T) {
     }
   ]
 }`))
+	userAgent := "userAgent123"
 	require.NoError(t, err, "make request to test server")
 	req.Header.Add("Authorization", "Bearer "+apiKey.Key)
 	req.Header.Add("Accept", "application/json")
+	req.Header.Add("User-Agent", userAgent)
 
+	require.Equal(t, userAgent, req.UserAgent())
 	// When: aibridged handles the request.
 	rec := httptest.NewRecorder()
 	srv.ServeHTTP(rec, req)
@@ -252,6 +256,13 @@ func TestIntegration(t *testing.T) {
 	require.False(t, intc0.EndedAt.Time.Before(intc0.StartedAt), "EndedAt should not be before StartedAt")
 	require.Less(t, intc0.EndedAt.Time.Sub(intc0.StartedAt), 5*time.Second)
 
+	require.True(t, intc0.Metadata.Valid)
+	meta := map[string]any{}
+	err = json.Unmarshal(intc0.Metadata.RawMessage, &meta)
+	require.NoError(t, err)
+	require.Contains(t, meta, "user-agent")
+	require.Equal(t, userAgent, meta["user-agent"])
+
 	prompts, err := db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptions[0].ID)
 	require.NoError(t, err)
 	require.Len(t, prompts, 1)

enterprise/aibridged/http.go

@@ -78,9 +78,10 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 	}
 
 	handler, err := s.GetRequestHandler(ctx, Request{
-		SessionKey:  key,
 		APIKeyID:    resp.ApiKeyId,
 		InitiatorID: id,
+		SessionKey:  key,
+		UserAgent:   r.UserAgent(),
 	})
 	if err != nil {
 		logger.Warn(ctx, "failed to acquire request handler", slog.Error(err))

enterprise/aibridged/pool.go

@@ -146,7 +146,10 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
 			return nil, xerrors.Errorf("acquire client: %w", err)
 		}
 
-		return &recorderTranslation{apiKeyID: req.APIKeyID, client: client}, nil
+		return &recorderTranslation{
+			apiKeyID: req.APIKeyID,
+			client:   client,
+		}, nil
 	})
 
 	// Slow path.

enterprise/aibridged/pool_test.go

@@ -57,9 +57,10 @@ func TestPool(t *testing.T) {
 
 	// ...and it will return it when acquired again.
 	instB, err := pool.Acquire(t.Context(), aibridged.Request{
-		SessionKey:  "key",
+		SessionKey:  "different key",
 		InitiatorID: id,
 		APIKeyID:    apiKeyID1.String(),
+		UserAgent:   "some user-agent",
 	}, clientFn, newMockMCPFactory(mcpProxy))
 	require.NoError(t, err, "acquire pool instance")
 	require.Same(t, inst, instB)

enterprise/aibridged/proto/aibridged.pb.go

@@ -43,6 +43,8 @@ message RecordInterceptionRequest {
   map<string, google.protobuf.Any> metadata = 5;
   google.protobuf.Timestamp started_at = 6;
   string api_key_id = 7;
+  string client = 8;
+  string user_agent = 9;
 }
 
 message RecordInterceptionResponse {}

enterprise/aibridged/proto/aibridged.proto

@@ -3,7 +3,8 @@ package aibridged
 import "github.com/google/uuid"
 
 type Request struct {
-	SessionKey  string
 	APIKeyID    string
 	InitiatorID uuid.UUID
+	SessionKey  string
+	UserAgent   string
 }

enterprise/aibridged/request.go

@@ -30,6 +30,7 @@ func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibri
 		InitiatorId: req.InitiatorID,
 		Provider:    req.Provider,
 		Model:       req.Model,
+		UserAgent:   req.UserAgent,
 		Metadata:    marshalForProto(req.Metadata),
 		StartedAt:   timestamppb.New(req.StartedAt),
 	})

enterprise/aibridged/translator.go

@@ -0,0 +1,91 @@
+package aibridged //nolint:testpackage
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+	"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"
+
+	"github.com/coder/aibridge"
+	abpb "github.com/coder/coder/v2/enterprise/aibridged/proto"
+	"github.com/coder/coder/v2/testutil"
+)
+
+const (
+	MetaKeyUserAgent = "user-agent"
+)
+
+type mockClient struct {
+	abpb.DRPCRecorderClient
+	got *abpb.RecordInterceptionRequest
+}
+
+func (mc *mockClient) RecordInterception(ctx context.Context, in *abpb.RecordInterceptionRequest) (*abpb.RecordInterceptionResponse, error) {
+	mc.got = in
+	return &abpb.RecordInterceptionResponse{}, nil
+}
+
+func mustAnypbNew(t *testing.T, src proto.Message) *anypb.Any {
+	ret, err := anypb.New(src)
+	require.NoError(t, err)
+	return ret
+}
+
+func TestRecordInterception(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		apiKeyID  string
+		userAgent string
+		in        aibridge.InterceptionRecord
+		expect    *abpb.RecordInterceptionRequest
+	}{
+		{
+			name:     "ok",
+			apiKeyID: "key",
+			in: aibridge.InterceptionRecord{
+				ID:          uuid.UUID{1}.String(),
+				InitiatorID: uuid.UUID{2}.String(),
+				Provider:    "prov",
+				Model:       "model",
+				UserAgent:   "user-agent",
+				Metadata:    map[string]any{"some": "data"},
+				StartedAt:   time.UnixMicro(123),
+			},
+			expect: &abpb.RecordInterceptionRequest{
+				Id:          uuid.UUID{1}.String(),
+				ApiKeyId:    "key",
+				InitiatorId: uuid.UUID{2}.String(),
+				Provider:    "prov",
+				Model:       "model",
+				UserAgent:   "user-agent",
+				Metadata: map[string]*anypb.Any{
+					"some": mustAnypbNew(t, structpb.NewStringValue("data")),
+				},
+				StartedAt: timestamppb.New(time.UnixMicro(123)),
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			ctx := testutil.Context(t, testutil.WaitShort)
+
+			mc := &mockClient{}
+			rt := &recorderTranslation{
+				apiKeyID: tc.apiKeyID,
+				client:   mc,
+			}
+			rt.RecordInterception(ctx, &tc.in)
+			require.Equal(t, tc.expect, mc.got)
+		})
+	}
+}

enterprise/aibridged/translator_test.go

@@ -27,6 +27,10 @@ import (
 	"github.com/coder/coder/v2/enterprise/aibridged/proto"
 )
 
+const (
+	MetaKeyUserAgent = "user-agent"
+)
+
 var (
 	ErrExpiredOrInvalidOAuthToken = xerrors.New("expired or invalid OAuth2 token")
 	ErrNoMCPConfigFound           = xerrors.New("no MCP config found")
@@ -144,6 +148,15 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce
 		)
 	}
 
+	if in.UserAgent != "" {
+		md, err := anypb.New(structpb.NewStringValue(in.UserAgent))
+		if err != nil {
+			s.logger.Warn(ctx, "failed to convert user agent to proto", slog.Error(err))
+			return nil, xerrors.Errorf("invalid user agent")
+		}
+		metadata[MetaKeyUserAgent] = md
+	}
+
 	out, err := json.Marshal(metadata)
 	if err != nil {
 		s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))