@sixb/pg 0.1.21 → 0.1.22

package/src/agents/threads.ts

@@ -1,24 +1,40 @@
 import {
+  assertConversationThreadRow,
+  participantsToAdd,
+  planAgentThreadFork,
+  remapArtifactIds,
   snapshotAgentThreadSandboxParams,
   transitionAgentThreadSandbox,
 } from "@sixb/core/internal/agent-run-storage-provider"
 import {
+  type AddAgentThreadParticipantsInput,
+  type AgentContextCheckpointRecord,
   AgentStorageError,
+  type AgentThreadParticipantRecord,
   type AgentThreadRecord,
   type AgentThreadSandboxState,
   type AgentThreadStore,
   type CreateAgentThreadInput,
+  type ForkAgentThreadInput,
+  type ListAgentThreadChildrenInput,
+  type ListAgentThreadParticipantsInput,
   type ListAgentThreadsInput,
   type ListAgentThreadsResult,
+  type RemoveAgentThreadParticipantInput,
   type TransitionAgentThreadSandboxInput,
 } from "@sixb/core/storage"
-import type { SqlParameter } from "../pg-client"
-import { isUniqueViolation } from "../storage-errors"
+import type { SQLClient, SqlParameter } from "../pg-client"
+import { isUniqueViolation, isUniqueViolationOf } from "../storage-errors"
 import { type PgStoreClient, runPgTransaction } from "../transactions"
 import {
+  type AgentContextCheckpointRow,
+  type AgentMessageRow,
   type AgentRunRow,
+  type AgentThreadParticipantRow,
   type AgentThreadRow,
   queryAgentList,
+  rowToContextCheckpointRecord,
+  rowToParticipantRecord,
   rowToRunRecord,
   rowToThreadRecord,
 } from "./rows"
@@ -100,6 +116,165 @@ export class PgAgentThreadStore implements AgentThreadStore {
     }
   }
 
