@sixb/pg 0.1.20 → 0.1.22

package/src/agents/memory.ts

@@ -0,0 +1,175 @@
+import type { Principal } from "@sixb/core"
+import {
+  agentMemoryConflicts,
+  agentMemoryHistoryCutoff,
+  assertSaveAgentMemoryInput,
+} from "@sixb/core/internal/agent-run-storage-provider"
+import type {
+  AgentMemoryRepoRecord,
+  AgentMemoryRevisionRecord,
+  AgentMemoryStore,
+  ListAgentMemoryRevisionsInput,
+  ListAgentMemoryRevisionsResult,
+  SaveAgentMemoryInput,
+  SaveAgentMemoryResult,
+} from "@sixb/core/storage"
+import { lockAdvisoryKeys, type PgStoreClient, runPgTransaction } from "../transactions"
+
+interface RevisionRow {
+  readonly revision: number
+  readonly message: string
+  readonly author_principal_type: Principal["type"]
+  readonly author_principal_id: string
+  readonly source: string | null
+  readonly created_at: Date
+}
+
+export class PgAgentMemoryStore implements AgentMemoryStore {
+  constructor(private readonly sql: PgStoreClient) {}
+
+  async get(input: {
+    projectId: string
+    repoId: string
+    asOf?: Date
+    revision?: number
+  }): Promise<AgentMemoryRepoRecord> {
+    const asOf = input.asOf ?? null
+    const upTo = input.revision ?? null
+    // One statement, so the head and the files come from one snapshot.
+    const rows = await this.sql<{ head: number; path: string | null; contents: string | null }[]>`
+      WITH head AS (
+        SELECT COALESCE(MAX(revision), 0)::int AS revision FROM agent_memory_revisions
+        WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+          AND (${asOf}::timestamptz IS NULL OR created_at <= ${asOf}::timestamptz)
+          AND (${upTo}::int IS NULL OR revision <= ${upTo}::int)
+      ), latest AS (
+        SELECT DISTINCT ON (f.path) f.path, f.contents FROM agent_memory_files f, head
+        WHERE f.project_id = ${input.projectId} AND f.repo_id = ${input.repoId}
+          AND f.revision <= head.revision
+        ORDER BY f.path, f.revision DESC
+      )
+      SELECT head.revision AS head, latest.path, latest.contents
+      FROM head
+      LEFT JOIN latest ON latest.contents IS NOT NULL
+      ORDER BY latest.path
+    `
+    return {
+      revision: rows[0]?.head ?? 0,
+      files: rows.flatMap(({ path, contents }) =>
+        path === null || contents === null ? [] : [{ path, contents }]
+      ),
+    }
+  }
+
+  async save(input: SaveAgentMemoryInput): Promise<SaveAgentMemoryResult> {
+    assertSaveAgentMemoryInput(input, "SixbPg")
+    return runPgTransaction(this.sql, async (tx) => {
+      // A repo's first save has no row to lock, so saves to one repo take its advisory lock.
+      await lockAdvisoryKeys(tx, [`agent-memory:${input.projectId}:${input.repoId}`])
+      const [{ head } = { head: 0 }] = await tx<{ head: number }[]>`
+        SELECT COALESCE(MAX(revision), 0)::int AS head FROM agent_memory_revisions
+        WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+      `
+      const paths = input.changes.map((change) => change.path)
+      const latest = await tx<{ path: string; revision: number; contents: string | null }[]>`
+        SELECT DISTINCT ON (path) path, revision, contents FROM agent_memory_files
+        WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+          AND path IN ${tx(paths)}
+        ORDER BY path, revision DESC
+      `
+      const conflicts = agentMemoryConflicts(
+        input,
+        head,
+        new Map(latest.map((file) => [file.path, file]))
+      )
+      if (conflicts.length > 0) return { saved: false, revision: head, conflicts }
+
+      const revision = head + 1
+      const createdAt = input.createdAt ?? new Date()
+      await tx`
+        INSERT INTO agent_memory_revisions (
+          project_id, repo_id, revision, message, author_principal_type, author_principal_id,
+          source, created_at
+        ) VALUES (
+          ${input.projectId}, ${input.repoId}, ${revision}, ${input.message},
+          ${input.author.type}, ${input.author.id}, ${input.source ?? null}, ${createdAt}
+        )
+      `
+      await tx`
+        INSERT INTO agent_memory_files (project_id, repo_id, path, revision, contents)
+        SELECT ${input.projectId}, ${input.repoId}, change.path, ${revision}, change.contents
+        FROM jsonb_to_recordset(${JSON.stringify(input.changes)}::text::jsonb)
+          AS change(path TEXT, contents TEXT)
+      `
+      await tx`
+        DELETE FROM agent_memory_files AS f
+        WHERE f.project_id = ${input.projectId} AND f.repo_id = ${input.repoId}
+          AND f.revision IN (
+            SELECT revision FROM agent_memory_revisions
+            WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+              AND created_at < ${agentMemoryHistoryCutoff(createdAt)}
+          )
+          AND EXISTS (
+            SELECT 1 FROM agent_memory_files n
+            WHERE n.project_id = f.project_id AND n.repo_id = f.repo_id AND n.path = f.path
+              AND n.revision > f.revision
+          )
+      `
+      return { saved: true, revision }
+    })
+  }
+
+  async listRevisions(
+    input: ListAgentMemoryRevisionsInput
+  ): Promise<ListAgentMemoryRevisionsResult> {
+    return runPgTransaction(
+      this.sql,
+      async (tx) => {
+        const [counted] = await tx<{ total: number }[]>`
+          SELECT COUNT(*)::int AS total FROM agent_memory_revisions
+          WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+        `
+        const offset = input.offset ?? 0
+        const rows = await tx<RevisionRow[]>`
+          SELECT revision, message, author_principal_type, author_principal_id, source, created_at
+          FROM agent_memory_revisions
+          WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+          ORDER BY revision DESC
+          LIMIT ${input.limit ?? null} OFFSET ${offset}
+        `
+        const changes =
+          rows.length === 0
+            ? []
+            : await tx<{ revision: number; path: string; deleted: boolean }[]>`
+                SELECT revision, path, contents IS NULL AS deleted FROM agent_memory_files
+                WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+                  AND revision IN ${tx(rows.map((row) => row.revision))}
+                ORDER BY path
+              `
+        const total = counted?.total ?? 0
+        const revisions = rows.map(
+          (row): AgentMemoryRevisionRecord => ({
+            revision: row.revision,
+            message: row.message,
+            author: { type: row.author_principal_type, id: row.author_principal_id },
+            ...(row.source === null ? {} : { source: row.source }),
+            changes: changes
+              .filter((change) => change.revision === row.revision)
+              .map((change) => ({ path: change.path, deleted: change.deleted })),
+            createdAt: new Date(row.created_at),
+          })
+        )
+        return { revisions, hasMore: offset + rows.length < total, total }
+      },
+      { isolation: "repeatableRead" }
+    )
+  }
+
+  async delete(input: { projectId: string; repoId: string }): Promise<void> {
+    // File versions go with their revisions through ON DELETE CASCADE.
+    await this.sql`
+      DELETE FROM agent_memory_revisions
+      WHERE project_id = ${input.projectId} AND repo_id = ${input.repoId}
+    `
+  }
+}

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}'.`
+  )
+}