Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

chore: record user agent in interception metadata#20695

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Open
pawbana wants to merge1 commit intomain
base:main
Choose a base branch
Loading
from10-30-pb-aibridge-intc-user-agent-metadata
Open
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletionsenterprise/aibridged/aibridged_integration_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
Expand All@@ -20,7 +21,7 @@
"github.com/coder/coder/v2/coderd/externalauth"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/enterprise/aibridged"

Check failure on line 24 in enterprise/aibridged/aibridged_integration_test.go

View workflow job for this annotation

GitHub Actions/ lint

could not import github.com/coder/coder/v2/enterprise/aibridged (-: # github.com/coder/coder/v2/enterprise/aibridged [github.com/coder/coder/v2/enterprise/aibridged.test]
"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
"github.com/coder/coder/v2/testutil"
)
Expand DownExpand Up@@ -208,10 +209,13 @@
}
]
}`))
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)
Expand All@@ -234,6 +238,13 @@
require.True(t, intc0.StartedAt.Before(intc0.EndedAt.Time))
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)
Expand Down
3 changes: 2 additions & 1 deletionenterprise/aibridged/http.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -67,9 +67,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))
Expand Down
5 changes: 4 additions & 1 deletionenterprise/aibridged/pool.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -129,7 +129,10 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
returnnil,xerrors.Errorf("acquire client: %w",err)
}

return&recorderTranslation{apiKeyID:req.APIKeyID,client:client},nil
return&recorderTranslation{
apiKeyID:req.APIKeyID,
client:client,
},nil
})

// Slow path.
Expand Down
3 changes: 2 additions & 1 deletionenterprise/aibridged/pool_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -56,9 +56,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: "differentkey",
InitiatorID: id,
APIKeyID: apiKeyID1.String(),
UserAgent: "some user-agent",
}, clientFn, newMockMCPFactory(mcpProxy))
require.NoError(t, err, "acquire pool instance")
require.Same(t, inst, instB)
Expand Down
426 changes: 218 additions & 208 deletionsenterprise/aibridged/proto/aibridged.pb.go
View file
Open in desktop

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletionsenterprise/aibridged/proto/aibridged.proto
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -43,6 +43,7 @@ message RecordInterceptionRequest {
map<string, google.protobuf.Any> metadata = 5;
google.protobuf.Timestamp started_at = 6;
string api_key_id = 7;
string user_agent = 8;
}

message RecordInterceptionResponse {}
Expand Down
3 changes: 2 additions & 1 deletionenterprise/aibridged/request.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
}
1 change: 1 addition & 0 deletionsenterprise/aibridged/translator.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,6 +31,7 @@
InitiatorId: req.InitiatorID,
Provider: req.Provider,
Model: req.Model,
UserAgent: req.UserAgent,

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ lint

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)) (typecheck)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ lint

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ lint

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ lint

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)) (typecheck)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ test-go-pg-17

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ test-go-pg (ubuntu-latest)

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ test-go-pg (ubuntu-latest)

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ test-e2e

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ test-go-race-pg

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)

Check failure on line 34 in enterprise/aibridged/translator.go

View workflow job for this annotation

GitHub Actions/ check-build

req.UserAgent undefined (type *aibridge.InterceptionRecord has no field or method UserAgent)
Metadata: marshalForProto(req.Metadata),
StartedAt: timestamppb.New(req.StartedAt),
})
Expand Down
91 changes: 91 additions & 0 deletionsenterprise/aibridged/translator_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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",

Check failure on line 58 in enterprise/aibridged/translator_test.go

View workflow job for this annotation

GitHub Actions/ lint

unknown field UserAgent in struct literal of type aibridge.InterceptionRecord) (typecheck)

Check failure on line 58 in enterprise/aibridged/translator_test.go

View workflow job for this annotation

GitHub Actions/ lint

unknown field UserAgent in struct literal of type aibridge.InterceptionRecord (typecheck)
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)
})
}
}
16 changes: 16 additions & 0 deletionsenterprise/aibridgedserver/aibridgedserver.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -24,10 +24,14 @@
"github.com/coder/coder/v2/coderd/httpmw"
codermcp "github.com/coder/coder/v2/coderd/mcp"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/enterprise/aibridged"

Check failure on line 27 in enterprise/aibridgedserver/aibridgedserver.go

View workflow job for this annotation

GitHub Actions/ lint

could not import github.com/coder/coder/v2/enterprise/aibridged (-: # github.com/coder/coder/v2/enterprise/aibridged
"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")
Expand DownExpand Up@@ -118,6 +122,18 @@
return nil, xerrors.Errorf("empty API key ID")
}

if in.UserAgent != "" {
if in.Metadata == nil {
in.Metadata = map[string]*anypb.Any{}
}
md, err := anypb.New(structpb.NewStringValue(in.UserAgent))
if err != nil {
s.logger.Warn(ctx, "failed to convert user agent to proto: %v", err)
return nil, xerrors.Errorf("invalid user agent")
}
in.Metadata[MetaKeyUserAgent] = md
}

_, err = s.store.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{
ID: intcID,
APIKeyID: sql.NullString{String: in.ApiKeyId, Valid: true},
Expand Down
Loading
[8]ページ先頭
©2009-2025 Movatter.jp