+  async fork(input: ForkAgentThreadInput): Promise<AgentThreadRecord> {
+    try {
+      return await runPgTransaction(this.sql, async (tx) => {
+        // A share lock holds back the writes that take the thread for update: an admitted request
+        // or a finalized answer lands before the copy or after it, never inside it.
+        const [source] = await tx<Pick<AgentThreadRow, "id" | "kind" | "active_run_id">[]>`
+          SELECT id, kind, active_run_id FROM agent_threads
+          WHERE project_id = ${input.projectId} AND id = ${input.sourceThreadId}
+          FOR SHARE
+        `
+
+        if (!source) {
+          throw new AgentStorageError(
+            "thread_not_found",
+            `[SixbPg] Agent thread '${input.sourceThreadId}' not found for project '${input.projectId}'.`
+          )
+        }
+        assertConversationThreadRow(source, "can be forked", "SixbPg")
+
+        const [activeRun] =
+          source.active_run_id === null
+            ? []
+            : await tx<{ trigger_message_id: string | null }[]>`
+                SELECT trigger_message_id FROM agent_runs
+                WHERE project_id = ${input.projectId} AND id = ${source.active_run_id}
+              `
+
+        // Only what the plan needs: the parts are copied by the database, not through here. A
+        // message of approval outcomes only is not counted (`countsAsThreadMessage`).
+        const messages = await tx<(Pick<AgentMessageRow, "id" | "seq"> & { counted: boolean })[]>`
+          SELECT id, seq,
+            NOT (
+              jsonb_array_length(parts) > 0 AND NOT EXISTS (
+                SELECT 1 FROM jsonb_array_elements(parts) AS part
+                WHERE part->>'type' <> 'approval-outcome'
+              )
+            ) AS counted
+          FROM agent_messages
+          WHERE project_id = ${input.projectId} AND thread_id = ${source.id}
+          ORDER BY seq ASC
+        `
+        const checkpoints = await tx<AgentContextCheckpointRow[]>`
+          SELECT * FROM agent_context_checkpoints
+          WHERE project_id = ${input.projectId} AND thread_id = ${source.id}
+          ORDER BY summarized_through_seq DESC, created_at DESC, id DESC
+        `
+
+        const plan = planAgentThreadFork(
+          input,
+          {
+            ...(activeRun?.trigger_message_id
+              ? { triggerMessageId: activeRun.trigger_message_id }
+              : {}),
+            messages: messages.map((message) => ({
+              id: message.id,
+              seq: Number(message.seq),
+              counted: message.counted,
+            })),
+            checkpoints: checkpoints.map(rowToContextCheckpointRecord),
+          },
+          "SixbPg"
+        )
+        const { thread } = plan
+
+        const [row] = await tx<AgentThreadRow[]>`
+          INSERT INTO agent_threads (
+            project_id,
+            id,
+            owner_principal_type,
+            owner_principal_id,
+            title,
+            sandbox_params,
+            sandbox_state,
+            status,
+            last_message_at,
+            message_count,
+            created_at,
+            updated_at
+          ) VALUES (
+            ${thread.projectId},
+            ${thread.id},
+            ${thread.ownerPrincipal.type},
+            ${thread.ownerPrincipal.id},
+            ${thread.title ?? null},
+            ${thread.sandboxParams === undefined ? null : JSON.stringify(thread.sandboxParams)}::text::jsonb,
+            ${thread.sandboxState === undefined ? null : JSON.stringify(thread.sandboxState)}::text::jsonb,
+            ${thread.status},
+            ${thread.lastMessageAt ?? null},
+            ${thread.messageCount},
+            ${thread.createdAt},
+            ${thread.updatedAt}
+          )
+          RETURNING *
+        `
+
+        if (plan.messages.length > 0) {
+          await tx`
+            INSERT INTO agent_messages (
+              project_id,
+              id,
+              thread_id,
+              run_id,
+              role,
+              author_principal_type,
+              author_principal_id,
+              seq,
+              parts,
+              metadata,
+              content_version,
+              created_at,
+              completed_at
+            )
+            SELECT
+              message.project_id,
+              copy.id,
+              ${thread.id},
+              message.run_id,
+              message.role,
+              message.author_principal_type,
+              message.author_principal_id,
+              message.seq,
+              message.parts,
+              message.metadata,
+              message.content_version,
+              message.created_at,
+              message.completed_at
+            FROM unnest(
+              ${tx.array(plan.messages.map((copy) => copy.sourceId))}::text[],
+              ${tx.array(plan.messages.map((copy) => copy.id))}::text[]
+            ) AS copy(source_id, id)
+            JOIN agent_messages AS message
+              ON message.project_id = ${input.projectId} AND message.id = copy.source_id
+          `
+          await remapCopiedParts(
+            tx,
+            input.projectId,
+            plan.messages.map((copy) => copy.id),
+            input.artifactCopies
+          )
+        }
+
+        if (plan.checkpoint) {
+          await insertCopiedCheckpoint(tx, plan.checkpoint)
+        }
+
+        return rowToThreadRecord(row)
+      })
+    } catch (error) {
+      if (isUniqueViolationOf(error, "agent_threads_pkey")) {
+        throw new AgentStorageError(
+          "duplicate_id",
+          `[SixbPg] Agent thread '${input.id}' already exists for project '${input.projectId}'.`
+        )
+      }
+
+      throw error
+    }
+  }
+
   async getById(params: { projectId: string; id: string }): Promise<AgentThreadRecord | null> {
     const [row] = await this.sql<AgentThreadRow[]>`
       SELECT * FROM agent_threads WHERE project_id = ${params.projectId} AND id = ${params.id}
@@ -131,12 +306,29 @@ export class PgAgentThreadStore implements AgentThreadStore {
     })
   }
 
+  async listChildren(input: ListAgentThreadChildrenInput): Promise<readonly AgentThreadRecord[]> {
+    // By `idx_agent_threads_children`, which holds sub-agents' threads only.
+    const rows = await this.sql<AgentThreadRow[]>`
+      SELECT *
+      FROM agent_threads
+      WHERE project_id = ${input.projectId}
+        AND parent_thread_id = ${input.parentThreadId}
+        AND kind = 'subagent'
+        ${input.parentRunId === undefined ? this.sql`` : this.sql`AND parent_run_id = ${input.parentRunId}`}
+      ORDER BY created_at ASC, id ASC
+      LIMIT ${input.limit ?? 100}
+      OFFSET ${input.offset ?? 0}
+    `
+    return rows.map(rowToThreadRecord)
+  }
+
   async list(input: ListAgentThreadsInput): Promise<ListAgentThreadsResult> {
     if (input.statuses !== undefined && input.statuses.length === 0) {
       return { threads: [], hasMore: false, total: 0 }
     }
 
-    const whereClauses = ["project_id = $1"]
+    // A sub-agent's thread is read through its conversation; the list indexes leave it out too.
+    const whereClauses = ["project_id = $1", "kind = 'conversation'"]
     const params: SqlParameter[] = [input.projectId]
     let index = 2
 
@@ -151,6 +343,17 @@ export class PgAgentThreadStore implements AgentThreadStore {
       params.push(input.ownerPrincipal.type, input.ownerPrincipal.id)
     }
 
+    if (input.participantPrincipal) {
+      whereClauses.push(`EXISTS (
+        SELECT 1 FROM agent_thread_participants participant
+        WHERE participant.project_id = agent_threads.project_id
+          AND participant.thread_id = agent_threads.id
+          AND participant.principal_type = $${index++}
+          AND participant.principal_id = $${index++}
+      )`)
+      params.push(input.participantPrincipal.type, input.participantPrincipal.id)
+    }
+
     const dir = input.order === "asc" ? "ASC" : "DESC"
     const { rows, total, hasMore } = await queryAgentList<AgentThreadRow>({
       sql: this.sql,
@@ -165,4 +368,154 @@ export class PgAgentThreadStore implements AgentThreadStore {
 
     return { threads: rows.map(rowToThreadRecord), hasMore, total }
   }
+
+  async addParticipants(
+    input: AddAgentThreadParticipantsInput
+  ): Promise<readonly AgentThreadParticipantRecord[]> {
+    return runPgTransaction(this.sql, async (tx) => {
+      // Concurrent adds to one thread queue on its row, so the limit counts every participant.
+      const [thread] = await tx<Pick<AgentThreadRow, "id" | "kind">[]>`
+        SELECT id, kind FROM agent_threads
+        WHERE project_id = ${input.projectId} AND id = ${input.threadId}
+        FOR UPDATE
+      `
+      if (!thread) {
+        throw new AgentStorageError(
+          "thread_not_found",
+          `[SixbPg] Agent thread '${input.threadId}' not found for project '${input.projectId}'.`
+        )
+      }
+      assertConversationThreadRow(thread, "has participants", "SixbPg")
+
+      const existing = await selectParticipants(tx, input.projectId, [input.threadId])
+      const added = participantsToAdd(input, existing, "SixbPg")
+      const addedAt = input.addedAt ?? new Date()
+
+      for (const principal of added) {
+        await tx`
+          INSERT INTO agent_thread_participants (
+            project_id,
+            thread_id,
+            principal_type,
+            principal_id,
+            added_at,
+            added_by_type,
+            added_by_id
+          ) VALUES (
+            ${input.projectId},
+            ${input.threadId},
+            ${principal.type},
+            ${principal.id},
+            ${addedAt},
+            ${input.addedBy.type},
+            ${input.addedBy.id}
+          )
+        `
+      }
+
+      return added.length === 0
+        ? existing
+        : selectParticipants(tx, input.projectId, [input.threadId])
+    })
+  }
+
+  async removeParticipant(input: RemoveAgentThreadParticipantInput): Promise<boolean> {
+    const rows = await this.sql<{ principal_id: string }[]>`
+      DELETE FROM agent_thread_participants
+      WHERE project_id = ${input.projectId}
+        AND thread_id = ${input.threadId}
+        AND principal_type = ${input.principal.type}
+        AND principal_id = ${input.principal.id}
+      RETURNING principal_id
+    `
+
+    return rows.length > 0
+  }
+
+  async listParticipants(
+    input: ListAgentThreadParticipantsInput
+  ): Promise<readonly AgentThreadParticipantRecord[]> {
+    return selectParticipants(this.sql, input.projectId, [...new Set(input.threadIds)])
+  }
+}
+
+async function selectParticipants(
+  sql: PgStoreClient,
+  projectId: string,
+  threadIds: readonly string[]
+): Promise<AgentThreadParticipantRecord[]> {
+  if (threadIds.length === 0) {
+    return []
+  }
+
+  const rows = await sql<AgentThreadParticipantRow[]>`
+    SELECT *
+    FROM agent_thread_participants
+    WHERE project_id = ${projectId}
+      AND thread_id = ANY(${sql.array([...threadIds])}::text[])
+    ORDER BY added_at, principal_id COLLATE "C", thread_id COLLATE "C"
+  `
+
+  return rows.map(rowToParticipantRecord)
+}
+
+async function insertCopiedCheckpoint(
+  tx: SQLClient,
+  checkpoint: AgentContextCheckpointRecord
+): Promise<void> {
+  await tx`
+    INSERT INTO agent_context_checkpoints (
+      project_id,
+      id,
+      thread_id,
+      created_by_run_id,
+      reason,
+      summary,
+      summary_format_version,
+      summarized_through_seq,
+      observed_head_seq,
+      estimated_input_tokens_before,
+      estimated_input_tokens_after,
+      summary_model_id,
+      created_at
+    ) VALUES (
+      ${checkpoint.projectId},
+      ${checkpoint.id},
+      ${checkpoint.threadId},
+      ${checkpoint.createdByRunId},
+      ${checkpoint.reason},
+      ${checkpoint.summary},
+      ${checkpoint.summaryFormatVersion},
+      ${checkpoint.summarizedThroughSeq},
+      ${checkpoint.observedHeadSeq},
+      ${checkpoint.estimatedInputTokensBefore},
+      ${checkpoint.estimatedInputTokensAfter},
+      ${checkpoint.summaryModelId},
+      ${checkpoint.createdAt}
+    )
+  `
+}
+
+/** The copied messages name the fork's copies of the source's artifacts. */
+async function remapCopiedParts(
+  sql: SQLClient,
+  projectId: string,
+  messageIds: readonly string[],
+  copies: Readonly<Record<string, string>> | undefined
+): Promise<void> {
+  if (!copies || Object.keys(copies).length === 0 || messageIds.length === 0) return
+
+  const rows = await sql<{ id: string; parts: string }[]>`
+    SELECT id, parts::text AS parts FROM agent_messages
+    WHERE project_id = ${projectId} AND id = ANY(${sql.array([...messageIds])}::text[])
+  `
+  for (const row of rows) {
+    const parts = remapArtifactIds(row.parts, copies)
+    if (parts === row.parts) continue
+
+    await sql`
+      UPDATE agent_messages SET parts = ${parts}::text::jsonb
+      WHERE project_id = ${projectId} AND id = ${row.id}
+    `
+  }
 }

package/src/ai-cost-groups.ts

@@ -3,6 +3,7 @@ import {
   aiModelCallGroupFragmentFromRow,
   buildAiModelCallGroupsFromFragments,
   normalizeAiModelCallAccountingQuery,
+  TOKENLESS_MODEL_KINDS,
 } from "@sixb/core/internal/ai-cost-storage-provider"
 import type { ListAiModelCallGroupsInput, ListAiModelCallGroupsResult } from "@sixb/core/storage"
 import type { PgStoreClient } from "./transactions"
@@ -34,11 +35,15 @@ export async function listAiModelCallGroups(
   return { items, total, hasMore: query.offset + items.length < total }
 }
 
+// Fixed kind names, never user input.
+const TOKENLESS_KINDS = TOKENLESS_MODEL_KINDS.map((kind) => `'${kind}'`).join(", ")
+
 const FILTERED_SQL = `
 WITH filtered AS (
   SELECT COALESCE(root_agent.execution_id, usage.execution_id) AS root_execution_id,
     usage.execution_id, usage.provider_id, usage.requested_model_id, usage.occurred_at,
-    usage.total_tokens,
+    usage.total_tokens, usage.model_kind, usage.output_images, usage.input_characters,
+    usage.audio_duration_ms,
     execution.executor_kind, execution.executor_id, execution.authority_primitive_id AS primitive_id,
     COALESCE(cost.status, 'unvalued') AS valuation_status,
     cost.currency AS amount_currency, cost.amount_nanos,
@@ -84,7 +89,12 @@ SELECT filtered.root_execution_id, execution_id, provider_id, requested_model_id
   MIN(occurred_at) AS first_call_at, MAX(occurred_at) AS last_call_at,
   CAST(COUNT(*) AS TEXT) AS call_count,
   CAST(COUNT(total_tokens) AS TEXT) AS token_count,
+  CAST(COUNT(CASE WHEN total_tokens IS NULL AND model_kind IN (${TOKENLESS_KINDS}) THEN 1 END) AS TEXT)
+    AS tokenless_count,
   CAST(SUM(total_tokens) AS TEXT) AS total_tokens,
+  CAST(SUM(output_images) AS TEXT) AS output_images,
+  CAST(SUM(input_characters) AS TEXT) AS input_characters,
+  CAST(SUM(audio_duration_ms) AS TEXT) AS audio_duration_ms,
   valuation_status, amount_currency,
   CAST(SUM(amount_nanos / 1000000000000) AS TEXT) AS amount_high,
   CAST(SUM((amount_nanos / 1000000) % 1000000) AS TEXT) AS amount_middle,

package/src/auth-storage/auth-storage.ts

@@ -11,6 +11,7 @@ import type {
   InvitationRecord,
   SessionRecord,
   SuspendUserAndRevokeSessionsInput,
+  UpdateAuthUserResult,
   UserIdentityRecord,
   UserRecord,
 } from "@sixb/core/storage"
@@ -256,16 +257,17 @@ export class PgAuthStorage implements AuthStorage {
           consumedAt: completedAt,
         })
 
-        const user = existingUserRow
-          ? rowToUserRecord(existingUserRow)
-          : await this.insertUser(tx, {
-              id: newUserId ?? assertNonEmpty(input.newUserId, "User id"),
-              projectId,
-              email,
-              displayName: input.newUserDisplayName,
-              avatarUrl: input.newUserAvatarUrl,
-              createdAt: completedAt,
-            })
+        const previousUser = existingUserRow ? rowToUserRecord(existingUserRow) : null
+        const user =
+          previousUser ??
+          (await this.insertUser(tx, {
+            id: newUserId ?? assertNonEmpty(input.newUserId, "User id"),
+            projectId,
+            email,
+            displayName: input.newUserDisplayName,
+            avatarUrl: input.newUserAvatarUrl,
+            createdAt: completedAt,
+          }))
 
         const invitation = await this.acceptInvitationAndApplyGroups(tx, {
           activeInvitation,
@@ -293,6 +295,7 @@ export class PgAuthStorage implements AuthStorage {
           session,
           invitation: invitation.invitation,
           groupMemberships,
+          previousUser,
         }
       }
     )
@@ -410,6 +413,7 @@ export class PgAuthStorage implements AuthStorage {
         ])
         await this.consumeOidcAttempt(input, completedAt, projectId, tx)
 
