Compare commits

...

4 Commits

Author SHA1 Message Date
Paweł Banaszewski 42b467acf7 migration number fix + telemetry test fix 2026-02-17 09:43:08 +00:00
Paweł Banaszewski 98bc5f04fb migration fix 2026-02-17 09:43:08 +00:00
Paweł Banaszewski 02aeef5a91 review: dafault + unknown -> Unknown 2026-02-17 09:43:08 +00:00
Paweł Banaszewski 1cbc8929ab feat: add client columns to aibridge_interceptions
Also adds interception filtering by client.
2026-02-17 09:43:08 +00:00
16 changed files with 97 additions and 28 deletions
+1
View File
@@ -1590,6 +1590,7 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
Model: takeFirst(seed.Model, "model"),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
StartedAt: takeFirst(seed.StartedAt, dbtime.Now()),
Client: seed.Client,
})
if endedAt != nil {
interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
+4 -1
View File
@@ -1023,7 +1023,8 @@ CREATE TABLE aibridge_interceptions (
started_at timestamp with time zone NOT NULL,
metadata jsonb,
ended_at timestamp with time zone,
api_key_id text
api_key_id text,
client character varying(64) DEFAULT 'Unknown'::character varying
);
COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge';
@@ -3274,6 +3275,8 @@ CREATE INDEX idx_agent_stats_created_at ON workspace_agent_stats USING btree (cr
CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_id);
CREATE INDEX idx_aibridge_interceptions_client ON aibridge_interceptions USING btree (client);
CREATE INDEX idx_aibridge_interceptions_initiator_id ON aibridge_interceptions USING btree (initiator_id);
CREATE INDEX idx_aibridge_interceptions_model ON aibridge_interceptions USING btree (model);
@@ -0,0 +1,2 @@
ALTER TABLE aibridge_interceptions
DROP COLUMN client;
@@ -0,0 +1,5 @@
ALTER TABLE aibridge_interceptions
ADD COLUMN client VARCHAR(64)
DEFAULT 'Unknown';
CREATE INDEX idx_aibridge_interceptions_client ON aibridge_interceptions (client);
+3
View File
@@ -790,6 +790,7 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
arg.InitiatorID,
arg.Provider,
arg.Model,
arg.Client,
arg.AfterID,
arg.Offset,
arg.Limit,
@@ -810,6 +811,7 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.AIBridgeInterception.Client,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
@@ -847,6 +849,7 @@ func (q *sqlQuerier) CountAuthorizedAIBridgeInterceptions(ctx context.Context, a
arg.InitiatorID,
arg.Provider,
arg.Model,
arg.Client,
)
if err != nil {
return 0, err
+1
View File
@@ -3642,6 +3642,7 @@ type AIBridgeInterception struct {
Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"`
EndedAt sql.NullTime `db:"ended_at" json:"ended_at"`
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
Client sql.NullString `db:"client" json:"client"`
}
// Audit log of tokens used by intercepted requests in AI Bridge
+3
View File
@@ -8070,12 +8070,15 @@ func TestUpdateAIBridgeInterceptionEnded(t *testing.T) {
ID: uid,
InitiatorID: user.ID,
Metadata: json.RawMessage("{}"),
Client: sql.NullString{String: "client", Valid: true},
}
intc, err := db.InsertAIBridgeInterception(ctx, insertParams)
require.NoError(t, err)
require.Equal(t, uid, intc.ID)
require.False(t, intc.EndedAt.Valid)
require.True(t, intc.Client.Valid)
require.Equal(t, "client", intc.Client.String)
interceptions = append(interceptions, intc)
}
+35 -16
View File
@@ -123,8 +123,7 @@ WITH interceptions_in_range AS (
WHERE
provider = $1::text
AND model = $2::text
-- TODO: use the client value once we have it (see https://github.com/coder/aibridge/issues/31)
AND 'unknown' = $3::text
AND COALESCE(client, 'Unknown') = $3::text
AND ended_at IS NOT NULL -- incomplete interceptions are not included in summaries
AND ended_at >= $4::timestamptz
AND ended_at < $5::timestamptz
@@ -301,6 +300,11 @@ WHERE
WHEN $5::text != '' THEN aibridge_interceptions.model = $5::text
ELSE true
END
-- Filter client
AND CASE
WHEN $6::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = $6::text
ELSE true
END
-- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeInterceptions
-- @authorize_filter
`
@@ -311,6 +315,7 @@ type CountAIBridgeInterceptionsParams struct {
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
Client string `db:"client" json:"client"`
}
func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAIBridgeInterceptionsParams) (int64, error) {
@@ -320,6 +325,7 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI
arg.InitiatorID,
arg.Provider,
arg.Model,
arg.Client,
)
var count int64
err := row.Scan(&count)
@@ -372,7 +378,7 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
FROM
aibridge_interceptions
WHERE
@@ -391,13 +397,14 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
&i.Client,
)
return i, err
}
const getAIBridgeInterceptions = `-- name: GetAIBridgeInterceptions :many
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
FROM
aibridge_interceptions
`
@@ -420,6 +427,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
&i.Client,
); err != nil {
return nil, err
}
@@ -565,11 +573,11 @@ func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context,
const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, api_key_id, initiator_id, provider, model, metadata, started_at
id, api_key_id, initiator_id, provider, model, metadata, started_at, client
) VALUES (
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7, $8
)
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
`
type InsertAIBridgeInterceptionParams struct {
@@ -580,6 +588,7 @@ type InsertAIBridgeInterceptionParams struct {
Model string `db:"model" json:"model"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
Client sql.NullString `db:"client" json:"client"`
}
func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) {
@@ -591,6 +600,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
arg.Model,
arg.Metadata,
arg.StartedAt,
arg.Client,
)
var i AIBridgeInterception
err := row.Scan(
@@ -602,6 +612,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
&i.Client,
)
return i, err
}
@@ -740,7 +751,7 @@ func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIB
const listAIBridgeInterceptions = `-- name: ListAIBridgeInterceptions :many
SELECT
aibridge_interceptions.id, aibridge_interceptions.initiator_id, aibridge_interceptions.provider, aibridge_interceptions.model, aibridge_interceptions.started_at, aibridge_interceptions.metadata, aibridge_interceptions.ended_at, aibridge_interceptions.api_key_id,
aibridge_interceptions.id, aibridge_interceptions.initiator_id, aibridge_interceptions.provider, aibridge_interceptions.model, aibridge_interceptions.started_at, aibridge_interceptions.metadata, aibridge_interceptions.ended_at, aibridge_interceptions.api_key_id, aibridge_interceptions.client,
visible_users.id, visible_users.username, visible_users.name, visible_users.avatar_url
FROM
aibridge_interceptions
@@ -773,9 +784,14 @@ WHERE
WHEN $5::text != '' THEN aibridge_interceptions.model = $5::text
ELSE true
END
-- Filter client
AND CASE
WHEN $6::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = $6::text
ELSE true
END
-- Cursor pagination
AND CASE
WHEN $6::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN (
WHEN $7::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN (
-- The pagination cursor is the last ID of the previous page.
-- The query is ordered by the started_at field, so select all
-- rows before the cursor and before the after_id UUID.
@@ -783,8 +799,8 @@ WHERE
-- "after_id" terminology comes from our pagination parser in
-- coderd.
(aibridge_interceptions.started_at, aibridge_interceptions.id) < (
(SELECT started_at FROM aibridge_interceptions WHERE id = $6),
$6::uuid
(SELECT started_at FROM aibridge_interceptions WHERE id = $7),
$7::uuid
)
)
ELSE true
@@ -794,8 +810,8 @@ WHERE
ORDER BY
aibridge_interceptions.started_at DESC,
aibridge_interceptions.id DESC
LIMIT COALESCE(NULLIF($8::integer, 0), 100)
OFFSET $7
LIMIT COALESCE(NULLIF($9::integer, 0), 100)
OFFSET $8
`
type ListAIBridgeInterceptionsParams struct {
@@ -804,6 +820,7 @@ type ListAIBridgeInterceptionsParams struct {
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
Client string `db:"client" json:"client"`
AfterID uuid.UUID `db:"after_id" json:"after_id"`
Offset int32 `db:"offset_" json:"offset_"`
Limit int32 `db:"limit_" json:"limit_"`
@@ -821,6 +838,7 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr
arg.InitiatorID,
arg.Provider,
arg.Model,
arg.Client,
arg.AfterID,
arg.Offset,
arg.Limit,
@@ -841,6 +859,7 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.AIBridgeInterception.Client,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
@@ -864,8 +883,7 @@ SELECT
DISTINCT ON (provider, model, client)
provider,
model,
-- TODO: use the client value once we have it (see https://github.com/coder/aibridge/issues/31)
'unknown' AS client
COALESCE(client, 'Unknown') AS client
FROM
aibridge_interceptions
WHERE
@@ -1047,7 +1065,7 @@ UPDATE aibridge_interceptions
WHERE
id = $2::uuid
AND ended_at IS NULL
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
`
type UpdateAIBridgeInterceptionEndedParams struct {
@@ -1067,6 +1085,7 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
&i.Client,
)
return i, err
}
+14 -6
View File
@@ -1,8 +1,8 @@
-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, api_key_id, initiator_id, provider, model, metadata, started_at
id, api_key_id, initiator_id, provider, model, metadata, started_at, client
) VALUES (
@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at
@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at, @client
)
RETURNING *;
@@ -115,6 +115,11 @@ WHERE
WHEN @model::text != '' THEN aibridge_interceptions.model = @model::text
ELSE true
END
-- Filter client
AND CASE
WHEN @client::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = @client::text
ELSE true
END
-- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeInterceptions
-- @authorize_filter
;
@@ -154,6 +159,11 @@ WHERE
WHEN @model::text != '' THEN aibridge_interceptions.model = @model::text
ELSE true
END
-- Filter client
AND CASE
WHEN @client::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = @client::text
ELSE true
END
-- Cursor pagination
AND CASE
WHEN @after_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN (
@@ -219,8 +229,7 @@ SELECT
DISTINCT ON (provider, model, client)
provider,
model,
-- TODO: use the client value once we have it (see https://github.com/coder/aibridge/issues/31)
'unknown' AS client
COALESCE(client, 'Unknown') AS client
FROM
aibridge_interceptions
WHERE
@@ -242,8 +251,7 @@ WITH interceptions_in_range AS (
WHERE
provider = @provider::text
AND model = @model::text
-- TODO: use the client value once we have it (see https://github.com/coder/aibridge/issues/31)
AND 'unknown' = @client::text
AND COALESCE(client, 'Unknown') = @client::text
AND ended_at IS NOT NULL -- incomplete interceptions are not included in summaries
AND ended_at >= @ended_at_after::timestamptz
AND ended_at < @ended_at_before::timestamptz
+1
View File
@@ -385,6 +385,7 @@ func AIBridgeInterceptions(ctx context.Context, db database.Store, query string,
filter.InitiatorID = parseUser(ctx, db, parser, values, "initiator", actorID)
filter.Provider = parser.String(values, "", "provider")
filter.Model = parser.String(values, "", "model")
filter.Client = parser.String(values, "", "client")
// Time must be between started_after and started_before.
filter.StartedAfter = parser.Time3339Nano(values, time.Time{}, "started_after")
+2 -2
View File
@@ -376,7 +376,7 @@ func TestTelemetry(t *testing.T) {
require.Equal(t, snapshot1.Provider, aiBridgeInterception1.Provider)
require.Equal(t, snapshot1.Model, aiBridgeInterception1.Model)
require.Equal(t, snapshot1.Client, "unknown") // no client info yet
require.Equal(t, snapshot1.Client, "Unknown") // no client info yet
require.EqualValues(t, snapshot1.InterceptionCount, 2)
require.EqualValues(t, snapshot1.InterceptionsByRoute, map[string]int64{}) // no route info yet
require.EqualValues(t, snapshot1.InterceptionDurationMillis.P50, 90_000)
@@ -396,7 +396,7 @@ func TestTelemetry(t *testing.T) {
require.Equal(t, snapshot2.Provider, aiBridgeInterception3.Provider)
require.Equal(t, snapshot2.Model, aiBridgeInterception3.Model)
require.Equal(t, snapshot2.Client, "unknown") // no client info yet
require.Equal(t, snapshot2.Client, "Unknown") // no client info yet
require.EqualValues(t, snapshot2.InterceptionCount, 1)
require.EqualValues(t, snapshot2.InterceptionsByRoute, map[string]int64{}) // no route info yet
require.EqualValues(t, snapshot2.InterceptionDurationMillis.P50, 180_000)
+4
View File
@@ -75,6 +75,7 @@ type AIBridgeListInterceptionsFilter struct {
StartedAfter time.Time `json:"started_after,omitempty" format:"date-time"`
Provider string `json:"provider,omitempty"`
Model string `json:"model,omitempty"`
Client string `json:"client,omitempty"`
FilterQuery string `json:"q,omitempty"`
}
@@ -101,6 +102,9 @@ func (f AIBridgeListInterceptionsFilter) asRequestOption() RequestOption {
if f.Model != "" {
params = append(params, fmt.Sprintf("model:%q", f.Model))
}
if f.Client != "" {
params = append(params, fmt.Sprintf("client:%q", f.Client))
}
if f.FilterQuery != "" {
// If custom stuff is added, just add it on here.
params = append(params, f.FilterQuery)
+1
View File
@@ -135,6 +135,7 @@ func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Reques
InitiatorID: filter.InitiatorID,
Provider: filter.Provider,
Model: filter.Model,
Client: filter.Client,
})
if err != nil {
return xerrors.Errorf("count authorized aibridge interceptions: %w", err)
+18
View File
@@ -10,6 +10,7 @@ import (
"github.com/google/uuid"
"github.com/stretchr/testify/require"
aiblib "github.com/coder/aibridge"
"github.com/coder/coder/v2/coderd/coderdtest"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/db2sdk"
@@ -433,6 +434,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
Provider: "two",
Model: "two",
StartedAt: now.Add(-time.Hour),
Client: sql.NullString{String: aiblib.ClientCursor, Valid: true},
}, &now)
i3 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{
ID: uuid.MustParse("00000000-0000-0000-0000-000000000003"),
@@ -440,6 +442,7 @@ func TestAIBridgeListInterceptions(t *testing.T) {
Provider: "three",
Model: "three",
StartedAt: now.Add(-2 * time.Hour),
Client: sql.NullString{String: aiblib.ClientClaude, Valid: true},
}, &now)
// Convert to SDK types for response comparison. We don't care about the
@@ -498,6 +501,21 @@ func TestAIBridgeListInterceptions(t *testing.T) {
filter: codersdk.AIBridgeListInterceptionsFilter{Model: "three"},
want: []codersdk.AIBridgeInterception{i3SDK},
},
{
name: "Client/Unknown",
filter: codersdk.AIBridgeListInterceptionsFilter{Client: "Unknown"},
want: []codersdk.AIBridgeInterception{i1SDK},
},
{
name: "Client/Match",
filter: codersdk.AIBridgeListInterceptionsFilter{Client: aiblib.ClientCursor},
want: []codersdk.AIBridgeInterception{i2SDK},
},
{
name: "Client/NoMatch",
filter: codersdk.AIBridgeListInterceptionsFilter{Client: "nonsense"},
want: []codersdk.AIBridgeInterception{},
},
{
name: "StartedAfter/NoMatch",
filter: codersdk.AIBridgeListInterceptionsFilter{
+1 -1
View File
@@ -473,7 +473,7 @@ require (
github.com/anthropics/anthropic-sdk-go v1.19.0
github.com/brianvoe/gofakeit/v7 v7.14.0
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225
github.com/coder/aibridge v1.0.3
github.com/coder/aibridge v1.0.4
github.com/coder/aisdk-go v0.0.9
github.com/coder/boundary v0.8.0
github.com/coder/preview v1.0.4
+2 -2
View File
@@ -927,8 +927,8 @@ github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5 h1:6xNmx7iTtyBRev0+D/T
github.com/cncf/xds/go v0.0.0-20251210132809-ee656c7534f5/go.mod h1:KdCmV+x/BuvyMxRnYBlmVaq4OLiKW6iRQfvC62cvdkI=
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225 h1:tRIViZ5JRmzdOEo5wUWngaGEFBG8OaE1o2GIHN5ujJ8=
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225/go.mod h1:rNLVpYgEVeu1Zk29K64z6Od8RBP9DwqCu9OfCzh8MR4=
github.com/coder/aibridge v1.0.3 h1:gt3XKbnFBJ/jyls/yanU/iWZO5yhd6LVYuTQbEZ/SxQ=
github.com/coder/aibridge v1.0.3/go.mod h1:c7Of2xfAksZUrPWN180Eh60fiKgzs7dyOjniTjft6AE=
github.com/coder/aibridge v1.0.4 h1:i2GiDFfqRE6t86FKSv9N/5BM0yN865E8lAeBIi+hLp8=
github.com/coder/aibridge v1.0.4/go.mod h1:c7Of2xfAksZUrPWN180Eh60fiKgzs7dyOjniTjft6AE=
github.com/coder/aisdk-go v0.0.9 h1:Vzo/k2qwVGLTR10ESDeP2Ecek1SdPfZlEjtTfMveiVo=
github.com/coder/aisdk-go v0.0.9/go.mod h1:KF6/Vkono0FJJOtWtveh5j7yfNrSctVTpwgweYWSp5M=
github.com/coder/boundary v0.8.0 h1:g/H6VIGY4IoWeKkbvao7zhO1BAQe7upSHfHzoAZxdik=