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: session detail API (#23203)

Browse Source
This commit is contained in:
Danny Kopping
2026-03-26 18:09:53 +02:00
committed by GitHub
parent e937f89081
commit 801e57d430
19 changed files with 2759 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/apidoc/docs.go
Generated
+224
View File
@@ -214,6 +214,58 @@ const docTemplate = `{
]
}
},
"/aibridge/sessions/{session_id}": {
"get": {
"produces": [
"application/json"
],
"tags": [
"AI Bridge"
],
"summary": "Get AI Bridge session threads",
"operationId": "get-ai-bridge-session-threads",
"parameters": [
{
"type": "string",
"description": "Session ID (client_session_id or interception UUID)",
"name": "session_id",
"in": "path",
"required": true
},
{
"type": "string",
"description": "Thread pagination cursor (forward/older)",
"name": "after_id",
"in": "query"
},
{
"type": "string",
"description": "Thread pagination cursor (backward/newer)",
"name": "before_id",
"in": "query"
},
{
"type": "integer",
"description": "Number of threads per page (default 50)",
"name": "limit",
"in": "query"
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsResponse"
}
}
},
"security": [
{
"CoderSessionToken": []
}
]
}
},
"/appearance": {
"get": {
"produces": [
@@ -12675,6 +12727,29 @@ const docTemplate = `{
}
}
},
"codersdk.AIBridgeAgenticAction": {
"type": "object",
"properties": {
"model": {
"type": "string"
},
"thinking": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeModelThought"
}
},
"token_usage": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
},
"tool_calls": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeToolCall"
}
}
}
},
"codersdk.AIBridgeAnthropicConfig": {
"type": "object",
"properties": {
@@ -12843,6 +12918,14 @@ const docTemplate = `{
}
}
},
"codersdk.AIBridgeModelThought": {
"type": "object",
"properties": {
"text": {
"type": "string"
}
}
},
"codersdk.AIBridgeOpenAIConfig": {
"type": "object",
"properties": {
@@ -12942,6 +13025,76 @@ const docTemplate = `{
}
}
},
"codersdk.AIBridgeSessionThreadsResponse": {
"type": "object",
"properties": {
"client": {
"type": "string"
},
"ended_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string"
},
"initiator": {
"$ref": "#/definitions/codersdk.MinimalUser"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"models": {
"type": "array",
"items": {
"type": "string"
}
},
"page_ended_at": {
"type": "string",
"format": "date-time"
},
"page_started_at": {
"type": "string",
"format": "date-time"
},
"providers": {
"type": "array",
"items": {
"type": "string"
}
},
"started_at": {
"type": "string",
"format": "date-time"
},
"threads": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeThread"
}
},
"token_usage_summary": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
}
}
},
"codersdk.AIBridgeSessionThreadsTokenUsage": {
"type": "object",
"properties": {
"input_tokens": {
"type": "integer"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"output_tokens": {
"type": "integer"
}
}
},
"codersdk.AIBridgeSessionTokenUsageSummary": {
"type": "object",
"properties": {
@@ -12953,6 +13106,41 @@ const docTemplate = `{
}
}
},
"codersdk.AIBridgeThread": {
"type": "object",
"properties": {
"agentic_actions": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeAgenticAction"
}
},
"ended_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string",
"format": "uuid"
},
"model": {
"type": "string"
},
"prompt": {
"type": "string"
},
"provider": {
"type": "string"
},
"started_at": {
"type": "string",
"format": "date-time"
},
"token_usage": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
}
}
},
"codersdk.AIBridgeTokenUsage": {
"type": "object",
"properties": {
@@ -12983,6 +13171,42 @@ const docTemplate = `{
}
}
},
"codersdk.AIBridgeToolCall": {
"type": "object",
"properties": {
"created_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string",
"format": "uuid"
},
"injected": {
"type": "boolean"
},
"input": {
"type": "string"
},
"interception_id": {
"type": "string",
"format": "uuid"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"provider_response_id": {
"type": "string"
},
"server_url": {
"type": "string"
},
"tool": {
"type": "string"
}
}
},
"codersdk.AIBridgeToolUsage": {
"type": "object",
"properties": {
coderd/apidoc/swagger.json
Generated
+220
View File
@@ -183,6 +183,54 @@
]
}
},
"/aibridge/sessions/{session_id}": {
"get": {
"produces": ["application/json"],
"tags": ["AI Bridge"],
"summary": "Get AI Bridge session threads",
"operationId": "get-ai-bridge-session-threads",
"parameters": [
{
"type": "string",
"description": "Session ID (client_session_id or interception UUID)",
"name": "session_id",
"in": "path",
"required": true
},
{
"type": "string",
"description": "Thread pagination cursor (forward/older)",
"name": "after_id",
"in": "query"
},
{
"type": "string",
"description": "Thread pagination cursor (backward/newer)",
"name": "before_id",
"in": "query"
},
{
"type": "integer",
"description": "Number of threads per page (default 50)",
"name": "limit",
"in": "query"
}
],
"responses": {
"200": {
"description": "OK",
"schema": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsResponse"
}
}
},
"security": [
{
"CoderSessionToken": []
}
]
}
},
"/appearance": {
"get": {
"produces": ["application/json"],
@@ -11261,6 +11309,29 @@
}
}
},
"codersdk.AIBridgeAgenticAction": {
"type": "object",
"properties": {
"model": {
"type": "string"
},
"thinking": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeModelThought"
}
},
"token_usage": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
},
"tool_calls": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeToolCall"
}
}
}
},
"codersdk.AIBridgeAnthropicConfig": {
"type": "object",
"properties": {
@@ -11429,6 +11500,14 @@
}
}
},
"codersdk.AIBridgeModelThought": {
"type": "object",
"properties": {
"text": {
"type": "string"
}
}
},
"codersdk.AIBridgeOpenAIConfig": {
"type": "object",
"properties": {
@@ -11528,6 +11607,76 @@
}
}
},
"codersdk.AIBridgeSessionThreadsResponse": {
"type": "object",
"properties": {
"client": {
"type": "string"
},
"ended_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string"
},
"initiator": {
"$ref": "#/definitions/codersdk.MinimalUser"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"models": {
"type": "array",
"items": {
"type": "string"
}
},
"page_ended_at": {
"type": "string",
"format": "date-time"
},
"page_started_at": {
"type": "string",
"format": "date-time"
},
"providers": {
"type": "array",
"items": {
"type": "string"
}
},
"started_at": {
"type": "string",
"format": "date-time"
},
"threads": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeThread"
}
},
"token_usage_summary": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
}
}
},
"codersdk.AIBridgeSessionThreadsTokenUsage": {
"type": "object",
"properties": {
"input_tokens": {
"type": "integer"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"output_tokens": {
"type": "integer"
}
}
},
"codersdk.AIBridgeSessionTokenUsageSummary": {
"type": "object",
"properties": {
@@ -11539,6 +11688,41 @@
}
}
},
"codersdk.AIBridgeThread": {
"type": "object",
"properties": {
"agentic_actions": {
"type": "array",
"items": {
"$ref": "#/definitions/codersdk.AIBridgeAgenticAction"
}
},
"ended_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string",
"format": "uuid"
},
"model": {
"type": "string"
},
"prompt": {
"type": "string"
},
"provider": {
"type": "string"
},
"started_at": {
"type": "string",
"format": "date-time"
},
"token_usage": {
"$ref": "#/definitions/codersdk.AIBridgeSessionThreadsTokenUsage"
}
}
},
"codersdk.AIBridgeTokenUsage": {
"type": "object",
"properties": {
@@ -11569,6 +11753,42 @@
}
}
},
"codersdk.AIBridgeToolCall": {
"type": "object",
"properties": {
"created_at": {
"type": "string",
"format": "date-time"
},
"id": {
"type": "string",
"format": "uuid"
},
"injected": {
"type": "boolean"
},
"input": {
"type": "string"
},
"interception_id": {
"type": "string",
"format": "uuid"
},
"metadata": {
"type": "object",
"additionalProperties": {}
},
"provider_response_id": {
"type": "string"
},
"server_url": {
"type": "string"
},
"tool": {
"type": "string"
}
}
},
"codersdk.AIBridgeToolUsage": {
"type": "object",
"properties": {
coderd/database/db2sdk/db2sdk.go
+281
View File
@@ -1097,6 +1097,287 @@ func AIBridgeToolUsage(usage database.AIBridgeToolUsage) codersdk.AIBridgeToolUs
}
}
// AIBridgeSessionThreads converts session metadata and thread interceptions
// into the threads response. It groups interceptions into threads, builds
// agentic actions from tool usages and model thoughts, and aggregates
// token usage with metadata.
func AIBridgeSessionThreads(
session database.ListAIBridgeSessionsRow,
interceptions []database.ListAIBridgeSessionThreadsRow,
tokenUsages []database.AIBridgeTokenUsage,
toolUsages []database.AIBridgeToolUsage,
userPrompts []database.AIBridgeUserPrompt,
modelThoughts []database.AIBridgeModelThought,
) codersdk.AIBridgeSessionThreadsResponse {
// Index subresources by interception ID.
tokensByInterception := make(map[uuid.UUID][]database.AIBridgeTokenUsage, len(interceptions))
for _, tu := range tokenUsages {
tokensByInterception[tu.InterceptionID] = append(tokensByInterception[tu.InterceptionID], tu)
}
toolsByInterception := make(map[uuid.UUID][]database.AIBridgeToolUsage, len(interceptions))
for _, tu := range toolUsages {
toolsByInterception[tu.InterceptionID] = append(toolsByInterception[tu.InterceptionID], tu)
}
promptsByInterception := make(map[uuid.UUID][]database.AIBridgeUserPrompt, len(interceptions))
for _, up := range userPrompts {
promptsByInterception[up.InterceptionID] = append(promptsByInterception[up.InterceptionID], up)
}
thoughtsByInterception := make(map[uuid.UUID][]database.AIBridgeModelThought, len(interceptions))
for _, mt := range modelThoughts {
thoughtsByInterception[mt.InterceptionID] = append(thoughtsByInterception[mt.InterceptionID], mt)
}
// Group interceptions by thread_id, preserving the order returned by the
// SQL query.
interceptionsByThread := make(map[uuid.UUID][]database.AIBridgeInterception, len(interceptions))
var threadIDs []uuid.UUID
for _, row := range interceptions {
if _, ok := interceptionsByThread[row.ThreadID]; !ok {
threadIDs = append(threadIDs, row.ThreadID)
}
interceptionsByThread[row.ThreadID] = append(interceptionsByThread[row.ThreadID], row.AIBridgeInterception)
}
// Build threads and track page time bounds.
threads := make([]codersdk.AIBridgeThread, 0, len(threadIDs))
var pageStartedAt, pageEndedAt *time.Time
for _, threadID := range threadIDs {
intcs := interceptionsByThread[threadID]
thread := buildAIBridgeThread(threadID, intcs, tokensByInterception, toolsByInterception, promptsByInterception, thoughtsByInterception)
for _, intc := range intcs {
if pageStartedAt == nil || intc.StartedAt.Before(*pageStartedAt) {
t := intc.StartedAt
pageStartedAt = &t
}
if intc.EndedAt.Valid {
if pageEndedAt == nil || intc.EndedAt.Time.After(*pageEndedAt) {
t := intc.EndedAt.Time
pageEndedAt = &t
}
}
}
threads = append(threads, thread)
}
// Aggregate session-level token usage metadata from all token
// usages in the session (not just the page).
sessionTokenMeta := aggregateTokenMetadata(tokenUsages)
resp := codersdk.AIBridgeSessionThreadsResponse{
ID: session.SessionID,
Initiator: MinimalUserFromVisibleUser(database.VisibleUser{
ID: session.UserID,
Username: session.UserUsername,
Name: session.UserName,
AvatarURL: session.UserAvatarUrl,
}),
Providers: session.Providers,
Models: session.Models,
Metadata: jsonOrEmptyMap(pqtype.NullRawMessage{RawMessage: session.Metadata, Valid: len(session.Metadata) > 0}),
StartedAt: session.StartedAt,
PageStartedAt: pageStartedAt,
PageEndedAt: pageEndedAt,
TokenUsageSummary: codersdk.AIBridgeSessionThreadsTokenUsage{
InputTokens: session.InputTokens,
OutputTokens: session.OutputTokens,
Metadata: sessionTokenMeta,
},
Threads: threads,
}
if resp.Providers == nil {
resp.Providers = []string{}
}
if resp.Models == nil {
resp.Models = []string{}
}
if session.Client != "" {
resp.Client = &session.Client
}
if !session.EndedAt.IsZero() {
resp.EndedAt = &session.EndedAt
}
return resp
}
func buildAIBridgeThread(
threadID uuid.UUID,
interceptions []database.AIBridgeInterception,
tokensByInterception map[uuid.UUID][]database.AIBridgeTokenUsage,
toolsByInterception map[uuid.UUID][]database.AIBridgeToolUsage,
promptsByInterception map[uuid.UUID][]database.AIBridgeUserPrompt,
thoughtsByInterception map[uuid.UUID][]database.AIBridgeModelThought,
) codersdk.AIBridgeThread {
// Find the root interception (where id == threadID) to get the
// thread prompt and model.
var rootIntc *database.AIBridgeInterception
for i := range interceptions {
if interceptions[i].ID == threadID {
rootIntc = &interceptions[i]
break
}
}
// Fallback to first interception if root not found.
if rootIntc == nil && len(interceptions) > 0 {
rootIntc = &interceptions[0]
}
thread := codersdk.AIBridgeThread{
ID: threadID,
}
if rootIntc != nil {
thread.Model = rootIntc.Model
thread.Provider = rootIntc.Provider
// Get first user prompt from root interception.
// A thread can only have one prompt, by definition, since we currently
// only store the last prompt observed in an interception.
if prompts := promptsByInterception[rootIntc.ID]; len(prompts) > 0 {
thread.Prompt = &prompts[0].Prompt
}
}
// Compute thread time bounds from interceptions.
for _, intc := range interceptions {
if thread.StartedAt.IsZero() || intc.StartedAt.Before(thread.StartedAt) {
thread.StartedAt = intc.StartedAt
}
if intc.EndedAt.Valid {
if thread.EndedAt == nil || intc.EndedAt.Time.After(*thread.EndedAt) {
t := intc.EndedAt.Time
thread.EndedAt = &t
}
}
}
// Build agentic actions grouped by interception. Each interception that
// has tool calls produces one action with all its tool calls, thinking
// blocks, and token usage.
var actions []codersdk.AIBridgeAgenticAction
for _, intc := range interceptions {
tools := toolsByInterception[intc.ID]
if len(tools) == 0 {
continue
}
// Thinking blocks for this interception.
thoughts := thoughtsByInterception[intc.ID]
thinking := make([]codersdk.AIBridgeModelThought, 0, len(thoughts))
for _, mt := range thoughts {
thinking = append(thinking, codersdk.AIBridgeModelThought{
Text: mt.Content,
})
}
// Token usage for the interception.
actionTokenUsage := aggregateTokenUsage(tokensByInterception[intc.ID])
// Build tool call list.
toolCalls := make([]codersdk.AIBridgeToolCall, 0, len(tools))
for _, tu := range tools {
toolCalls = append(toolCalls, codersdk.AIBridgeToolCall{
ID: tu.ID,
InterceptionID: tu.InterceptionID,
ProviderResponseID: tu.ProviderResponseID,
ServerURL: tu.ServerUrl.String,
Tool: tu.Tool,
Injected: tu.Injected,
Input: tu.Input,
Metadata: jsonOrEmptyMap(tu.Metadata),
CreatedAt: tu.CreatedAt,
})
}
actions = append(actions, codersdk.AIBridgeAgenticAction{
Model: intc.Model,
TokenUsage: actionTokenUsage,
Thinking: thinking,
ToolCalls: toolCalls,
})
}
if actions == nil {
// Make an empty slice so we don't serialize `null`.
actions = make([]codersdk.AIBridgeAgenticAction, 0)
}
thread.AgenticActions = actions
// Aggregate thread-level token usage.
var threadTokens []database.AIBridgeTokenUsage
for _, intc := range interceptions {
threadTokens = append(threadTokens, tokensByInterception[intc.ID]...)
}
thread.TokenUsage = aggregateTokenUsage(threadTokens)
return thread
}
// aggregateTokenUsage sums token usage rows and aggregates metadata.
func aggregateTokenUsage(tokens []database.AIBridgeTokenUsage) codersdk.AIBridgeSessionThreadsTokenUsage {
var inputTokens, outputTokens int64
for _, tu := range tokens {
inputTokens += tu.InputTokens
outputTokens += tu.OutputTokens
// TODO: once https://github.com/coder/aibridge/issues/150 lands we
// should aggregate the other token types.
}
return codersdk.AIBridgeSessionThreadsTokenUsage{
InputTokens: inputTokens,
OutputTokens: outputTokens,
Metadata: aggregateTokenMetadata(tokens),
}
}
// aggregateTokenMetadata sums all numeric values from the metadata
// JSONB across the given token usage rows by key. Nested objects are
// flattened using dot-notation (e.g. {"cache": {"read_tokens": 10}}
// becomes "cache.read_tokens"). Non-numeric leaves (strings,
// booleans, arrays, nulls) are silently skipped.
func aggregateTokenMetadata(tokens []database.AIBridgeTokenUsage) map[string]any {
sums := make(map[string]int64)
for _, tu := range tokens {
if !tu.Metadata.Valid || len(tu.Metadata.RawMessage) == 0 {
continue
}
var m map[string]json.RawMessage
if err := json.Unmarshal(tu.Metadata.RawMessage, &m); err != nil {
continue
}
flattenAndSum(sums, "", m)
}
result := make(map[string]any, len(sums))
for k, v := range sums {
result[k] = v
}
return result
}
// flattenAndSum recursively walks a JSON object and sums all numeric
// leaf values into sums, using dot-separated keys for nested objects.
func flattenAndSum(sums map[string]int64, prefix string, m map[string]json.RawMessage) {
for k, raw := range m {
key := k
if prefix != "" {
key = prefix + "." + k
}
// Try as a number first.
var n json.Number
if err := json.Unmarshal(raw, &n); err == nil {
if v, err := n.Int64(); err == nil {
sums[key] += v
}
continue
}
// Try as a nested object.
var nested map[string]json.RawMessage
if err := json.Unmarshal(raw, &nested); err == nil {
flattenAndSum(sums, key, nested)
}
// Arrays, strings, booleans, nulls are skipped.
}
}
func InvalidatedPresets(invalidatedPresets []database.UpdatePresetsLastInvalidatedAtRow) []codersdk.InvalidatedPreset {
var presets []codersdk.InvalidatedPreset
for _, p := range invalidatedPresets {
coderd/database/db2sdk/db2sdk_internal_test.go
+308
View File
@@ -0,0 +1,308 @@
package db2sdk
import (
"encoding/json"
"testing"
"github.com/google/uuid"
"github.com/sqlc-dev/pqtype"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/coderd/database"
)
func TestAggregateTokenMetadata(t *testing.T) {
t.Parallel()
t.Run("empty_input", func(t *testing.T) {
t.Parallel()
result := aggregateTokenMetadata(nil)
require.Empty(t, result)
})
t.Run("sums_across_rows", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"cache_read_tokens":100,"reasoning_tokens":50}`),
Valid: true,
},
},
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"cache_read_tokens":200,"reasoning_tokens":75}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(300), result["cache_read_tokens"])
require.Equal(t, int64(125), result["reasoning_tokens"])
require.Len(t, result, 2)
})
t.Run("skips_null_and_invalid_metadata", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{Valid: false},
},
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: nil,
Valid: true,
},
},
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"tokens":42}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(42), result["tokens"])
require.Len(t, result, 1)
})
t.Run("skips_non_integer_values", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
// Float values fail json.Number.Int64(), so they
// are silently dropped.
RawMessage: json.RawMessage(`{"good":10,"fractional":1.5}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(10), result["good"])
_, hasFractional := result["fractional"]
require.False(t, hasFractional)
})
t.Run("skips_malformed_json", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`not json`),
Valid: true,
},
},
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"tokens":5}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
// The malformed row is skipped, the valid one is counted.
require.Equal(t, int64(5), result["tokens"])
require.Len(t, result, 1)
})
t.Run("flattens_nested_objects", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"cache_read_tokens": 100,
"cache": {"creation_tokens": 40, "read_tokens": 60},
"reasoning_tokens": 50,
"tags": ["a", "b"]
}`),
Valid: true,
},
},
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"cache_read_tokens": 200,
"cache": {"creation_tokens": 10}
}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(300), result["cache_read_tokens"])
require.Equal(t, int64(50), result["reasoning_tokens"])
require.Equal(t, int64(50), result["cache.creation_tokens"])
require.Equal(t, int64(60), result["cache.read_tokens"])
// Arrays are skipped.
_, hasTags := result["tags"]
require.False(t, hasTags)
require.Len(t, result, 4)
})
t.Run("flattens_deeply_nested_objects", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"provider": {
"anthropic": {"cache_creation_tokens": 100, "cache_read_tokens": 200},
"openai": {"reasoning_tokens": 50}
},
"total": 500
}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(100), result["provider.anthropic.cache_creation_tokens"])
require.Equal(t, int64(200), result["provider.anthropic.cache_read_tokens"])
require.Equal(t, int64(50), result["provider.openai.reasoning_tokens"])
require.Equal(t, int64(500), result["total"])
require.Len(t, result, 4)
})
// Real-world provider metadata shapes from
// https://github.com/coder/aibridge/issues/150.
t.Run("aggregates_real_provider_metadata", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
// Anthropic-style: cache fields are top-level.
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 23490
}`),
Valid: true,
},
},
{
// OpenAI-style: cache fields are nested inside
// input_tokens_details.
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"input_tokens_details": {"cached_tokens": 11904}
}`),
Valid: true,
},
},
{
// Second Anthropic row to verify summing.
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{
"cache_creation_input_tokens": 500,
"cache_read_input_tokens": 10000
}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
// Anthropic fields are summed across two rows.
require.Equal(t, int64(500), result["cache_creation_input_tokens"])
require.Equal(t, int64(33490), result["cache_read_input_tokens"])
// OpenAI nested field is flattened with dot notation.
require.Equal(t, int64(11904), result["input_tokens_details.cached_tokens"])
require.Len(t, result, 3)
})
t.Run("skips_string_boolean_null_values", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"tokens":10,"name":"test","enabled":true,"nothing":null}`),
Valid: true,
},
},
}
result := aggregateTokenMetadata(tokens)
require.Equal(t, int64(10), result["tokens"])
require.Len(t, result, 1)
})
}
func TestAggregateTokenUsage(t *testing.T) {
t.Parallel()
t.Run("empty_input", func(t *testing.T) {
t.Parallel()
result := aggregateTokenUsage(nil)
require.Equal(t, int64(0), result.InputTokens)
require.Equal(t, int64(0), result.OutputTokens)
require.Empty(t, result.Metadata)
})
t.Run("sums_tokens_and_metadata", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
InputTokens: 100,
OutputTokens: 50,
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"reasoning_tokens":20}`),
Valid: true,
},
},
{
ID: uuid.New(),
InputTokens: 200,
OutputTokens: 75,
Metadata: pqtype.NullRawMessage{
RawMessage: json.RawMessage(`{"reasoning_tokens":30}`),
Valid: true,
},
},
}
result := aggregateTokenUsage(tokens)
require.Equal(t, int64(300), result.InputTokens)
require.Equal(t, int64(125), result.OutputTokens)
require.Equal(t, int64(50), result.Metadata["reasoning_tokens"])
})
t.Run("handles_rows_without_metadata", func(t *testing.T) {
t.Parallel()
tokens := []database.AIBridgeTokenUsage{
{
ID: uuid.New(),
InputTokens: 500,
OutputTokens: 200,
Metadata: pqtype.NullRawMessage{Valid: false},
},
}
result := aggregateTokenUsage(tokens)
require.Equal(t, int64(500), result.InputTokens)
require.Equal(t, int64(200), result.OutputTokens)
require.Empty(t, result.Metadata)
})
}
coderd/database/dbauthz/dbauthz.go
+19
View File
@@ -5363,6 +5363,13 @@ func (q *querier) ListAIBridgeInterceptionsTelemetrySummaries(ctx context.Contex
return q.db.ListAIBridgeInterceptionsTelemetrySummaries(ctx, arg)
}
func (q *querier) ListAIBridgeModelThoughtsByInterceptionIDs(ctx context.Context, interceptionIDs []uuid.UUID) ([]database.AIBridgeModelThought, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAibridgeInterception); err != nil {
return nil, err
}
return q.db.ListAIBridgeModelThoughtsByInterceptionIDs(ctx, interceptionIDs)
}
func (q *querier) ListAIBridgeModels(ctx context.Context, arg database.ListAIBridgeModelsParams) ([]string, error) {
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type)
if err != nil {
@@ -5371,6 +5378,14 @@ func (q *querier) ListAIBridgeModels(ctx context.Context, arg database.ListAIBri
return q.db.ListAuthorizedAIBridgeModels(ctx, arg, prep)
}
func (q *querier) ListAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams) ([]database.ListAIBridgeSessionThreadsRow, error) {
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type)
if err != nil {
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
}
return q.db.ListAuthorizedAIBridgeSessionThreads(ctx, arg, prep)
}
func (q *querier) ListAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams) ([]database.ListAIBridgeSessionsRow, error) {
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type)
if err != nil {
@@ -7221,6 +7236,10 @@ func (q *querier) CountAuthorizedAIBridgeSessions(ctx context.Context, arg datab
return q.db.CountAuthorizedAIBridgeSessions(ctx, arg, prepared)
}
func (q *querier) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionThreadsRow, error) {
return q.db.ListAuthorizedAIBridgeSessionThreads(ctx, arg, prepared)
}
func (q *querier) GetAuthorizedChats(ctx context.Context, arg database.GetChatsParams, _ rbac.PreparedAuthorized) ([]database.Chat, error) {
return q.GetChats(ctx, arg)
}
coderd/database/dbauthz/dbauthz_test.go
+20
View File
@@ -5528,6 +5528,26 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(ids).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns([]database.AIBridgeToolUsage{})
}))
s.Run("ListAIBridgeModelThoughtsByInterceptionIDs", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
ids := []uuid.UUID{{1}}
db.EXPECT().ListAIBridgeModelThoughtsByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeModelThought{}, nil).AnyTimes()
check.Args(ids).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns([]database.AIBridgeModelThought{})
}))
s.Run("ListAIBridgeSessionThreads", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
params := database.ListAIBridgeSessionThreadsParams{}
db.EXPECT().ListAuthorizedAIBridgeSessionThreads(gomock.Any(), params, gomock.Any()).Return([]database.ListAIBridgeSessionThreadsRow{}, nil).AnyTimes()
// No asserts here because SQLFilter.
check.Args(params).Asserts()
}))
s.Run("ListAuthorizedAIBridgeSessionThreads", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
params := database.ListAIBridgeSessionThreadsParams{}
db.EXPECT().ListAuthorizedAIBridgeSessionThreads(gomock.Any(), params, gomock.Any()).Return([]database.ListAIBridgeSessionThreadsRow{}, nil).AnyTimes()
// No asserts here because SQLFilter.
check.Args(params, emptyPreparedAuthorized{}).Asserts()
}))
s.Run("UpdateAIBridgeInterceptionEnded", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
intcID := uuid.UUID{1}
params := database.UpdateAIBridgeInterceptionEndedParams{ID: intcID}
coderd/database/dbgen/dbgen.go
+11
View File
@@ -1663,6 +1663,17 @@ func AIBridgeToolUsage(t testing.TB, db database.Store, seed database.InsertAIBr
return toolUsage
}
func AIBridgeModelThought(t testing.TB, db database.Store, seed database.InsertAIBridgeModelThoughtParams) database.AIBridgeModelThought {
thought, err := db.InsertAIBridgeModelThought(genCtx, database.InsertAIBridgeModelThoughtParams{
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
Content: takeFirst(seed.Content, ""),
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
})
require.NoError(t, err, "insert aibridge model thought")
return thought
}
func Task(t testing.TB, db database.Store, orig database.TaskTable) database.Task {
t.Helper()
coderd/database/dbmetrics/querymetrics.go
+24
View File
@@ -3760,6 +3760,14 @@ func (m queryMetricsStore) ListAIBridgeInterceptionsTelemetrySummaries(ctx conte
return r0, r1
}
func (m queryMetricsStore) ListAIBridgeModelThoughtsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeModelThought, error) {
start := time.Now()
r0, r1 := m.s.ListAIBridgeModelThoughtsByInterceptionIDs(ctx, interceptionIds)
m.queryLatencies.WithLabelValues("ListAIBridgeModelThoughtsByInterceptionIDs").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListAIBridgeModelThoughtsByInterceptionIDs").Inc()
return r0, r1
}
func (m queryMetricsStore) ListAIBridgeModels(ctx context.Context, arg database.ListAIBridgeModelsParams) ([]string, error) {
start := time.Now()
r0, r1 := m.s.ListAIBridgeModels(ctx, arg)
@@ -3768,6 +3776,14 @@ func (m queryMetricsStore) ListAIBridgeModels(ctx context.Context, arg database.
return r0, r1
}
func (m queryMetricsStore) ListAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams) ([]database.ListAIBridgeSessionThreadsRow, error) {
start := time.Now()
r0, r1 := m.s.ListAIBridgeSessionThreads(ctx, arg)
m.queryLatencies.WithLabelValues("ListAIBridgeSessionThreads").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListAIBridgeSessionThreads").Inc()
return r0, r1
}
func (m queryMetricsStore) ListAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams) ([]database.ListAIBridgeSessionsRow, error) {
start := time.Now()
r0, r1 := m.s.ListAIBridgeSessions(ctx, arg)
@@ -5240,6 +5256,14 @@ func (m queryMetricsStore) CountAuthorizedAIBridgeSessions(ctx context.Context,
return r0, r1
}
func (m queryMetricsStore) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionThreadsRow, error) {
start := time.Now()
r0, r1 := m.s.ListAuthorizedAIBridgeSessionThreads(ctx, arg, prepared)
m.queryLatencies.WithLabelValues("ListAuthorizedAIBridgeSessionThreads").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListAuthorizedAIBridgeSessionThreads").Inc()
return r0, r1
}
func (m queryMetricsStore) GetAuthorizedChats(ctx context.Context, arg database.GetChatsParams, prepared rbac.PreparedAuthorized) ([]database.Chat, error) {
start := time.Now()
r0, r1 := m.s.GetAuthorizedChats(ctx, arg, prepared)
coderd/database/dbmock/dbmock.go
+45
View File
@@ -7037,6 +7037,21 @@ func (mr *MockStoreMockRecorder) ListAIBridgeInterceptionsTelemetrySummaries(ctx
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeInterceptionsTelemetrySummaries", reflect.TypeOf((*MockStore)(nil).ListAIBridgeInterceptionsTelemetrySummaries), ctx, arg)
}
// ListAIBridgeModelThoughtsByInterceptionIDs mocks base method.
func (m *MockStore) ListAIBridgeModelThoughtsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeModelThought, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListAIBridgeModelThoughtsByInterceptionIDs", ctx, interceptionIds)
ret0, _ := ret[0].([]database.AIBridgeModelThought)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListAIBridgeModelThoughtsByInterceptionIDs indicates an expected call of ListAIBridgeModelThoughtsByInterceptionIDs.
func (mr *MockStoreMockRecorder) ListAIBridgeModelThoughtsByInterceptionIDs(ctx, interceptionIds any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeModelThoughtsByInterceptionIDs", reflect.TypeOf((*MockStore)(nil).ListAIBridgeModelThoughtsByInterceptionIDs), ctx, interceptionIds)
}
// ListAIBridgeModels mocks base method.
func (m *MockStore) ListAIBridgeModels(ctx context.Context, arg database.ListAIBridgeModelsParams) ([]string, error) {
m.ctrl.T.Helper()
@@ -7052,6 +7067,21 @@ func (mr *MockStoreMockRecorder) ListAIBridgeModels(ctx, arg any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeModels", reflect.TypeOf((*MockStore)(nil).ListAIBridgeModels), ctx, arg)
}
// ListAIBridgeSessionThreads mocks base method.
func (m *MockStore) ListAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams) ([]database.ListAIBridgeSessionThreadsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListAIBridgeSessionThreads", ctx, arg)
ret0, _ := ret[0].([]database.ListAIBridgeSessionThreadsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListAIBridgeSessionThreads indicates an expected call of ListAIBridgeSessionThreads.
func (mr *MockStoreMockRecorder) ListAIBridgeSessionThreads(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeSessionThreads", reflect.TypeOf((*MockStore)(nil).ListAIBridgeSessionThreads), ctx, arg)
}
// ListAIBridgeSessions mocks base method.
func (m *MockStore) ListAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams) ([]database.ListAIBridgeSessionsRow, error) {
m.ctrl.T.Helper()
@@ -7142,6 +7172,21 @@ func (mr *MockStoreMockRecorder) ListAuthorizedAIBridgeModels(ctx, arg, prepared
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAuthorizedAIBridgeModels", reflect.TypeOf((*MockStore)(nil).ListAuthorizedAIBridgeModels), ctx, arg, prepared)
}
// ListAuthorizedAIBridgeSessionThreads mocks base method.
func (m *MockStore) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionThreadsRow, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ListAuthorizedAIBridgeSessionThreads", ctx, arg, prepared)
ret0, _ := ret[0].([]database.ListAIBridgeSessionThreadsRow)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ListAuthorizedAIBridgeSessionThreads indicates an expected call of ListAuthorizedAIBridgeSessionThreads.
func (mr *MockStoreMockRecorder) ListAuthorizedAIBridgeSessionThreads(ctx, arg, prepared any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAuthorizedAIBridgeSessionThreads", reflect.TypeOf((*MockStore)(nil).ListAuthorizedAIBridgeSessionThreads), ctx, arg, prepared)
}
// ListAuthorizedAIBridgeSessions mocks base method.
func (m *MockStore) ListAuthorizedAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionsRow, error) {
m.ctrl.T.Helper()
coderd/database/modelqueries.go
+57 -1
View File
@@ -813,6 +813,7 @@ type aibridgeQuerier interface {
ListAuthorizedAIBridgeModels(ctx context.Context, arg ListAIBridgeModelsParams, prepared rbac.PreparedAuthorized) ([]string, error)
ListAuthorizedAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionsRow, error)
CountAuthorizedAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) (int64, error)
ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionThreadsRow, error)
}
func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeInterceptionsRow, error) {
@@ -1050,11 +1051,66 @@ func (q *sqlQuerier) CountAuthorizedAIBridgeSessions(ctx context.Context, arg Co
return count, nil
}
func (q *sqlQuerier) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionThreadsRow, error) {
authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{
VariableConverter: regosql.AIBridgeInterceptionConverter(),
})
if err != nil {
return nil, xerrors.Errorf("compile authorized filter: %w", err)
}
filtered, err := insertAuthorizedFilter(listAIBridgeSessionThreads, fmt.Sprintf(" AND %s", authorizedFilter))
if err != nil {
return nil, xerrors.Errorf("insert authorized filter: %w", err)
}
query := fmt.Sprintf("-- name: ListAuthorizedAIBridgeSessionThreads :many\n%s", filtered)
rows, err := q.db.QueryContext(ctx, query,
arg.SessionID,
arg.AfterID,
arg.BeforeID,
arg.Limit,
)
if err != nil {
return nil, err
}
defer rows.Close()
var items []ListAIBridgeSessionThreadsRow
for rows.Next() {
var i ListAIBridgeSessionThreadsRow
if err := rows.Scan(
&i.ThreadID,
&i.AIBridgeInterception.ID,
&i.AIBridgeInterception.InitiatorID,
&i.AIBridgeInterception.Provider,
&i.AIBridgeInterception.Model,
&i.AIBridgeInterception.StartedAt,
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.AIBridgeInterception.Client,
&i.AIBridgeInterception.ThreadParentID,
&i.AIBridgeInterception.ThreadRootID,
&i.AIBridgeInterception.ClientSessionID,
&i.AIBridgeInterception.SessionID,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
func insertAuthorizedFilter(query string, replaceWith string) (string, error) {
if !strings.Contains(query, authorizedQueryPlaceholder) {
return "", xerrors.Errorf("query does not contain authorized replace string, this is not an authorized query")
}
filtered := strings.Replace(query, authorizedQueryPlaceholder, replaceWith, 1)
filtered := strings.ReplaceAll(query, authorizedQueryPlaceholder, replaceWith)
return filtered, nil
}
coderd/database/querier.go
+4
View File
@@ -778,7 +778,11 @@ type sqlcQuerier interface {
// Finds all unique AI Bridge interception telemetry summaries combinations
// (provider, model, client) in the given timeframe for telemetry reporting.
ListAIBridgeInterceptionsTelemetrySummaries(ctx context.Context, arg ListAIBridgeInterceptionsTelemetrySummariesParams) ([]ListAIBridgeInterceptionsTelemetrySummariesRow, error)
ListAIBridgeModelThoughtsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeModelThought, error)
ListAIBridgeModels(ctx context.Context, arg ListAIBridgeModelsParams) ([]string, error)
// Returns all interceptions belonging to paginated threads within a session.
// Threads are paginated by (started_at, thread_id) cursor.
ListAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams) ([]ListAIBridgeSessionThreadsRow, error)
// Returns paginated sessions with aggregated metadata, token counts, and
// the most recent user prompt. A "session" is a logical grouping of
// interceptions that share the same session_id (set by the client).
coderd/database/queries.sql.go
+148
View File
@@ -1092,6 +1092,45 @@ func (q *sqlQuerier) ListAIBridgeInterceptionsTelemetrySummaries(ctx context.Con
return items, nil
}
const listAIBridgeModelThoughtsByInterceptionIDs = `-- name: ListAIBridgeModelThoughtsByInterceptionIDs :many
SELECT
interception_id, content, metadata, created_at
FROM
aibridge_model_thoughts
WHERE
interception_id = ANY($1::uuid[])
ORDER BY
created_at ASC
`
func (q *sqlQuerier) ListAIBridgeModelThoughtsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeModelThought, error) {
rows, err := q.db.QueryContext(ctx, listAIBridgeModelThoughtsByInterceptionIDs, pq.Array(interceptionIds))
if err != nil {
return nil, err
}
defer rows.Close()
var items []AIBridgeModelThought
for rows.Next() {
var i AIBridgeModelThought
if err := rows.Scan(
&i.InterceptionID,
&i.Content,
&i.Metadata,
&i.CreatedAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listAIBridgeModels = `-- name: ListAIBridgeModels :many
SELECT
model
@@ -1146,6 +1185,115 @@ func (q *sqlQuerier) ListAIBridgeModels(ctx context.Context, arg ListAIBridgeMod
return items, nil
}
const listAIBridgeSessionThreads = `-- name: ListAIBridgeSessionThreads :many
WITH paginated_threads AS (
SELECT
-- Find thread root interceptions (thread_root_id IS NULL), apply cursor
-- pagination, and return the page.
aibridge_interceptions.id AS thread_id,
aibridge_interceptions.started_at
FROM
aibridge_interceptions
WHERE
aibridge_interceptions.session_id = $1::text
AND aibridge_interceptions.ended_at IS NOT NULL
AND aibridge_interceptions.thread_root_id IS NULL
-- Pagination cursor.
AND ($2::uuid = '00000000-0000-0000-0000-000000000000'::uuid OR
(aibridge_interceptions.started_at, aibridge_interceptions.id) > (
(SELECT started_at FROM aibridge_interceptions ai2 WHERE ai2.id = $2),
$2::uuid
)
)
AND ($3::uuid = '00000000-0000-0000-0000-000000000000'::uuid OR
(aibridge_interceptions.started_at, aibridge_interceptions.id) < (
(SELECT started_at FROM aibridge_interceptions ai2 WHERE ai2.id = $3),
$3::uuid
)
)
-- @authorize_filter
ORDER BY
aibridge_interceptions.started_at ASC,
aibridge_interceptions.id ASC
LIMIT COALESCE(NULLIF($4::integer, 0), 50)
)
SELECT
COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id) AS thread_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, aibridge_interceptions.thread_parent_id, aibridge_interceptions.thread_root_id, aibridge_interceptions.client_session_id, aibridge_interceptions.session_id
FROM
aibridge_interceptions
JOIN
paginated_threads pt
ON pt.thread_id = COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id)
WHERE
aibridge_interceptions.session_id = $1::text
AND aibridge_interceptions.ended_at IS NOT NULL
-- @authorize_filter
ORDER BY
-- Ensure threads and their associated interceptions (agentic loops) are sorted chronologically.
pt.started_at ASC,
pt.thread_id ASC,
aibridge_interceptions.started_at ASC,
aibridge_interceptions.id ASC
`
type ListAIBridgeSessionThreadsParams struct {
SessionID string `db:"session_id" json:"session_id"`
AfterID uuid.UUID `db:"after_id" json:"after_id"`
BeforeID uuid.UUID `db:"before_id" json:"before_id"`
Limit int32 `db:"limit_" json:"limit_"`
}
type ListAIBridgeSessionThreadsRow struct {
ThreadID uuid.UUID `db:"thread_id" json:"thread_id"`
AIBridgeInterception AIBridgeInterception `db:"aibridge_interception" json:"aibridge_interception"`
}
// Returns all interceptions belonging to paginated threads within a session.
// Threads are paginated by (started_at, thread_id) cursor.
func (q *sqlQuerier) ListAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams) ([]ListAIBridgeSessionThreadsRow, error) {
rows, err := q.db.QueryContext(ctx, listAIBridgeSessionThreads,
arg.SessionID,
arg.AfterID,
arg.BeforeID,
arg.Limit,
)
if err != nil {
return nil, err
}
defer rows.Close()
var items []ListAIBridgeSessionThreadsRow
for rows.Next() {
var i ListAIBridgeSessionThreadsRow
if err := rows.Scan(
&i.ThreadID,
&i.AIBridgeInterception.ID,
&i.AIBridgeInterception.InitiatorID,
&i.AIBridgeInterception.Provider,
&i.AIBridgeInterception.Model,
&i.AIBridgeInterception.StartedAt,
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.AIBridgeInterception.Client,
&i.AIBridgeInterception.ThreadParentID,
&i.AIBridgeInterception.ThreadRootID,
&i.AIBridgeInterception.ClientSessionID,
&i.AIBridgeInterception.SessionID,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listAIBridgeSessions = `-- name: ListAIBridgeSessions :many
WITH filtered_interceptions AS (
SELECT
coderd/database/queries/aibridge.sql
+64
View File
@@ -592,6 +592,70 @@ LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100)
OFFSET @offset_
;
-- name: ListAIBridgeSessionThreads :many
-- Returns all interceptions belonging to paginated threads within a session.
-- Threads are paginated by (started_at, thread_id) cursor.
WITH paginated_threads AS (
SELECT
-- Find thread root interceptions (thread_root_id IS NULL), apply cursor
-- pagination, and return the page.
aibridge_interceptions.id AS thread_id,
aibridge_interceptions.started_at
FROM
aibridge_interceptions
WHERE
aibridge_interceptions.session_id = @session_id::text
AND aibridge_interceptions.ended_at IS NOT NULL
AND aibridge_interceptions.thread_root_id IS NULL
-- Pagination cursor.
AND (@after_id::uuid = '00000000-0000-0000-0000-000000000000'::uuid OR
(aibridge_interceptions.started_at, aibridge_interceptions.id) > (
(SELECT started_at FROM aibridge_interceptions ai2 WHERE ai2.id = @after_id),
@after_id::uuid
)
)
AND (@before_id::uuid = '00000000-0000-0000-0000-000000000000'::uuid OR
(aibridge_interceptions.started_at, aibridge_interceptions.id) < (
(SELECT started_at FROM aibridge_interceptions ai2 WHERE ai2.id = @before_id),
@before_id::uuid
)
)
-- @authorize_filter
ORDER BY
aibridge_interceptions.started_at ASC,
aibridge_interceptions.id ASC
LIMIT COALESCE(NULLIF(@limit_::integer, 0), 50)
)
SELECT
COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id) AS thread_id,
sqlc.embed(aibridge_interceptions)
FROM
aibridge_interceptions
JOIN
paginated_threads pt
ON pt.thread_id = COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id)
WHERE
aibridge_interceptions.session_id = @session_id::text
AND aibridge_interceptions.ended_at IS NOT NULL
-- @authorize_filter
ORDER BY
-- Ensure threads and their associated interceptions (agentic loops) are sorted chronologically.
pt.started_at ASC,
pt.thread_id ASC,
aibridge_interceptions.started_at ASC,
aibridge_interceptions.id ASC
;
-- name: ListAIBridgeModelThoughtsByInterceptionIDs :many
SELECT
*
FROM
aibridge_model_thoughts
WHERE
interception_id = ANY(@interception_ids::uuid[])
ORDER BY
created_at ASC;
-- name: ListAIBridgeModels :many
SELECT
model