Skip to content

Commit 5fb52dc

Browse files
jakehwlldannykopping
authored and
david-fraley
committed
fix: remove inflight interceptions from aibridge returned values (#20852)
Addresses [`aibridge#54`](coder/aibridge#54) When querying against the values in the database for `/api/experimental/aibridge/interceptions` we found strange behaviour wherein there was interceptions that lacked prompting and other various fields we want. Generally this was as a result of the data not actually existing for these values (as they were inflight). The simple solution to this was to hide them if they didn't exist. This PR addresses that. --------- Co-authored-by: Danny Kopping <[email protected]>
1 parent b19e6ce commit 5fb52dc

4 files changed

Lines changed: 62 additions & 12 deletions

File tree

‎coderd/database/queries.sql.go‎

Lines changed: 6 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/aibridge.sql‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,8 +89,10 @@ SELECT
8989
FROM
9090
aibridge_interceptions
9191
WHERE
92+
-- Remove inflight interceptions (ones which lack an ended_at value).
93+
aibridge_interceptions.ended_at IS NOT NULL
9294
-- Filter by time frame
93-
CASE
95+
AND CASE
9496
WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= @started_after::timestamptz
9597
ELSE true
9698
END
@@ -126,8 +128,10 @@ FROM
126128
JOIN
127129
visible_users ON visible_users.id = aibridge_interceptions.initiator_id
128130
WHERE
131+
-- Remove inflight interceptions (ones which lack an ended_at value).
132+
aibridge_interceptions.ended_at IS NOT NULL
129133
-- Filter by time frame
130-
CASE
134+
AND CASE
131135
WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= @started_after::timestamptz
132136
ELSE true
133137
END

‎enterprise/cli/aibridge_test.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,10 +43,11 @@ func TestAIBridgeListInterceptions(t *testing.T) {
4343
InitiatorID: member.ID,
4444
StartedAt: now.Add(-time.Hour),
4545
}, &now)
46+
interception2EndedAt := now.Add(time.Minute)
4647
interception2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
4748
InitiatorID: member.ID,
4849
StartedAt: now,
49-
}, nil)
50+
}, &interception2EndedAt)
5051
// Should not be returned because the user can't see it.
5152
_ = dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
5253
InitiatorID: owner.UserID,
@@ -91,12 +92,13 @@ func TestAIBridgeListInterceptions(t *testing.T) {
9192
now := dbtime.Now()
9293

9394
// This interception should be returned since it matches all filters.
95+
goodInterceptionEndedAt := now.Add(time.Minute)
9496
goodInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
9597
InitiatorID: member.ID,
9698
Provider: "real-provider",
9799
Model: "real-model",
98100
StartedAt: now,
99-
}, nil)
101+
}, &goodInterceptionEndedAt)
100102

101103
// These interceptions should not be returned since they don't match the
102104
// filters.
@@ -173,10 +175,11 @@ func TestAIBridgeListInterceptions(t *testing.T) {
173175
memberClient, member := coderdtest.CreateAnotherUser(t, client, owner.OrganizationID)
174176

175177
now := dbtime.Now()
178+
firstInterceptionEndedAt := now.Add(time.Minute)
176179
firstInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
177180
InitiatorID: member.ID,
178181
StartedAt: now,
179-
}, nil)
182+
}, &firstInterceptionEndedAt)
180183
returnedInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
181184
InitiatorID: member.ID,
182185
StartedAt: now.Add(-time.Hour),

‎enterprise/coderd/aibridge_test.go‎

