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

feat: add db query for setting interception ended_at field#20437

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

Merged
pawbana merged 7 commits intomainfrompb/aibridge-intc-ended-at-db-method
Oct 27, 2025
Merged
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
4 changes: 4 additions & 0 deletionscoderd/apidoc/docs.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

4 changes: 4 additions & 0 deletionscoderd/apidoc/swagger.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

6 changes: 5 additions & 1 deletioncoderd/database/db2sdk/db2sdk.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -962,7 +962,7 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
// created_at ASC
return sdkToolUsages[i].CreatedAt.Before(sdkToolUsages[j].CreatedAt)
})
return codersdk.AIBridgeInterception{
intc := codersdk.AIBridgeInterception{
ID: interception.ID,
Initiator: MinimalUserFromVisibleUser(initiator),
Provider: interception.Provider,
Expand All@@ -973,6 +973,10 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
UserPrompts: sdkUserPrompts,
ToolUsages: sdkToolUsages,
}
if interception.EndedAt.Valid {
intc.EndedAt = &interception.EndedAt.Time
}
return intc
}

func AIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage {
Expand Down
7 changes: 7 additions & 0 deletionscoderd/database/dbauthz/dbauthz.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -4711,6 +4711,13 @@ func (q *querier) UnfavoriteWorkspace(ctx context.Context, id uuid.UUID) error {
return update(q.log, q.auth, fetch, q.db.UnfavoriteWorkspace)(ctx, id)
}

func (q *querier) UpdateAIBridgeInterceptionEnded(ctx context.Context, params database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception, error) {
if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, params.ID); err != nil {
return database.AIBridgeInterception{}, err
}
return q.db.UpdateAIBridgeInterceptionEnded(ctx, params)
}

func (q *querier) UpdateAPIKeyByID(ctx context.Context, arg database.UpdateAPIKeyByIDParams) error {
fetch := func(ctx context.Context, arg database.UpdateAPIKeyByIDParams) (database.APIKey, error) {
return q.db.GetAPIKeyByID(ctx, arg.ID)
Expand Down
9 changes: 9 additions & 0 deletionscoderd/database/dbauthz/dbauthz_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -4617,4 +4617,13 @@ func (s *MethodTestSuite) TestAIBridge() {
db.EXPECT().ListAIBridgeToolUsagesByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeToolUsage{}, nil).AnyTimes()
check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns([]database.AIBridgeToolUsage{})
}))

s.Run("UpdateAIBridgeInterceptionEnded", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intcID := uuid.UUID{1}
params := database.UpdateAIBridgeInterceptionEndedParams{ID: intcID}
intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intcID})
db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intcID).Return(intc, nil).AnyTimes() // Validation.
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), params).Return(intc, nil).AnyTimes()
check.Args(params).Asserts(intc, policy.ActionUpdate).Returns(intc)
}))
}
9 changes: 8 additions & 1 deletioncoderd/database/dbgen/dbgen.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -1495,7 +1495,7 @@ func ClaimPrebuild(
return claimedWorkspace
}

func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams) database.AIBridgeInterception {
func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams, endedAt *time.Time) database.AIBridgeInterception {
interception, err := db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{
ID: takeFirst(seed.ID, uuid.New()),
InitiatorID: takeFirst(seed.InitiatorID, uuid.New()),
Expand All@@ -1504,6 +1504,13 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
StartedAt: takeFirst(seed.StartedAt, dbtime.Now()),
})
if endedAt != nil {
interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
ID: interception.ID,
EndedAt: *endedAt,
})
require.NoError(t, err, "insert aibridge interception")
}
require.NoError(t, err, "insert aibridge interception")
return interception
}
Expand Down
7 changes: 7 additions & 0 deletionscoderd/database/dbmetrics/querymetrics.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

15 changes: 15 additions & 0 deletionscoderd/database/dbmock/dbmock.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

1 change: 1 addition & 0 deletionscoderd/database/querier.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

65 changes: 65 additions & 0 deletionscoderd/database/querier_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -7724,3 +7724,68 @@ func TestUpdateTaskWorkspaceID(t *testing.T) {
})
}
}

func TestUpdateAIBridgeInterceptionEnded(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)

t.Run("NonExistingInterception", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)

got, err := db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
ID: uuid.New(),
EndedAt: time.Now(),
})
require.ErrorContains(t, err, "no rows in result set")
require.EqualValues(t, database.AIBridgeInterception{}, got)
})

