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

Commit0a79817

Browse files
authored
feat: initializeaibridged & mount API handler (#19798)
Addressescoder/internal#987
1 parent6971f61 commit0a79817

File tree

14 files changed

+798
-39
lines changed

14 files changed

+798
-39
lines changed

‎coderd/coderd.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@ import (
4444
"tailscale.com/types/key"
4545
"tailscale.com/util/singleflight"
4646

47+
"github.com/coder/coder/v2/provisionerd/proto"
48+
4749
"cdr.dev/slog"
4850
"github.com/coder/quartz"
4951
"github.com/coder/serpent"
@@ -95,7 +97,6 @@ import (
9597
"github.com/coder/coder/v2/coderd/workspacestats"
9698
"github.com/coder/coder/v2/codersdk"
9799
"github.com/coder/coder/v2/codersdk/healthsdk"
98-
"github.com/coder/coder/v2/provisionerd/proto"
99100
"github.com/coder/coder/v2/provisionersdk"
100101
"github.com/coder/coder/v2/site"
101102
"github.com/coder/coder/v2/tailnet"
@@ -999,6 +1000,11 @@ func New(options *Options) *API {
9991000

10001001
// Experimental routes are not guaranteed to be stable and may change at any time.
10011002
r.Route("/api/experimental",func(r chi.Router) {
1003+
r.NotFound(func(rw http.ResponseWriter,_*http.Request) {httpapi.RouteNotFound(rw) })
1004+
1005+
// Only this group should be subject to apiKeyMiddleware; aibridged will mount its own
1006+
// router and handles key validation in a different fashion.
1007+
// See enterprise/x/aibridged/http.go.
10021008
r.Group(func(r chi.Router) {
10031009
r.Use(apiKeyMiddleware)
10041010
r.Route("/aitasks",func(r chi.Router) {

‎coderd/database/dbauthz/dbauthz.go‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -175,15 +175,15 @@ func (q *querier) authorizePrebuiltWorkspace(ctx context.Context, action policy.
175175
returnxerrors.Errorf("authorize context: %w",workspaceErr)
176176
}
177177

178-
//authorizeAIBridgeInterceptionUpdate validates that the context's actor matches the initiator of the AIBridgeInterception.
178+
//authorizeAIBridgeInterceptionAction validates that the context's actor matches the initiator of the AIBridgeInterception.
179179
// This is used by all of the sub-resources which fall under the [ResourceAibridgeInterception] umbrella.
180-
func (q*querier)authorizeAIBridgeInterceptionUpdate(ctx context.Context,interceptionID uuid.UUID)error {
180+
func (q*querier)authorizeAIBridgeInterceptionAction(ctx context.Context,action policy.Action,interceptionID uuid.UUID)error {
181181
inter,err:=q.db.GetAIBridgeInterceptionByID(ctx,interceptionID)
182182
iferr!=nil {
183183
returnxerrors.Errorf("fetch aibridge interception %q: %w",interceptionID,err)
184184
}
185185

186-
err=q.authorizeContext(ctx,policy.ActionUpdate,inter.RBACObject())
186+
err=q.authorizeContext(ctx,action,inter.RBACObject())
187187
iferr!=nil {
188188
returnerr
189189
}
@@ -1928,6 +1928,37 @@ func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID)
19281928
returnfetch(q.log,q.auth,q.db.GetAIBridgeInterceptionByID)(ctx,id)
19291929
}
19301930

1931+
func (q*querier)GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBridgeInterception,error) {
1932+
fetch:=func(ctx context.Context,_any) ([]database.AIBridgeInterception,error) {
1933+
returnq.db.GetAIBridgeInterceptions(ctx)
1934+
}
1935+
returnfetchWithPostFilter(q.auth,policy.ActionRead,fetch)(ctx,nil)
1936+
}
1937+
1938+
func (q*querier)GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context,interceptionID uuid.UUID) ([]database.AIBridgeTokenUsage,error) {
1939+
// All aibridge_token_usages records belong to the initiator of their associated interception.
1940+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionRead,interceptionID);err!=nil {
1941+
returnnil,err
1942+
}
1943+
returnq.db.GetAIBridgeTokenUsagesByInterceptionID(ctx,interceptionID)
1944+
}
1945+
1946+
func (q*querier)GetAIBridgeToolUsagesByInterceptionID(ctx context.Context,interceptionID uuid.UUID) ([]database.AIBridgeToolUsage,error) {
1947+
// All aibridge_token_usages records belong to the initiator of their associated interception.
1948+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionRead,interceptionID);err!=nil {
1949+
returnnil,err
1950+
}
1951+
returnq.db.GetAIBridgeToolUsagesByInterceptionID(ctx,interceptionID)
1952+
}
1953+
1954+
func (q*querier)GetAIBridgeUserPromptsByInterceptionID(ctx context.Context,interceptionID uuid.UUID) ([]database.AIBridgeUserPrompt,error) {
1955+
// All aibridge_token_usages records belong to the initiator of their associated interception.
1956+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionRead,interceptionID);err!=nil {
1957+
returnnil,err
1958+
}
1959+
returnq.db.GetAIBridgeUserPromptsByInterceptionID(ctx,interceptionID)
1960+
}
1961+
19311962
func (q*querier)GetAPIKeyByID(ctx context.Context,idstring) (database.APIKey,error) {
19321963
returnfetch(q.log,q.auth,q.db.GetAPIKeyByID)(ctx,id)
19331964
}
@@ -3813,23 +3844,23 @@ func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.I
38133844

38143845
func (q*querier)InsertAIBridgeTokenUsage(ctx context.Context,arg database.InsertAIBridgeTokenUsageParams)error {
38153846
// All aibridge_token_usages records belong to the initiator of their associated interception.
3816-
iferr:=q.authorizeAIBridgeInterceptionUpdate(ctx,arg.InterceptionID);err!=nil {
3847+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionUpdate,arg.InterceptionID);err!=nil {
38173848
returnerr
38183849
}
38193850
returnq.db.InsertAIBridgeTokenUsage(ctx,arg)
38203851
}
38213852

38223853
func (q*querier)InsertAIBridgeToolUsage(ctx context.Context,arg database.InsertAIBridgeToolUsageParams)error {
38233854
// All aibridge_tool_usages records belong to the initiator of their associated interception.
3824-
iferr:=q.authorizeAIBridgeInterceptionUpdate(ctx,arg.InterceptionID);err!=nil {
3855+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionUpdate,arg.InterceptionID);err!=nil {
38253856
returnerr
38263857
}
38273858
returnq.db.InsertAIBridgeToolUsage(ctx,arg)
38283859
}
38293860

38303861
func (q*querier)InsertAIBridgeUserPrompt(ctx context.Context,arg database.InsertAIBridgeUserPromptParams)error {
38313862
// All aibridge_user_prompts records belong to the initiator of their associated interception.
3832-
iferr:=q.authorizeAIBridgeInterceptionUpdate(ctx,arg.InterceptionID);err!=nil {
3863+
iferr:=q.authorizeAIBridgeInterceptionAction(ctx,policy.ActionUpdate,arg.InterceptionID);err!=nil {
38333864
returnerr
38343865
}
38353866
returnq.db.InsertAIBridgeUserPrompt(ctx,arg)

‎coderd/database/dbauthz/dbauthz_test.go‎

Lines changed: 61 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -4335,10 +4335,10 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) {
43354335

43364336
func (s*MethodTestSuite)TestAIBridge() {
43374337
s.Run("GetAIBridgeInterceptionByID",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4338-
sessID:= uuid.UUID{2}
4339-
sess:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:sessID})
4340-
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),sessID).Return(sess,nil).AnyTimes()
4341-
check.Args(sessID).Asserts(sess,policy.ActionRead).Returns(sess)
4338+
intID:= uuid.UUID{2}
4339+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4340+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()
4341+
check.Args(intID).Asserts(intc,policy.ActionRead).Returns(intc)
43424342
}))
43434343

43444344
s.Run("InsertAIBridgeInterception",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
@@ -4348,39 +4348,76 @@ func (s *MethodTestSuite) TestAIBridge() {
43484348
user.IsSystem=false
43494349
user.Deleted=false
43504350

4351-
sessID:= uuid.UUID{2}
4352-
sess:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:sessID,InitiatorID:initID})
4351+
intID:= uuid.UUID{2}
4352+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID,InitiatorID:initID})
43534353

