@sixb/pg 0.1.21 → 0.1.22

package/src/agents/messages.ts

@@ -1,3 +1,4 @@
+import { countsAsThreadMessage } from "@sixb/core/internal/agent-run-storage-provider"
 import { AGENT_MESSAGE_CONTENT_VERSION } from "@sixb/core/internal/agents"
 import {
   type AgentMessageRecord,
@@ -91,7 +92,10 @@ export class PgAgentMessageStore implements AgentMessageStore {
 
         await tx`
           UPDATE agent_threads
-          SET message_count = message_count + 1, last_message_at = ${createdAt}, updated_at = ${createdAt}
+          SET
+            message_count = message_count + ${countsAsThreadMessage(input) ? 1 : 0},
+            last_message_at = ${createdAt},
+            updated_at = ${createdAt}
           WHERE project_id = ${input.projectId} AND id = ${input.threadId}
         `
 
@@ -124,25 +128,28 @@ export class PgAgentMessageStore implements AgentMessageStore {
         )
       }
 
-      const deleted = await tx<{ id: string }[]>`
+      const deleted = await tx<{ id: string; parts: AgentMessageRecord["parts"] | string }[]>`
         DELETE FROM agent_messages
         WHERE project_id = ${input.projectId}
           AND thread_id = ${input.threadId}
           AND run_id = ${input.runId}
-        RETURNING id
+        RETURNING id, parts
       `
       if (deleted.length === 0) {
         return 0
       }
+      // The thread counts only the messages the conversation shows: take back the ones it counted.
+      const counted = deleted.filter((row) =>
+        countsAsThreadMessage({
+          parts: typeof row.parts === "string" ? JSON.parse(row.parts) : row.parts,
+        })
+      ).length
 
       const updatedAt = new Date()
       await tx`
         UPDATE agent_threads
         SET
-          message_count = (
-            SELECT COUNT(*) FROM agent_messages
-            WHERE project_id = ${input.projectId} AND thread_id = ${input.threadId}
-          ),
+          message_count = GREATEST(message_count - ${counted}, 0),
           last_message_at = (
             SELECT created_at FROM agent_messages
             WHERE project_id = ${input.projectId} AND thread_id = ${input.threadId}

package/src/agents/rows.ts

@@ -7,6 +7,7 @@ import {
   type AgentMessageRecord,
   type AgentRunDiagnostic,
   type AgentRunRecord,
+  type AgentThreadParticipantRecord,
   type AgentThreadRecord,
   type AgentThreadSandboxState,
   type ConversationAgentRunSpec,
@@ -22,6 +23,10 @@ import type { SQLClient, SqlParameter } from "../pg-client"
 export interface AgentThreadRow {
   project_id: string
   id: string
+  kind: AgentThreadRecord["kind"]
+  parent_thread_id: string | null
+  parent_run_id: string | null
+  spawn_key: string | null
   owner_principal_type: Principal["type"]
   owner_principal_id: string
   title: string | null
@@ -35,6 +40,16 @@ export interface AgentThreadRow {
   updated_at: Date | string
 }
 
+export interface AgentThreadParticipantRow {
+  project_id: string
+  thread_id: string
+  principal_type: "user"
+  principal_id: string
+  added_at: Date | string
+  added_by_type: Principal["type"]
+  added_by_id: string
+}
+
 export interface AgentRunRow {
   project_id: string
   id: string
@@ -98,6 +113,16 @@ export function rowToThreadRecord(row: AgentThreadRow): AgentThreadRecord {
   return {
     id: row.id,
     projectId: row.project_id,
+    kind: row.kind,
+    ...(row.parent_thread_id && row.parent_run_id && row.spawn_key
+      ? {
+          parent: {
+            threadId: row.parent_thread_id,
+            spawnedBy: row.parent_run_id,
+            key: row.spawn_key,
+          },
+        }
+      : {}),
     ownerPrincipal: { type: row.owner_principal_type, id: row.owner_principal_id },
     title: row.title ?? undefined,
     ...(row.sandbox_state == null
@@ -160,12 +185,20 @@ export function rowToRunRecord(row: AgentRunRow): AgentRunRecord {
       ...(row.spec === null ? {} : { spec: parseJsonColumn<ConversationAgentRunSpec>(row.spec) }),
     }
   }
-  if (!row.parent_run_id || !row.spawn_key || row.spec === null) {
+  if (
+    !row.thread_id ||
+    !row.trigger_message_id ||
+    !row.parent_run_id ||
+    !row.spawn_key ||
+    row.spec === null
+  ) {
     throw new Error(`[SixbPg] Subagent run '${row.id}' has invalid persisted fields.`)
   }
   return {
     ...base,
     kind: "subagent",
+    threadId: row.thread_id,
+    triggerMessageId: row.trigger_message_id,
     parentRunId: row.parent_run_id,
     spawnKey: row.spawn_key,
     spec: parseJsonColumn<SubagentRunSpec>(row.spec),
@@ -220,6 +253,18 @@ export function rowToContextCheckpointRecord(
   }
 }
 
+export function rowToParticipantRecord(
+  row: AgentThreadParticipantRow
+): AgentThreadParticipantRecord {
+  return {
+    projectId: row.project_id,
+    threadId: row.thread_id,
+    principal: { type: "user", id: row.principal_id },
+    addedAt: new Date(row.added_at),
+    addedBy: { type: row.added_by_type, id: row.added_by_id },
+  }
+}
+
 function principalFromColumns(
   type: Principal["type"] | null | undefined,
   id: string | null | undefined

package/src/agents/runs.ts

@@ -1,10 +1,16 @@
 import {
+  agentRunMayHoldCalls,
   assertAgentRunExecution,
+  assertAgentRunFinishPayload,
+  assertAwaitingApprovalFinish,
   assertConversationAgentRunSpec,
+  assertConversationThreadRow,
   assertCreateSubagentRunInput,
-  assertSubagentRunResult,
+  assertSubagentThreadTakesRun,
+  newSubagentThreadRecord,
   subagentRunMatchesCreateInput,
 } from "@sixb/core/internal/agent-run-storage-provider"
+import { AGENT_MESSAGE_CONTENT_VERSION } from "@sixb/core/internal/agents"
 import { serializeSixbFailure } from "@sixb/core/internal/errors"
 import {
   AGENT_RUN_FAILURE_CODES,
@@ -20,6 +26,7 @@ import {
   type FinishQueuedAgentRunInput,
   type ListAgentRunsInput,
   type ListAgentRunsResult,
+  type ListThreadSubagentsInput,
   type ReclaimAgentRunInput,
   type StartAgentRunInput,
   type SubagentRunRecord,
@@ -28,7 +35,8 @@ import type { SQLClient, SqlParameter } from "../pg-client"
 import { appendRunListFilters, hasEmptyStatuses, queryRunList } from "../run-list-query"
 import { isUniqueViolation } from "../storage-errors"
 import { type PgStoreClient, runPgTransaction } from "../transactions"
-import { type AgentRunRow, rowToRunRecord } from "./rows"
+import { pgHoldsRun } from "./approvals"
+import { type AgentRunRow, type AgentThreadRow, rowToRunRecord, rowToThreadRecord } from "./rows"
 
 const PG_RUN_ID_BATCH_SIZE = 10_000
 
@@ -50,8 +58,8 @@ export class PgAgentRunStore implements AgentRunStore {
 
     try {
       return await runPgTransaction(this.sql, async (tx) => {
-        const [thread] = await tx<{ active_run_id: string | null }[]>`
-          SELECT active_run_id FROM agent_threads
+        const [thread] = 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.threadId}
           FOR UPDATE
         `
@@ -62,6 +70,7 @@ export class PgAgentRunStore implements AgentRunStore {
             `[SixbPg] Agent thread '${input.threadId}' not found for project '${input.projectId}'.`
           )
         }
+        assertConversationThreadRow(thread, "takes turns", "SixbPg")
         if (thread.active_run_id !== null) {
           throw new AgentStorageError(
             "active_run_exists",
@@ -190,18 +199,33 @@ export class PgAgentRunStore implements AgentRunStore {
           )
         }
 
+        const createdAt = input.createdAt ?? new Date()
+        await claimChildThread(tx, input, parent, createdAt)
         const [row] = await tx<AgentRunRow[]>`
           INSERT INTO agent_runs (
-            project_id, id, execution_id, kind, parent_run_id, spawn_key, spec,
-            status, attempt, created_at
+            project_id, id, execution_id, kind, thread_id, trigger_message_id, parent_run_id,
+            spawn_key, spec, status, attempt, created_at
           ) VALUES (
             ${input.projectId}, ${input.id}, ${input.executionId}, ${"subagent"},
-            ${input.parentRunId}, ${input.spawnKey},
+            ${input.thread.id}, ${input.trigger.id}, ${input.parentRunId}, ${input.spawnKey},
             ${JSON.stringify(input.spec)}::text::jsonb,
-            ${"queued"}, ${0}, ${input.createdAt ?? new Date()}
+            ${"queued"}, ${0}, ${createdAt}
           )
           RETURNING *
         `
+        await tx`
+          INSERT INTO agent_messages (
+            project_id, id, thread_id, run_id, role, seq, parts, metadata, content_version,
+            created_at
+          ) VALUES (
+            ${input.projectId}, ${input.trigger.id}, ${input.thread.id}, ${input.id}, ${"user"},
+            (SELECT COALESCE(MAX(seq), 0) + 1 FROM agent_messages
+              WHERE project_id = ${input.projectId} AND thread_id = ${input.thread.id}),
+            ${JSON.stringify(input.trigger.parts)}::text::jsonb,
+            ${input.trigger.metadata === undefined ? null : JSON.stringify(input.trigger.metadata)}::text::jsonb,
+            ${AGENT_MESSAGE_CONTENT_VERSION}, ${createdAt}
+          )
+        `
         const created = rowToRunRecord(row)
         if (created.kind !== "subagent") {
           throw new Error(`[SixbPg] Subagent insert '${input.id}' returned the wrong run kind.`)
@@ -309,19 +333,7 @@ export class PgAgentRunStore implements AgentRunStore {
     return runPgTransaction(this.sql, async (tx) => {
       const run = await this.lockRunning(tx, input.projectId, input.id)
 
-      if (input.status === "succeeded" && run.kind === "subagent") {
-        assertSubagentRunResult(input.result, run.id, "SixbPg")
-      }
-      if (
-        input.status === "succeeded" &&
-        run.kind === "conversation" &&
-        input.result !== undefined
-      ) {
-        throw new AgentStorageError(
-          "invalid_input",
-          `[SixbPg] Conversational Agent run '${run.id}' cannot persist a subagent result.`
-        )
-      }
+      assertAgentRunFinishPayload(run.kind, input, run.id, "SixbPg")
 
       if (run.execution_token !== input.executionToken) {
         throw new AgentStorageError(
@@ -329,6 +341,14 @@ export class PgAgentRunStore implements AgentRunStore {
           `[SixbPg] Execution token is no longer current on agent run '${input.id}'.`
         )
       }
+      if (input.finishReason === "awaiting_approval") {
+        assertAwaitingApprovalFinish({
+          runId: input.id,
+          status: input.status,
+          holdsApproval: await pgHoldsRun(tx, input.projectId, input.id),
+          prefix: "SixbPg",
+        })
+      }
 
       const [row] = await tx<AgentRunRow[]>`
         UPDATE agent_runs
@@ -347,6 +367,9 @@ export class PgAgentRunStore implements AgentRunStore {
       `
 
       await this.releaseThread(tx, run, completedAt)
+      if (agentRunMayHoldCalls(rowToRunRecord(run))) {
+        await settleFinishedRunApprovals(tx, run.kind, input, completedAt)
+      }
 
       return rowToRunRecord(row)
     })
@@ -387,6 +410,30 @@ export class PgAgentRunStore implements AgentRunStore {
     })
   }
 
+  async listThreadSubagents(
+    input: ListThreadSubagentsInput
+  ): Promise<readonly SubagentRunRecord[]> {
+    // Reached through the turns of the conversation, by thread, then each turn's children by key
+    // or, for a finish reason, by `idx_agent_runs_paused_on_question`.
+    const rows = await this.sql<AgentRunRow[]>`
+      SELECT child.* FROM agent_runs turn
+      JOIN agent_runs child
+        ON child.project_id = turn.project_id AND child.parent_run_id = turn.id
+      WHERE turn.project_id = ${input.projectId}
+        AND turn.thread_id = ${input.threadId}
+        AND child.kind = 'subagent'
+        ${input.spawnKey === undefined ? this.sql`` : this.sql`AND child.spawn_key = ${input.spawnKey}`}
+        ${input.finishReason === undefined ? this.sql`` : this.sql`AND child.finish_reason = ${input.finishReason}`}
+      ORDER BY COALESCE(turn.started_at, turn.created_at) DESC, turn.id DESC,
+        child.created_at DESC, child.id DESC
+      LIMIT ${input.limit ?? 100}
+    `
+    return rows.flatMap((row) => {
+      const run = rowToRunRecord(row)
+      return run.kind === "subagent" ? [run] : []
+    })
+  }
+
   async list(input: ListAgentRunsInput): Promise<ListAgentRunsResult> {
     if (hasEmptyStatuses(input) || input.kinds?.length === 0) {
       return { runs: [], hasMore: false, total: 0 }
@@ -466,11 +513,96 @@ export class PgAgentRunStore implements AgentRunStore {
     return row
   }
 
+  /** A turn releases its conversation, a child run its own thread. */
   private async releaseThread(sql: SQLClient, run: AgentRunRow, completedAt: Date): Promise<void> {
-    if (run.kind !== "conversation") return
     await sql`
       UPDATE agent_threads SET active_run_id = ${null}, updated_at = ${completedAt}
       WHERE project_id = ${run.project_id} AND id = ${run.thread_id} AND active_run_id = ${run.id}
     `
   }
 }
+
+/**
+ * The child's thread: its first run creates it, owned by the conversation's owner; a later run
+ * takes it while it is idle. Either way the run is its active run. Locked after the parent run, in
+ * the run -> thread order every writer follows.
+ */
+async function claimChildThread(
+  tx: SQLClient,
+  input: CreateSubagentRunInput,
+  parent: AgentRunRow,
+  createdAt: Date
+): Promise<void> {
+  if (!input.thread.create) {
+    const [row] = await tx<AgentThreadRow[]>`
+      SELECT * FROM agent_threads
+      WHERE project_id = ${input.projectId} AND id = ${input.thread.id}
+      FOR UPDATE
+    `
+    assertSubagentThreadTakesRun(
+      row ? rowToThreadRecord(row) : null,
+      input,
+      parent.thread_id ?? "",
+      "SixbPg"
+    )
+    await tx`
+      UPDATE agent_threads
+      SET active_run_id = ${input.id}, message_count = message_count + 1,
+        last_message_at = ${createdAt}, updated_at = ${createdAt}
+      WHERE project_id = ${input.projectId} AND id = ${input.thread.id}
+    `
+    return
+  }
+
+  const [conversation] = await tx<AgentThreadRow[]>`
+    SELECT * FROM agent_threads WHERE project_id = ${input.projectId} AND id = ${parent.thread_id}
+  `
+  if (!conversation) {
+    throw new AgentStorageError(
+      "thread_not_found",
+      `[SixbPg] Agent thread '${parent.thread_id}' not found for project '${input.projectId}'.`
+    )
+  }
+
+  const thread = newSubagentThreadRecord(input, rowToThreadRecord(conversation), createdAt)
+  await tx`
+    INSERT INTO agent_threads (
+      project_id, id, kind, parent_thread_id, parent_run_id, spawn_key, owner_principal_type,
+      owner_principal_id, status, active_run_id, last_message_at, message_count, created_at,
+      updated_at
+    ) VALUES (
+      ${thread.projectId}, ${thread.id}, ${"subagent"}, ${conversation.id}, ${input.parentRunId},
+      ${input.spawnKey}, ${thread.ownerPrincipal.type}, ${thread.ownerPrincipal.id},
+      ${thread.status}, ${input.id}, ${createdAt}, ${thread.messageCount}, ${createdAt},
+      ${createdAt}
+    )
+  `
+}
+
+/**
+ * A turn that pauses keeps its calls for their decision. A turn that ends any other way cancels
+ * the ones still waiting, a late one included, and never resumes on the rest. A child that fails
+ * or is cancelled takes its own pending calls with it.
+ */
+async function settleFinishedRunApprovals(
+  tx: SQLClient,
+  kind: AgentRunRecord["kind"],
+  input: FinishAgentRunInput,
+  completedAt: Date
+): Promise<void> {
+  if (input.finishReason === "awaiting_approval") return
+  const turn = kind === "conversation"
+  if (!turn && input.status === "succeeded") return
+
+  await tx`
+    UPDATE agent_approvals SET status = 'cancelled'
+    WHERE project_id = ${input.projectId} AND status = 'pending'
+      AND (run_id = ${input.id} OR caller_run_id = ${input.id})
+  `
+  if (!turn) return
+  await tx`
+    UPDATE agent_approvals SET turn_stopped_at = ${completedAt}
+    WHERE project_id = ${input.projectId} AND run_id = ${input.id}
+      AND reported_in_message_id IS NULL AND turn_stopped_at IS NULL
+  `
+}

package/src/agents/settings.ts

@@ -0,0 +1,105 @@
+import { assertAgentApprovalModeInput } from "@sixb/core/internal/agent-run-storage-provider"
+import {
+  type AgentApprovalMode,
+  type AgentApprovalModeStore,
+  AgentStorageError,
+  type Principal,
+} from "@sixb/core/storage"
+import { isForeignKeyViolation } from "../storage-errors"
+import { type PgStoreClient, runPgTransaction } from "../transactions"
+
+export class PgAgentApprovalModeStore implements AgentApprovalModeStore {
+  constructor(private readonly sql: PgStoreClient) {}
+
+  async getDefault(input: {
+    readonly projectId: string
+    readonly principal: Principal
+  }): Promise<AgentApprovalMode | null> {
+    const [row] = await this.sql<{ approval_mode: AgentApprovalMode }[]>`
+      SELECT approval_mode FROM agent_user_settings
+      WHERE project_id = ${input.projectId} AND principal_type = ${input.principal.type}
+        AND principal_id = ${input.principal.id}
+    `
+    return row?.approval_mode ?? null
+  }
+
+  async setDefault(input: {
+    readonly projectId: string
+    readonly principal: Principal
+    readonly mode: AgentApprovalMode
+    readonly updatedAt?: Date
+  }): Promise<void> {
+    assertAgentApprovalModeInput(input.mode, "SixbPg")
+    await this.sql`
+      INSERT INTO agent_user_settings (
+        project_id, principal_type, principal_id, approval_mode, updated_at
+      ) VALUES (
+        ${input.projectId}, ${input.principal.type}, ${input.principal.id}, ${input.mode},
+        ${input.updatedAt ?? new Date()}
+      )
+      ON CONFLICT (project_id, principal_type, principal_id)
+      DO UPDATE SET approval_mode = EXCLUDED.approval_mode, updated_at = EXCLUDED.updated_at
+    `
+  }
+
+  async getForThread(input: {
+    readonly projectId: string
+    readonly threadId: string
+    readonly principal: Principal
+  }): Promise<AgentApprovalMode | null> {
+    const [row] = await this.sql<{ approval_mode: AgentApprovalMode }[]>`
+      SELECT approval_mode FROM agent_thread_user_settings
+      WHERE project_id = ${input.projectId} AND thread_id = ${input.threadId}
+        AND principal_type = ${input.principal.type} AND principal_id = ${input.principal.id}
+    `
+    return row?.approval_mode ?? null
+  }
+
+  async setForThread(input: {
+    readonly projectId: string
+    readonly threadId: string
+    readonly principal: Principal
+    readonly mode: AgentApprovalMode | null
+    readonly updatedAt?: Date
+  }): Promise<void> {
+    if (input.mode !== null) assertAgentApprovalModeInput(input.mode, "SixbPg")
+    try {
+      await runPgTransaction(this.sql, async (tx) => {
+        const [thread] = await tx<{ id: string }[]>`
+          SELECT id FROM agent_threads
+          WHERE project_id = ${input.projectId} AND id = ${input.threadId}
+        `
+        if (!thread) throw threadNotFound(input)
+        if (input.mode === null) {
+          await tx`
+            DELETE FROM agent_thread_user_settings
+            WHERE project_id = ${input.projectId} AND thread_id = ${input.threadId}
+              AND principal_type = ${input.principal.type} AND principal_id = ${input.principal.id}
+          `
+          return
+        }
+        await tx`
+          INSERT INTO agent_thread_user_settings (
+            project_id, thread_id, principal_type, principal_id, approval_mode, updated_at
+          ) VALUES (
+            ${input.projectId}, ${input.threadId}, ${input.principal.type}, ${input.principal.id},
+            ${input.mode}, ${input.updatedAt ?? new Date()}
+          )
+          ON CONFLICT (project_id, thread_id, principal_type, principal_id)
+          DO UPDATE SET approval_mode = EXCLUDED.approval_mode, updated_at = EXCLUDED.updated_at
+        `
+      })
+    } catch (error) {
+      // A thread deleted between the read and the insert.
+      if (!isForeignKeyViolation(error)) throw error
+      throw threadNotFound(input)
+    }
+  }
+}
+
+function threadNotFound(input: { readonly projectId: string; readonly threadId: string }) {
+  return new AgentStorageError(
+    "thread_not_found",
+    `[SixbPg] Agent thread '${input.threadId}' not found for project '${input.projectId}'.`
+  )
+}