mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: store tool call IDs to determine interception lineage (#22246)
Adds database columns and server-side logic to track interception lineage via tool call IDs. When an interception ends, the server resolves the correlating tool call ID to find the parent interception and links them via `parent_id`. New `provider_tool_call_id` column on `aibridge_tool_usages` and `parent_id` column on `aibridge_interceptions`, with indexes for lookup. `findParentInterceptionID` queries by tool call ID and filters out the current interception to find the parent. Adapted from the [coder/coder `dk/prompt_provenance_poc`](https://github.com/coder/coder/compare/main...dk/prompt_provenance_poc) branch. Depends on [coder/aibridge#188](https://github.com/coder/aibridge/pull/188). Closes https://github.com/coder/internal/issues/1334
This commit is contained in:
Danny Kopping
@@ -378,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, client
|
||||
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
|
||||
FROM
|
||||
aibridge_interceptions
|
||||
WHERE
|
||||
@@ -398,13 +398,43 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU
|
||||
&i.EndedAt,
|
||||
&i.APIKeyID,
|
||||
&i.Client,
|
||||
&i.ThreadParentID,
|
||||
&i.ThreadRootID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIBridgeInterceptionLineageByToolCallID = `-- name: GetAIBridgeInterceptionLineageByToolCallID :one
|
||||
SELECT aibridge_interceptions.id AS thread_parent_id,
|
||||
COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id) AS thread_root_id
|
||||
FROM aibridge_interceptions
|
||||
WHERE aibridge_interceptions.id = (
|
||||
SELECT interception_id FROM aibridge_tool_usages
|
||||
WHERE provider_tool_call_id = $1::text
|
||||
ORDER BY created_at DESC
|
||||
LIMIT 1
|
||||
)
|
||||
`
|
||||
|
||||
type GetAIBridgeInterceptionLineageByToolCallIDRow struct {
|
||||
ThreadParentID uuid.UUID `db:"thread_parent_id" json:"thread_parent_id"`
|
||||
ThreadRootID uuid.UUID `db:"thread_root_id" json:"thread_root_id"`
|
||||
}
|
||||
|
||||
// Look up the parent interception and the root of the thread by finding
|
||||
// which interception recorded a tool usage with the given tool call ID.
|
||||
// COALESCE ensures that if the parent has no thread_root_id (i.e. it IS
|
||||
// the root), we return its own ID as the root.
|
||||
func (q *sqlQuerier) GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (GetAIBridgeInterceptionLineageByToolCallIDRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, getAIBridgeInterceptionLineageByToolCallID, toolCallID)
|
||||
var i GetAIBridgeInterceptionLineageByToolCallIDRow
|
||||
err := row.Scan(&i.ThreadParentID, &i.ThreadRootID)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getAIBridgeInterceptions = `-- name: GetAIBridgeInterceptions :many
|
||||
SELECT
|
||||
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
|
||||
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
|
||||
FROM
|
||||
aibridge_interceptions
|
||||
`
|
||||
@@ -428,6 +458,8 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
|
||||
&i.EndedAt,
|
||||
&i.APIKeyID,
|
||||
&i.Client,
|
||||
&i.ThreadParentID,
|
||||
&i.ThreadRootID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -485,7 +517,7 @@ func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context,
|
||||
|
||||
const getAIBridgeToolUsagesByInterceptionID = `-- name: GetAIBridgeToolUsagesByInterceptionID :many
|
||||
SELECT
|
||||
id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at
|
||||
id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at, provider_tool_call_id
|
||||
FROM
|
||||
aibridge_tool_usages
|
||||
WHERE
|
||||
@@ -515,6 +547,7 @@ func (q *sqlQuerier) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context,
|
||||
&i.InvocationError,
|
||||
&i.Metadata,
|
||||
&i.CreatedAt,
|
||||
&i.ProviderToolCallID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -573,22 +606,24 @@ 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, client
|
||||
id, api_key_id, initiator_id, provider, model, metadata, started_at, client, thread_parent_id, thread_root_id
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7, $8
|
||||
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7, $8, $9::uuid, $10::uuid
|
||||
)
|
||||
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client
|
||||
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
|
||||
`
|
||||
|
||||
type InsertAIBridgeInterceptionParams struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
|
||||
Provider string `db:"provider" json:"provider"`
|
||||
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"`
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
|
||||
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
|
||||
Provider string `db:"provider" json:"provider"`
|
||||
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"`
|
||||
ThreadParentInterceptionID uuid.NullUUID `db:"thread_parent_interception_id" json:"thread_parent_interception_id"`
|
||||
ThreadRootInterceptionID uuid.NullUUID `db:"thread_root_interception_id" json:"thread_root_interception_id"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) {
|
||||
@@ -601,6 +636,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
|
||||
arg.Metadata,
|
||||
arg.StartedAt,
|
||||
arg.Client,
|
||||
arg.ThreadParentInterceptionID,
|
||||
arg.ThreadRootInterceptionID,
|
||||
)
|
||||
var i AIBridgeInterception
|
||||
err := row.Scan(
|
||||
@@ -613,6 +650,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
|
||||
&i.EndedAt,
|
||||
&i.APIKeyID,
|
||||
&i.Client,
|
||||
&i.ThreadParentID,
|
||||
&i.ThreadRootID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -661,17 +700,18 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
|
||||
|
||||
const insertAIBridgeToolUsage = `-- name: InsertAIBridgeToolUsage :one
|
||||
INSERT INTO aibridge_tool_usages (
|
||||
id, interception_id, provider_response_id, tool, server_url, input, injected, invocation_error, metadata, created_at
|
||||
id, interception_id, provider_response_id, provider_tool_call_id, tool, server_url, input, injected, invocation_error, metadata, created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, COALESCE($9::jsonb, '{}'::jsonb), $10
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9, COALESCE($10::jsonb, '{}'::jsonb), $11
|
||||
)
|
||||
RETURNING id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at
|
||||
RETURNING id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at, provider_tool_call_id
|
||||
`
|
||||
|
||||
type InsertAIBridgeToolUsageParams struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
|
||||
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
|
||||
ProviderToolCallID sql.NullString `db:"provider_tool_call_id" json:"provider_tool_call_id"`
|
||||
Tool string `db:"tool" json:"tool"`
|
||||
ServerUrl sql.NullString `db:"server_url" json:"server_url"`
|
||||
Input string `db:"input" json:"input"`
|
||||
@@ -686,6 +726,7 @@ func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBr
|
||||
arg.ID,
|
||||
arg.InterceptionID,
|
||||
arg.ProviderResponseID,
|
||||
arg.ProviderToolCallID,
|
||||
arg.Tool,
|
||||
arg.ServerUrl,
|
||||
arg.Input,
|
||||
@@ -706,6 +747,7 @@ func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBr
|
||||
&i.InvocationError,
|
||||
&i.Metadata,
|
||||
&i.CreatedAt,
|
||||
&i.ProviderToolCallID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -751,7 +793,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.client,
|
||||
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,
|
||||
visible_users.id, visible_users.username, visible_users.name, visible_users.avatar_url
|
||||
FROM
|
||||
aibridge_interceptions
|
||||
@@ -860,6 +902,8 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr
|
||||
&i.AIBridgeInterception.EndedAt,
|
||||
&i.AIBridgeInterception.APIKeyID,
|
||||
&i.AIBridgeInterception.Client,
|
||||
&i.AIBridgeInterception.ThreadParentID,
|
||||
&i.AIBridgeInterception.ThreadRootID,
|
||||
&i.VisibleUser.ID,
|
||||
&i.VisibleUser.Username,
|
||||
&i.VisibleUser.Name,
|
||||
@@ -1027,7 +1071,7 @@ func (q *sqlQuerier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Contex
|
||||
|
||||
const listAIBridgeToolUsagesByInterceptionIDs = `-- name: ListAIBridgeToolUsagesByInterceptionIDs :many
|
||||
SELECT
|
||||
id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at
|
||||
id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at, provider_tool_call_id
|
||||
FROM
|
||||
aibridge_tool_usages
|
||||
WHERE
|
||||
@@ -1057,6 +1101,7 @@ func (q *sqlQuerier) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context
|
||||
&i.InvocationError,
|
||||
&i.Metadata,
|
||||
&i.CreatedAt,
|
||||
&i.ProviderToolCallID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1119,7 +1164,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, client
|
||||
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
|
||||
`
|
||||
|
||||
type UpdateAIBridgeInterceptionEndedParams struct {
|
||||
@@ -1140,6 +1185,8 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up
|
||||
&i.EndedAt,
|
||||
&i.APIKeyID,
|
||||
&i.Client,
|
||||
&i.ThreadParentID,
|
||||
&i.ThreadRootID,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user