opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: initialize aibridged & mount API handler (#19798)

Browse Source 
Addresses https://github.com/coder/internal/issues/987
This commit is contained in:
Danny Kopping
2025-09-25 16:37:28 +02:00
committed by GitHub
parent 6971f612be
commit 0a79817050
14 changed files with 798 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbauthz/dbauthz.go
+37 -6
View File
@@ -175,15 +175,15 @@ func (q *querier) authorizePrebuiltWorkspace(ctx context.Context, action policy.
return xerrors.Errorf("authorize context: %w", workspaceErr)
}
// authorizeAIBridgeInterceptionUpdate validates that the context's actor matches the initiator of the AIBridgeInterception.
// authorizeAIBridgeInterceptionAction validates that the context's actor matches the initiator of the AIBridgeInterception.
// This is used by all of the sub-resources which fall under the [ResourceAibridgeInterception] umbrella.
func (q *querier) authorizeAIBridgeInterceptionUpdate(ctx context.Context, interceptionID uuid.UUID) error {
func (q *querier) authorizeAIBridgeInterceptionAction(ctx context.Context, action policy.Action, interceptionID uuid.UUID) error {
inter, err := q.db.GetAIBridgeInterceptionByID(ctx, interceptionID)
if err != nil {
return xerrors.Errorf("fetch aibridge interception %q: %w", interceptionID, err)
}
err = q.authorizeContext(ctx, policy.ActionUpdate, inter.RBACObject())
err = q.authorizeContext(ctx, action, inter.RBACObject())
if err != nil {
return err
}
@@ -1928,6 +1928,37 @@ func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID)
return fetch(q.log, q.auth, q.db.GetAIBridgeInterceptionByID)(ctx, id)
}
func (q *querier) GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBridgeInterception, error) {
fetch := func(ctx context.Context, _ any) ([]database.AIBridgeInterception, error) {
return q.db.GetAIBridgeInterceptions(ctx)
}
return fetchWithPostFilter(q.auth, policy.ActionRead, fetch)(ctx, nil)
}
func (q *querier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeTokenUsage, error) {
// All aibridge_token_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionRead, interceptionID); err != nil {
return nil, err
}
return q.db.GetAIBridgeTokenUsagesByInterceptionID(ctx, interceptionID)
}
func (q *querier) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeToolUsage, error) {
// All aibridge_token_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionRead, interceptionID); err != nil {
return nil, err
}
return q.db.GetAIBridgeToolUsagesByInterceptionID(ctx, interceptionID)
}
func (q *querier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeUserPrompt, error) {
// All aibridge_token_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionRead, interceptionID); err != nil {
return nil, err
}
return q.db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptionID)
}
func (q *querier) GetAPIKeyByID(ctx context.Context, id string) (database.APIKey, error) {
return fetch(q.log, q.auth, q.db.GetAPIKeyByID)(ctx, id)
}
@@ -3813,7 +3844,7 @@ func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.I
func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error {
// All aibridge_token_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil {
return err
}
return q.db.InsertAIBridgeTokenUsage(ctx, arg)
@@ -3821,7 +3852,7 @@ func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.Ins
func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error {
// All aibridge_tool_usages records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil {
return err
}
return q.db.InsertAIBridgeToolUsage(ctx, arg)
@@ -3829,7 +3860,7 @@ func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.Inse
func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error {
// All aibridge_user_prompts records belong to the initiator of their associated interception.
if err := q.authorizeAIBridgeInterceptionUpdate(ctx, arg.InterceptionID); err != nil {
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil {
return err
}
return q.db.InsertAIBridgeUserPrompt(ctx, arg)
coderd/database/dbauthz/dbauthz_test.go
+61 -24
View File
@@ -4335,10 +4335,10 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) {
func (s *MethodTestSuite) TestAIBridge() {
s.Run("GetAIBridgeInterceptionByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
sessID := uuid.UUID{2}
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes()
check.Args(sessID).Asserts(sess, policy.ActionRead).Returns(sess)
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes()
check.Args(intID).Asserts(intc, policy.ActionRead).Returns(intc)
}))
s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
@@ -4348,39 +4348,76 @@ func (s *MethodTestSuite) TestAIBridge() {
user.IsSystem = false
user.Deleted = false
sessID := uuid.UUID{2}
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID, InitiatorID: initID})
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID, InitiatorID: initID})
params := database.InsertAIBridgeInterceptionParams{ID: sess.ID, InitiatorID: sess.InitiatorID, Provider: sess.Provider, Model: sess.Model}
params := database.InsertAIBridgeInterceptionParams{ID: intc.ID, InitiatorID: intc.InitiatorID, Provider: intc.Provider, Model: intc.Model}
db.EXPECT().GetUserByID(gomock.Any(), initID).Return(user, nil).AnyTimes() // Validation.
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), params).Return(sess, nil).AnyTimes()
check.Args(params).Asserts(sess, policy.ActionCreate)
db.EXPECT().InsertAIBridgeInterception(gomock.Any(), params).Return(intc, nil).AnyTimes()
check.Args(params).Asserts(intc, policy.ActionCreate)
}))
s.Run("InsertAIBridgeTokenUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
sessID := uuid.UUID{2}
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
params := database.InsertAIBridgeTokenUsageParams{InterceptionID: sess.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
params := database.InsertAIBridgeTokenUsageParams{InterceptionID: intc.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), params).Return(nil).AnyTimes()
check.Args(params).Asserts(sess, policy.ActionUpdate)
check.Args(params).Asserts(intc, policy.ActionUpdate)
}))
s.Run("InsertAIBridgeToolUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
sessID := uuid.UUID{2}
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
params := database.InsertAIBridgeToolUsageParams{InterceptionID: sess.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
params := database.InsertAIBridgeToolUsageParams{InterceptionID: intc.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), params).Return(nil).AnyTimes()
check.Args(params).Asserts(sess, policy.ActionUpdate)
check.Args(params).Asserts(intc, policy.ActionUpdate)
}))
s.Run("InsertAIBridgeUserPrompt", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
sessID := uuid.UUID{2}
sess := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: sessID})
params := database.InsertAIBridgeUserPromptParams{InterceptionID: sess.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), sessID).Return(sess, nil).AnyTimes() // Validation.
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
params := database.InsertAIBridgeUserPromptParams{InterceptionID: intc.ID}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), params).Return(nil).AnyTimes()
check.Args(params).Asserts(sess, policy.ActionUpdate)
check.Args(params).Asserts(intc, policy.ActionUpdate)
}))
s.Run("GetAIBridgeInterceptions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
a := testutil.Fake(s.T(), faker, database.AIBridgeInterception{})
b := testutil.Fake(s.T(), faker, database.AIBridgeInterception{})
db.EXPECT().GetAIBridgeInterceptions(gomock.Any()).Return([]database.AIBridgeInterception{a, b}, nil).AnyTimes()
check.Args().Asserts(a, policy.ActionRead, b, policy.ActionRead).Returns([]database.AIBridgeInterception{a, b})
}))
s.Run("GetAIBridgeTokenUsagesByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
tok := testutil.Fake(s.T(), faker, database.AIBridgeTokenUsage{InterceptionID: intID})
toks := []database.AIBridgeTokenUsage{tok}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().GetAIBridgeTokenUsagesByInterceptionID(gomock.Any(), intID).Return(toks, nil).AnyTimes()
check.Args(intID).Asserts(intc, policy.ActionRead).Returns(toks)
}))
s.Run("GetAIBridgeToolUsagesByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
tool := testutil.Fake(s.T(), faker, database.AIBridgeToolUsage{InterceptionID: intID})
tools := []database.AIBridgeToolUsage{tool}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().GetAIBridgeToolUsagesByInterceptionID(gomock.Any(), intID).Return(tools, nil).AnyTimes()
check.Args(intID).Asserts(intc, policy.ActionRead).Returns(tools)
}))
s.Run("GetAIBridgeUserPromptsByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intID := uuid.UUID{2}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID})
pr := testutil.Fake(s.T(), faker, database.AIBridgeUserPrompt{InterceptionID: intID})
prs := []database.AIBridgeUserPrompt{pr}
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().GetAIBridgeUserPromptsByInterceptionID(gomock.Any(), intID).Return(prs, nil).AnyTimes()
check.Args(intID).Asserts(intc, policy.ActionRead).Returns(prs)
}))
}