Lines changed: 44 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,12 @@ func TestAIBridgeListInterceptions(t *testing.T) {
103103
// Insert a bunch of test data.
104104
now := dbtime.Now()
105105
i1ApiKey := sql.NullString{String: "some-api-key", Valid: true}
106+
i1EndedAt := now.Add(-time.Hour + time.Minute)
106107
i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
107108
APIKeyID: i1ApiKey,
108109
InitiatorID: user1.ID,
109110
StartedAt: now.Add(-time.Hour),
110-
}, nil)
111+
}, &i1EndedAt)
111112
i1tok1 := dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{
112113
InterceptionID: i1.ID,
113114
CreatedAt: now,
@@ -175,9 +176,11 @@ func TestAIBridgeListInterceptions(t *testing.T) {
175176
// Time comparison
176177
require.Len(t, res.Results, 2)
177178
require.Equal(t, res.Results[0].ID, i2SDK.ID)
178-
require.NotNil(t, now, res.Results[0].EndedAt)
179+
require.NotNil(t, res.Results[0].EndedAt)
179180
require.WithinDuration(t, now, *res.Results[0].EndedAt, 5*time.Second)
180181
res.Results[0].EndedAt = i2SDK.EndedAt
182+
require.NotNil(t, res.Results[1].EndedAt)
183+
res.Results[1].EndedAt = i1SDK.EndedAt
181184

182185
require.Equal(t, []codersdk.AIBridgeInterception{i2SDK, i1SDK}, res.Results)
183186
})
@@ -217,11 +220,12 @@ func TestAIBridgeListInterceptions(t *testing.T) {
217220
randomOffset, err := cryptorand.Intn(10000)
218221
require.NoError(t, err)
219222
randomOffsetDur := time.Duration(randomOffset) * time.Second
223+
endedAt := now.Add(randomOffsetDur + time.Minute)
220224
interception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
221225
ID: uuid.UUID{byte(i + 10)},
222226
InitiatorID: firstUser.UserID,
223227
StartedAt: now.Add(randomOffsetDur),
224-
}, nil)
228+
}, &endedAt)
225229
allInterceptionIDs = append(allInterceptionIDs, interception.ID)
226230
}
227231

@@ -297,6 +301,39 @@ func TestAIBridgeListInterceptions(t *testing.T) {
297301
}
298302
})
299303

304+
t.Run("InflightInterceptions", func(t *testing.T) {
305+
t.Parallel()
306+
dv := coderdtest.DeploymentValues(t)
307+
client, db, firstUser := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{
308+
Options: &coderdtest.Options{
309+
DeploymentValues: dv,
310+
},
311+
LicenseOptions: &coderdenttest.LicenseOptions{
312+
Features: license.Features{
313+
codersdk.FeatureAIBridge: 1,
314+
},
315+
},
316+
})
317+
ctx := testutil.Context(t, testutil.WaitLong)
318+
319+
now := dbtime.Now()
320+
i1EndedAt := now.Add(time.Minute)
321+
i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
322+
InitiatorID: firstUser.UserID,
323+
StartedAt: now,
324+
}, &i1EndedAt)
325+
dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
326+
InitiatorID: firstUser.UserID,
327+
StartedAt: now.Add(-time.Hour),
328+
}, nil)
329+
330+
res, err := client.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{})
331+
require.NoError(t, err)
332+
require.EqualValues(t, 1, res.Count)
333+
require.Len(t, res.Results, 1)
334+
require.Equal(t, i1.ID, res.Results[0].ID)
335+
})
336+
300337
t.Run("Authorized", func(t *testing.T) {
301338
t.Parallel()
302339
dv := coderdtest.DeploymentValues(t)
@@ -315,10 +352,11 @@ func TestAIBridgeListInterceptions(t *testing.T) {
315352
secondUserClient, secondUser := coderdtest.CreateAnotherUser(t, adminClient, firstUser.OrganizationID)
316353

317354
now := dbtime.Now()
355+
i1EndedAt := now.Add(time.Minute)
318356
i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
319357
InitiatorID: firstUser.UserID,
320358
StartedAt: now,
321-
}, nil)
359+
}, &i1EndedAt)
322360
i2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
323361
InitiatorID: secondUser.ID,
324362
StartedAt: now.Add(-time.Hour),
@@ -374,13 +412,14 @@ func TestAIBridgeListInterceptions(t *testing.T) {
374412

375413
// Insert a bunch of test data with varying filterable fields.
376414
now := dbtime.Now()
415+
i1EndedAt := now.Add(time.Minute)
377416
i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
378417
ID: uuid.MustParse("00000000-0000-0000-0000-000000000001"),
379418
InitiatorID: user1.ID,
380419
Provider: "one",
381420
Model: "one",
382421
StartedAt: now,
383-
}, nil)
422+
}, &i1EndedAt)
384423
i2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
385424
ID: uuid.MustParse("00000000-0000-0000-0000-000000000002"),
386425
InitiatorID: user1.ID,

0 commit comments

Comments
 (0)