+        const previousUser = userRow ? rowToUserRecord(userRow) : null
         let user: UserRecord
         if (!userRow) {
           user = await this.insertUser(tx, {
@@ -431,7 +435,7 @@ export class PgAuthStorage implements AuthStorage {
             updated_at: user.updatedAt,
           }
         } else {
-          user = rowToUserRecord(userRow)
+          user = await this.refreshProviderProfile(tx, userRow, input, completedAt)
         }
 
         const invitation = await this.acceptInvitationAndApplyGroups(tx, {
@@ -476,6 +480,7 @@ export class PgAuthStorage implements AuthStorage {
           identity: nextIdentity,
           invitation: invitation.invitation,
           groupMemberships,
+          previousUser,
         }
       }
     )
@@ -489,7 +494,7 @@ export class PgAuthStorage implements AuthStorage {
 
   async suspendUserAndRevokeSessions(
     input: SuspendUserAndRevokeSessionsInput
-  ): Promise<UserRecord> {
+  ): Promise<UpdateAuthUserResult> {
     return runPgTransaction(this.sql, async (tx) => {
       await lockAdvisoryKeys(tx, [authLockKey("sessions", input.projectId, input.userId)])
       const existing = await getUserRowById(
@@ -523,7 +528,7 @@ export class PgAuthStorage implements AuthStorage {
         revokedAt: input.suspendedAt,
       })
 
-      return rowToUserRecord(updated)
+      return { user: rowToUserRecord(updated), previousUser: rowToUserRecord(existing) }
     })
   }
 
@@ -733,6 +738,33 @@ export class PgAuthStorage implements AuthStorage {
     })
   }
 
+  // The identity provider owns the name and picture it sends: every sign-in overwrites them, and
+  // one it leaves out keeps the stored value.
+  private async refreshProviderProfile(
+    sql: SQLClient,
+    row: PgAuthUserRow,
+    input: CompleteOidcSignInInput,
+    updatedAt: Date
+  ): Promise<UserRecord> {
+    const displayName = input.displayName ?? row.display_name
+    const avatarUrl = input.avatarUrl ?? row.avatar_url
+    if (displayName === row.display_name && avatarUrl === row.avatar_url) {
+      return rowToUserRecord(row)
+    }
+
+    const [updated] = await sql<PgAuthUserRow[]>`
+      UPDATE auth_users
+      SET display_name = ${displayName},
+          avatar_url = ${avatarUrl},
+          updated_at = ${updatedAt}
+      WHERE project_id = ${row.project_id}
+        AND id = ${row.id}
+      RETURNING *
+    `
+
+    return rowToUserRecord(updated)
+  }
+
   private async consumeOidcAttempt(
     input: CompleteOidcSignInInput,
     completedAt: Date,