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

coderd/database/dbauthz/dbauthz.go

@@ -2263,6 +2263,13 @@ func (q *querier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID)
 	return fetch(q.log, q.auth, q.db.GetAIBridgeInterceptionByID)(ctx, id)
 }
 
+func (q *querier) GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (database.GetAIBridgeInterceptionLineageByToolCallIDRow, error) {
+	if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAibridgeInterception); err != nil {
+		return database.GetAIBridgeInterceptionLineageByToolCallIDRow{}, err
+	}
+	return q.db.GetAIBridgeInterceptionLineageByToolCallID(ctx, toolCallID)
+}
+
 func (q *querier) GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBridgeInterception, error) {
 	fetch := func(ctx context.Context, _ any) ([]database.AIBridgeInterception, error) {
 		return q.db.GetAIBridgeInterceptions(ctx)

coderd/database/dbauthz/dbauthz_test.go

@@ -5074,6 +5074,16 @@ func (s *MethodTestSuite) TestAIBridge() {
 		check.Args(intID).Asserts(intc, policy.ActionRead).Returns(intc)
 	}))
 
+	s.Run("GetAIBridgeInterceptionLineageByToolCallID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
+		toolCallID := "call_123"
+		row := database.GetAIBridgeInterceptionLineageByToolCallIDRow{
+			ThreadParentID: uuid.UUID{1},
+			ThreadRootID:   uuid.UUID{2},
+		}
+		db.EXPECT().GetAIBridgeInterceptionLineageByToolCallID(gomock.Any(), toolCallID).Return(row, nil).AnyTimes()
+		check.Args(toolCallID).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns(row)
+	}))
+
 	s.Run("GetAIBridgeInterceptions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
 		a := testutil.Fake(s.T(), faker, database.AIBridgeInterception{})
 		b := testutil.Fake(s.T(), faker, database.AIBridgeInterception{})

coderd/database/dbgen/dbgen.go

@@ -1585,14 +1585,16 @@ func ClaimPrebuild(
 
 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()),
-		APIKeyID:    seed.APIKeyID,
-		InitiatorID: takeFirst(seed.InitiatorID, uuid.New()),
-		Provider:    takeFirst(seed.Provider, "provider"),
-		Model:       takeFirst(seed.Model, "model"),
-		Metadata:    takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
-		StartedAt:   takeFirst(seed.StartedAt, dbtime.Now()),
-		Client:      seed.Client,
+		ID:                         takeFirst(seed.ID, uuid.New()),
+		APIKeyID:                   seed.APIKeyID,
+		InitiatorID:                takeFirst(seed.InitiatorID, uuid.New()),
+		Provider:                   takeFirst(seed.Provider, "provider"),
+		Model:                      takeFirst(seed.Model, "model"),
+		Metadata:                   takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
+		StartedAt:                  takeFirst(seed.StartedAt, dbtime.Now()),
+		Client:                     seed.Client,
+		ThreadParentInterceptionID: seed.ThreadParentInterceptionID,
+		ThreadRootInterceptionID:   seed.ThreadRootInterceptionID,
 	})
 	if endedAt != nil {
 		interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
@@ -1645,6 +1647,7 @@ func AIBridgeToolUsage(t testing.TB, db database.Store, seed database.InsertAIBr
 		ID:                 takeFirst(seed.ID, uuid.New()),
 		InterceptionID:     takeFirst(seed.InterceptionID, uuid.New()),
 		ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
+		ProviderToolCallID: takeFirst(seed.ProviderToolCallID),
 		Tool:               takeFirst(seed.Tool, "tool"),
 		ServerUrl:          serverURL,
 		Input:              takeFirst(seed.Input, "input"),

coderd/database/dbmetrics/querymetrics.go

@@ -791,6 +791,14 @@ func (m queryMetricsStore) GetAIBridgeInterceptionByID(ctx context.Context, id u
 	return r0, r1
 }
 
+func (m queryMetricsStore) GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (database.GetAIBridgeInterceptionLineageByToolCallIDRow, error) {
+	start := time.Now()
+	r0, r1 := m.s.GetAIBridgeInterceptionLineageByToolCallID(ctx, toolCallID)
+	m.queryLatencies.WithLabelValues("GetAIBridgeInterceptionLineageByToolCallID").Observe(time.Since(start).Seconds())
+	m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBridgeInterceptionLineageByToolCallID").Inc()
+	return r0, r1
+}
+
 func (m queryMetricsStore) GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBridgeInterception, error) {
 	start := time.Now()
 	r0, r1 := m.s.GetAIBridgeInterceptions(ctx)

coderd/database/dbmock/dbmock.go

@@ -1327,6 +1327,21 @@ func (mr *MockStoreMockRecorder) GetAIBridgeInterceptionByID(ctx, id any) *gomoc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeInterceptionByID", reflect.TypeOf((*MockStore)(nil).GetAIBridgeInterceptionByID), ctx, id)
 }
 
+// GetAIBridgeInterceptionLineageByToolCallID mocks base method.
+func (m *MockStore) GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (database.GetAIBridgeInterceptionLineageByToolCallIDRow, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetAIBridgeInterceptionLineageByToolCallID", ctx, toolCallID)
+	ret0, _ := ret[0].(database.GetAIBridgeInterceptionLineageByToolCallIDRow)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetAIBridgeInterceptionLineageByToolCallID indicates an expected call of GetAIBridgeInterceptionLineageByToolCallID.
+func (mr *MockStoreMockRecorder) GetAIBridgeInterceptionLineageByToolCallID(ctx, toolCallID any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeInterceptionLineageByToolCallID", reflect.TypeOf((*MockStore)(nil).GetAIBridgeInterceptionLineageByToolCallID), ctx, toolCallID)
+}
+
 // GetAIBridgeInterceptions mocks base method.
 func (m *MockStore) GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBridgeInterception, error) {
 	m.ctrl.T.Helper()

coderd/database/dump.sql

@@ -1044,13 +1044,19 @@ CREATE TABLE aibridge_interceptions (
     metadata jsonb,
     ended_at timestamp with time zone,
     api_key_id text,
-    client character varying(64) DEFAULT 'Unknown'::character varying
+    client character varying(64) DEFAULT 'Unknown'::character varying,
+    thread_parent_id uuid,
+    thread_root_id uuid
 );
 
 COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge';
 
 COMMENT ON COLUMN aibridge_interceptions.initiator_id IS 'Relates to a users record, but FK is elided for performance.';
 
+COMMENT ON COLUMN aibridge_interceptions.thread_parent_id IS 'The interception which directly caused this interception to occur, usually through an agentic loop or threaded conversation.';
+
+COMMENT ON COLUMN aibridge_interceptions.thread_root_id IS 'The root interception of the thread that this interception belongs to.';
+
 CREATE TABLE aibridge_token_usages (
     id uuid NOT NULL,
     interception_id uuid NOT NULL,
@@ -1075,7 +1081,8 @@ CREATE TABLE aibridge_tool_usages (
     injected boolean DEFAULT false NOT NULL,
     invocation_error text,
     metadata jsonb,
-    created_at timestamp with time zone NOT NULL
+    created_at timestamp with time zone NOT NULL,
+    provider_tool_call_id text
 );
 
 COMMENT ON TABLE aibridge_tool_usages IS 'Audit log of tool calls in intercepted requests in AI Bridge';
@@ -3448,12 +3455,18 @@ CREATE INDEX idx_aibridge_interceptions_provider ON aibridge_interceptions USING
 
 CREATE INDEX idx_aibridge_interceptions_started_id_desc ON aibridge_interceptions USING btree (started_at DESC, id DESC);
 
+CREATE INDEX idx_aibridge_interceptions_thread_parent_id ON aibridge_interceptions USING btree (thread_parent_id);
+
+CREATE INDEX idx_aibridge_interceptions_thread_root_id ON aibridge_interceptions USING btree (thread_root_id);
+
 CREATE INDEX idx_aibridge_token_usages_interception_id ON aibridge_token_usages USING btree (interception_id);
 
 CREATE INDEX idx_aibridge_token_usages_provider_response_id ON aibridge_token_usages USING btree (provider_response_id);
 
 CREATE INDEX idx_aibridge_tool_usages_interception_id ON aibridge_tool_usages USING btree (interception_id);
 
+CREATE INDEX idx_aibridge_tool_usages_provider_tool_call_id ON aibridge_tool_usages USING btree (provider_tool_call_id);
+
 CREATE INDEX idx_aibridge_tool_usagesprovider_response_id ON aibridge_tool_usages USING btree (provider_response_id);
 
 CREATE INDEX idx_aibridge_user_prompts_interception_id ON aibridge_user_prompts USING btree (interception_id);

coderd/database/migrations/000426_aibridge_tool_call_id_correlation.down.sql

@@ -0,0 +1,12 @@
+DROP INDEX IF EXISTS idx_aibridge_tool_usages_provider_tool_call_id;
+
+ALTER TABLE aibridge_tool_usages
+DROP COLUMN provider_tool_call_id;
+
+DROP INDEX IF EXISTS idx_aibridge_interceptions_thread_root_id;
+DROP INDEX IF EXISTS idx_aibridge_interceptions_thread_parent_id;
+
+ALTER TABLE aibridge_interceptions
+DROP COLUMN thread_root_id;
+ALTER TABLE aibridge_interceptions
+DROP COLUMN thread_parent_id;

coderd/database/migrations/000426_aibridge_tool_call_id_correlation.up.sql

@@ -0,0 +1,14 @@
+ALTER TABLE aibridge_tool_usages
+ADD COLUMN provider_tool_call_id text NULL; -- nullable to allow existing data to be correct
+
+CREATE INDEX idx_aibridge_tool_usages_provider_tool_call_id ON aibridge_tool_usages (provider_tool_call_id);
+
+ALTER TABLE aibridge_interceptions
+ADD COLUMN thread_parent_id UUID NULL,
+ADD COLUMN thread_root_id UUID NULL;
+
+COMMENT ON COLUMN aibridge_interceptions.thread_parent_id IS 'The interception which directly caused this interception to occur, usually through an agentic loop or threaded conversation.';
+COMMENT ON COLUMN aibridge_interceptions.thread_root_id IS 'The root interception of the thread that this interception belongs to.';
+
+CREATE INDEX idx_aibridge_interceptions_thread_parent_id ON aibridge_interceptions (thread_parent_id);
+CREATE INDEX idx_aibridge_interceptions_thread_root_id ON aibridge_interceptions (thread_root_id);

coderd/database/modelqueries.go

@@ -813,6 +813,8 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
 			&i.AIBridgeInterception.EndedAt,
 			&i.AIBridgeInterception.APIKeyID,
 			&i.AIBridgeInterception.Client,
+			&i.AIBridgeInterception.ThreadParentID,
+			&i.AIBridgeInterception.ThreadRootID,
 			&i.VisibleUser.ID,
 			&i.VisibleUser.Username,
 			&i.VisibleUser.Name,

coderd/database/models.go

@@ -3789,6 +3789,10 @@ type AIBridgeInterception struct {
 	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"`
+	// The interception which directly caused this interception to occur, usually through an agentic loop or threaded conversation.
+	ThreadParentID uuid.NullUUID `db:"thread_parent_id" json:"thread_parent_id"`
+	// The root interception of the thread that this interception belongs to.
+	ThreadRootID uuid.NullUUID `db:"thread_root_id" json:"thread_root_id"`
 }
 
 // Audit log of tokens used by intercepted requests in AI Bridge
@@ -3816,9 +3820,10 @@ type AIBridgeToolUsage struct {
 	// Whether this tool was injected; i.e. Bridge injected these tools into the request from an MCP server. If false it means a tool was defined by the client and already existed in the request (MCP or built-in).
 	Injected bool `db:"injected" json:"injected"`
 	// Only injected tools are invoked.
-	InvocationError sql.NullString        `db:"invocation_error" json:"invocation_error"`
-	Metadata        pqtype.NullRawMessage `db:"metadata" json:"metadata"`
-	CreatedAt       time.Time             `db:"created_at" json:"created_at"`
+	InvocationError    sql.NullString        `db:"invocation_error" json:"invocation_error"`
+	Metadata           pqtype.NullRawMessage `db:"metadata" json:"metadata"`
+	CreatedAt          time.Time             `db:"created_at" json:"created_at"`
+	ProviderToolCallID sql.NullString        `db:"provider_tool_call_id" json:"provider_tool_call_id"`
 }
 
 // Audit log of prompts used by intercepted requests in AI Bridge

coderd/database/querier.go

@@ -172,6 +172,11 @@ type sqlcQuerier interface {
 	// and returns the preset with the most parameters (largest subset).
 	FindMatchingPresetID(ctx context.Context, arg FindMatchingPresetIDParams) (uuid.UUID, error)
 	GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (AIBridgeInterception, error)
+	// 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.
+	GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (GetAIBridgeInterceptionLineageByToolCallIDRow, error)
 	GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeInterception, error)
 	GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeTokenUsage, error)
 	GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeToolUsage, error)

coderd/database/queries.sql.go

@@ -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
 }

coderd/database/queries/aibridge.sql

@@ -1,8 +1,8 @@
 -- 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 (
-	@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at, @client
+	@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at, @client, sqlc.narg('thread_parent_interception_id')::uuid, sqlc.narg('thread_root_interception_id')::uuid
 )
 RETURNING *;
 
@@ -14,6 +14,21 @@ WHERE
 	AND ended_at IS NULL
 RETURNING *;
 
+-- name: GetAIBridgeInterceptionLineageByToolCallID :one
+-- 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.
+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 = @tool_call_id::text
+  ORDER BY created_at DESC
+  LIMIT 1
+);
+
 -- name: InsertAIBridgeTokenUsage :one
 INSERT INTO aibridge_token_usages (
   id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
@@ -32,9 +47,9 @@ RETURNING *;
 
 -- 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 (
-  @id, @interception_id, @provider_response_id, @tool, @server_url, @input, @injected, @invocation_error, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
+  @id, @interception_id, @provider_response_id, @provider_tool_call_id, @tool, @server_url, @input, @injected, @invocation_error, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
 )
 RETURNING *;