t.Run("OK", func(t *testing.T) {
t.Parallel()
ctx := testutil.Context(t, testutil.WaitLong)

user := dbgen.User(t, db, database.User{})
interceptions := []database.AIBridgeInterception{}

for _, uid := range []uuid.UUID{{1}, {2}, {3}} {
insertParams := database.InsertAIBridgeInterceptionParams{
ID: uid,
InitiatorID: user.ID,
Metadata: json.RawMessage("{}"),
}

intc, err := db.InsertAIBridgeInterception(ctx, insertParams)
require.NoError(t, err)
require.Equal(t, uid, intc.ID)
require.False(t, intc.EndedAt.Valid)
interceptions = append(interceptions, intc)
}

intc0 := interceptions[0]
endedAt := time.Now()
// Mark first interception as done
updated, err := db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
ID: intc0.ID,
EndedAt: endedAt,
})
require.NoError(t, err)
require.EqualValues(t, updated.ID, intc0.ID)
require.True(t, updated.EndedAt.Valid)
require.WithinDuration(t, endedAt, updated.EndedAt.Time, 5*time.Second)

// Updating first interception again should fail
updated, err = db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
ID: intc0.ID,
EndedAt: endedAt.Add(time.Hour),
})
require.ErrorIs(t, err, sql.ErrNoRows)

// Other interceptions should not have ended_at set
for _, intc := range interceptions[1:] {
got, err := db.GetAIBridgeInterceptionByID(ctx, intc.ID)
require.NoError(t, err)
require.False(t, got.EndedAt.Valid)
}
})
}
29 changes: 29 additions & 0 deletionscoderd/database/queries.sql.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

8 changes: 8 additions & 0 deletionscoderd/database/queries/aibridge.sql
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,6 +6,14 @@ INSERT INTO aibridge_interceptions (
)
RETURNING *;

-- name: UpdateAIBridgeInterceptionEnded :one
UPDATE aibridge_interceptions
SET ended_at = @ended_at::timestamptz
WHERE
id = @id::uuid
AND ended_at IS NULL
RETURNING *;

-- name: InsertAIBridgeTokenUsage :one
INSERT INTO aibridge_token_usages (
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
Expand Down
1 change: 1 addition & 0 deletionscodersdk/aibridge.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -18,6 +18,7 @@ type AIBridgeInterception struct {
Model string `json:"model"`
Metadata map[string]any `json:"metadata"`
StartedAt time.Time `json:"started_at" format:"date-time"`
EndedAt *time.Time `json:"ended_at" format:"date-time"`
TokenUsages []AIBridgeTokenUsage `json:"token_usages"`
UserPrompts []AIBridgeUserPrompt `json:"user_prompts"`
ToolUsages []AIBridgeToolUsage `json:"tool_usages"`
Expand Down
1 change: 1 addition & 0 deletionsdocs/reference/api/aibridge.md
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

3 changes: 3 additions & 0 deletionsdocs/reference/api/schemas.md
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

24 changes: 12 additions & 12 deletionsenterprise/cli/exp_aibridge_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -43,16 +43,16 @@ func TestAIBridgeListInterceptions(t *testing.T) {
interception1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: member.ID,
StartedAt: now.Add(-time.Hour),
})
}, &now)
interception2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: member.ID,
StartedAt: now,
})
}, nil)
// Should not be returned because the user can't see it.
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: owner.UserID,
StartedAt: now.Add(-2 * time.Hour),
})
}, nil)

args := []string{
"exp",
Expand DownExpand Up@@ -99,7 +99,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
Provider: "real-provider",
Model: "real-model",
StartedAt: now,
})
}, nil)

// These interceptions should not be returned since they don't match the
// filters.
Expand All@@ -108,33 +108,33 @@ func TestAIBridgeListInterceptions(t *testing.T) {
Provider: goodInterception.Provider,
Model: goodInterception.Model,
StartedAt: goodInterception.StartedAt,
})
}, nil)
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: goodInterception.InitiatorID,
Provider: "bad-provider",
Model: goodInterception.Model,
StartedAt: goodInterception.StartedAt,
})
}, nil)
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: goodInterception.InitiatorID,
Provider: goodInterception.Provider,
Model: "bad-model",
StartedAt: goodInterception.StartedAt,
})
}, nil)
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: goodInterception.InitiatorID,
Provider: goodInterception.Provider,
Model: goodInterception.Model,
// Violates the started after filter.
StartedAt: now.Add(-2 * time.Hour),
})
}, nil)
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: goodInterception.InitiatorID,
Provider: goodInterception.Provider,
Model: goodInterception.Model,
// Violates the started before filter.
StartedAt: now.Add(2 * time.Hour),
})
}, nil)

args := []string{
"exp",
Expand DownExpand Up@@ -181,15 +181,15 @@ func TestAIBridgeListInterceptions(t *testing.T) {
firstInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: member.ID,
StartedAt: now,
})
}, nil)
returnedInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: member.ID,
StartedAt: now.Add(-time.Hour),
})
}, &now)
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
InitiatorID: member.ID,
StartedAt: now.Add(-2 * time.Hour),
})
}, nil)

args := []string{
"exp",
Expand Down
Loading
Loading
[8]ページ先頭
©2009-2025 Movatter.jp