4354-
params:= database.InsertAIBridgeInterceptionParams{ID:sess.ID,InitiatorID:sess.InitiatorID,Provider:sess.Provider,Model:sess.Model}
4354+
params:= database.InsertAIBridgeInterceptionParams{ID:intc.ID,InitiatorID:intc.InitiatorID,Provider:intc.Provider,Model:intc.Model}
43554355
db.EXPECT().GetUserByID(gomock.Any(),initID).Return(user,nil).AnyTimes()// Validation.
4356-
db.EXPECT().InsertAIBridgeInterception(gomock.Any(),params).Return(sess,nil).AnyTimes()
4357-
check.Args(params).Asserts(sess,policy.ActionCreate)
4356+
db.EXPECT().InsertAIBridgeInterception(gomock.Any(),params).Return(intc,nil).AnyTimes()
4357+
check.Args(params).Asserts(intc,policy.ActionCreate)
43584358
}))
43594359

43604360
s.Run("InsertAIBridgeTokenUsage",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4361-
sessID:= uuid.UUID{2}
4362-
sess:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:sessID})
4363-
params:= database.InsertAIBridgeTokenUsageParams{InterceptionID:sess.ID}
4364-
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),sessID).Return(sess,nil).AnyTimes()// Validation.
4361+
intID:= uuid.UUID{2}
4362+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4363+
params:= database.InsertAIBridgeTokenUsageParams{InterceptionID:intc.ID}
4364+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
43654365
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(),params).Return(nil).AnyTimes()
4366-
check.Args(params).Asserts(sess,policy.ActionUpdate)
4366+
check.Args(params).Asserts(intc,policy.ActionUpdate)
43674367
}))
43684368

43694369
s.Run("InsertAIBridgeToolUsage",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4370-
sessID:= uuid.UUID{2}
4371-
sess:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:sessID})
4372-
params:= database.InsertAIBridgeToolUsageParams{InterceptionID:sess.ID}
4373-
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),sessID).Return(sess,nil).AnyTimes()// Validation.
4370+
intID:= uuid.UUID{2}
4371+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4372+
params:= database.InsertAIBridgeToolUsageParams{InterceptionID:intc.ID}
4373+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
43744374
db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(),params).Return(nil).AnyTimes()
4375-
check.Args(params).Asserts(sess,policy.ActionUpdate)
4375+
check.Args(params).Asserts(intc,policy.ActionUpdate)
43764376
}))
43774377

43784378
s.Run("InsertAIBridgeUserPrompt",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4379-
sessID:= uuid.UUID{2}
4380-
sess:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:sessID})
4381-
params:= database.InsertAIBridgeUserPromptParams{InterceptionID:sess.ID}
4382-
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),sessID).Return(sess,nil).AnyTimes()// Validation.
4379+
intID:= uuid.UUID{2}
4380+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4381+
params:= database.InsertAIBridgeUserPromptParams{InterceptionID:intc.ID}
4382+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
43834383
db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(),params).Return(nil).AnyTimes()
4384-
check.Args(params).Asserts(sess,policy.ActionUpdate)
4384+
check.Args(params).Asserts(intc,policy.ActionUpdate)
4385+
}))
4386+
4387+
s.Run("GetAIBridgeInterceptions",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4388+
a:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{})
4389+
b:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{})
4390+
db.EXPECT().GetAIBridgeInterceptions(gomock.Any()).Return([]database.AIBridgeInterception{a,b},nil).AnyTimes()
4391+
check.Args().Asserts(a,policy.ActionRead,b,policy.ActionRead).Returns([]database.AIBridgeInterception{a,b})
4392+
}))
4393+
4394+
s.Run("GetAIBridgeTokenUsagesByInterceptionID",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4395+
intID:= uuid.UUID{2}
4396+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4397+
tok:=testutil.Fake(s.T(),faker, database.AIBridgeTokenUsage{InterceptionID:intID})
4398+
toks:= []database.AIBridgeTokenUsage{tok}
4399+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
4400+
db.EXPECT().GetAIBridgeTokenUsagesByInterceptionID(gomock.Any(),intID).Return(toks,nil).AnyTimes()
4401+
check.Args(intID).Asserts(intc,policy.ActionRead).Returns(toks)
4402+
}))
4403+
4404+
s.Run("GetAIBridgeToolUsagesByInterceptionID",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4405+
intID:= uuid.UUID{2}
4406+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4407+
tool:=testutil.Fake(s.T(),faker, database.AIBridgeToolUsage{InterceptionID:intID})
4408+
tools:= []database.AIBridgeToolUsage{tool}
4409+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
4410+
db.EXPECT().GetAIBridgeToolUsagesByInterceptionID(gomock.Any(),intID).Return(tools,nil).AnyTimes()
4411+
check.Args(intID).Asserts(intc,policy.ActionRead).Returns(tools)
4412+
}))
4413+
4414+
s.Run("GetAIBridgeUserPromptsByInterceptionID",s.Mocked(func(db*dbmock.MockStore,faker*gofakeit.Faker,check*expects) {
4415+
intID:= uuid.UUID{2}
4416+
intc:=testutil.Fake(s.T(),faker, database.AIBridgeInterception{ID:intID})
4417+
pr:=testutil.Fake(s.T(),faker, database.AIBridgeUserPrompt{InterceptionID:intID})
4418+
prs:= []database.AIBridgeUserPrompt{pr}
4419+
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(),intID).Return(intc,nil).AnyTimes()// Validation.
4420+
db.EXPECT().GetAIBridgeUserPromptsByInterceptionID(gomock.Any(),intID).Return(prs,nil).AnyTimes()
4421+
check.Args(intID).Asserts(intc,policy.ActionRead).Returns(prs)
43854422
}))
43864423
}

‎coderd/database/dbmetrics/querymetrics.go‎

Lines changed: 28 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/dbmock/dbmock.go‎

Lines changed: 60 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎coderd/database/querier.go‎

Lines changed: 4 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp