agent-relay-server 0.33.0 → 0.34.0

package/src/db/workspaces.ts

@@ -0,0 +1,440 @@
+import { Database } from "bun:sqlite";
+import { randomUUID } from "node:crypto";
+import { isRecord, stringValue, isMechanicalMessageKind } from "agent-relay-sdk";
+import { ORCHESTRATOR_PROTOCOL_VERSION, VERSION } from "../config.ts";
+import { parseJson } from "../utils";
+import { isLiveIsolatedWorkspace } from "../workspace-phase";
+import {
+  CONTRACT_REQUIREMENTS,
+  contractCompatibility,
+  parseRuntimeCapabilities,
+  parseRuntimeContracts,
+  parseRuntimePackage,
+  type RuntimeContracts,
+} from "../contracts";
+import { STALE_TTL_MS, DAY_MS, CLAIM_LEASE_MS, POOL_CLAIM_LEASE_MS, WORKSPACE_MERGE_LEASE_MS } from "../config";
+import { matchAgents } from "../agent-ref";
+import { getDb } from "./connection.ts";
+import type {
+  AgentCard,
+  ActivityEvent,
+  ActivityEventInput,
+  AgentKind,
+  AgentSessionGuard,
+  Artifact,
+  ArtifactBlob,
+  ArtifactKind,
+  ArtifactLink,
+  ArtifactSensitivity,
+  ArtifactVisibility,
+  AttachmentRef,
+  ChannelBinding,
+  ChannelBindingMode,
+  ChannelRouteTarget,
+  ChatHistoryImport,
+  ChatHistoryImportEntry,
+  ChannelSummary,
+  ChannelTargetHealth,
+  CreatePairInput,
+  HealthCheck,
+  HealthReport,
+  ManagedAgent,
+  ManagedSessionExitDiagnostics,
+  Message,
+  MessageDeliveryAttempt,
+  MessageDeliveryStatus,
+  Orchestrator,
+  OrchestratorHealth,
+  OrchestratorRuntimeInput,
+  OrchestratorStatus,
+  OrchestratorUpgradeState,
+  PairActionInput,
+  PairMessageInput,
+  PairSession,
+  PairStatus,
+  RegisterAgentInput,
+  ReplyObligation,
+  RegisterOrchestratorInput,
+  SendMessageInput,
+  PollQuery,
+  SpawnApprovalMode,
+  SpawnProvider,
+  Task,
+  TaskEvent,
+  TaskSeverity,
+  TaskStatus,
+  IntegrationEventInput,
+  IntegrationSummary,
+  IntegrationTaskStats,
+  InboxDraft,
+  InboxState,
+  InboxThreadState,
+  ContextSnapshot,
+  ContextState,
+  ProviderCapabilities,
+  TaskStatusInput,
+  WorkspaceMetadata,
+  WorkspaceRecord,
+  WorkspaceStatus,
+} from "../types";
+
+export function rowToWorkspace(row: any): WorkspaceRecord {
+  return {
+    id: row.id,
+    repoRoot: row.repo_root,
+    sourceCwd: row.source_cwd,
+    worktreePath: row.worktree_path,
+    branch: row.branch ?? undefined,
+    baseRef: row.base_ref ?? undefined,
+    baseSha: row.base_sha ?? undefined,
+    mode: row.mode,
+    requestedMode: row.requested_mode ?? undefined,
+    status: row.status,
+    ownerAgentId: row.owner_agent_id ?? undefined,
+    ownerPolicyName: row.owner_policy_name ?? undefined,
+    ownerAutomationRunId: row.owner_automation_run_id ?? undefined,
+    stewardAgentId: row.steward_agent_id ?? undefined,
+    metadata: parseJson<Record<string, unknown>>(row.metadata, {}),
+    createdAt: row.created_at,
+    updatedAt: row.updated_at,
+    readyAt: row.ready_at ?? undefined,
+    cleanedAt: row.cleaned_at ?? undefined,
+  };
+}
+
+
+export function upsertWorkspaceFromManagedAgent(agent: ManagedAgent): WorkspaceRecord | null {
+  const workspace = agent.workspace;
+  if (!workspace) return null;
+  if (workspace.mode === "isolated" && workspace.id && workspace.repoRoot && workspace.worktreePath) {
+    // A live agent re-registers its workspace on every heartbeat. That must not
+    // clobber progress the workspace has made past the live editing state
+    // (a merge that opened a PR, ready/review/conflict, terminal states) nor
+    // wipe accumulated metadata such as the merge result / PR link. Only the
+    // "active" live state is refreshable from registration; everything else is
+    // preserved and metadata is merged, not replaced.
+    const existing = getWorkspace(workspace.id);
+    const preserveStatus = existing != null && existing.status !== "active";
+    // The branch and base change ONLY via the relay's own land-and-continue recycle
+    // (setWorkspaceBranch repoints to `<branch>-N` and bumps base_sha; base_ref is fixed
+    // at spawn and never re-targeted). The runner keeps re-reporting its spawn-time
+    // branch/base on every heartbeat, and the recycle returns status to "active" — so
+    // without this the next heartbeat clobbers the repoint back to the original branch,
+    // and the next land targets a deleted branch and strands the work (vent #62 follow-up).
+    // baseRef needs the same pin: a heartbeat carrying a stale/wrong base (e.g. a recycled
+    // `-N` branch when the agent was spawned from inside a managed worktree) would otherwise
+    // overwrite the true base, and the next land advances a local-only branch that never
+    // reaches origin/main (#285). Trust the existing row's branch/base over registration;
+    // only a brand-new row takes the runner's values.
+    return upsertWorkspace({
+      id: workspace.id,
+      repoRoot: workspace.repoRoot,
+      sourceCwd: workspace.sourceCwd ?? agent.cwd,
+      worktreePath: workspace.worktreePath,
+      branch: existing?.branch ?? workspace.branch,
+      baseRef: existing?.baseRef ?? workspace.baseRef,
+      baseSha: existing?.baseSha ?? workspace.baseSha,
+      mode: workspace.mode,
+      requestedMode: workspace.requestedMode,
+      status: preserveStatus ? existing!.status : (workspace.status ?? "active"),
+      ownerAgentId: agent.agentId || undefined,
+      ownerPolicyName: agent.policyName,
+      ownerAutomationRunId: agent.automationRunId,
+      metadata: { ...(existing?.metadata ?? {}), provider: agent.provider, label: agent.label, sessionName: agent.sessionName, tmuxSession: agent.tmuxSession },
+    });
+  }
+  // Shared-mode occupancy: one row per agent sharing a git repo, with no
+  // worktree. Records true repo occupancy in the Workspaces view (informational
+  // only — never a cleanup/merge target). Skip non-git launches.
+  const sharedRepoRoot = workspace.mode === "shared" ? (workspace.repoRoot ?? workspace.probe?.repoRoot) : undefined;
+  if (sharedRepoRoot && agent.agentId) {
+    return upsertWorkspace({
+      id: `shared-${agent.agentId}`,
+      repoRoot: sharedRepoRoot,
+      sourceCwd: workspace.sourceCwd ?? agent.cwd,
+      worktreePath: "",
+      mode: "shared",
+      requestedMode: workspace.requestedMode,
+      status: "active",
+      ownerAgentId: agent.agentId,
+      ownerPolicyName: agent.policyName,
+      ownerAutomationRunId: agent.automationRunId,
+      metadata: { occupancy: true, provider: agent.provider, label: agent.label, sessionName: agent.sessionName, tmuxSession: agent.tmuxSession },
+    });
+  }
+  return null;
+}
+
+export function upsertWorkspace(input: Omit<WorkspaceRecord, "createdAt" | "updatedAt"> & Partial<Pick<WorkspaceRecord, "createdAt" | "updatedAt">>): WorkspaceRecord {
+  const now = Date.now();
+  getDb().query(`
+    INSERT INTO workspaces (id, repo_root, source_cwd, worktree_path, branch, base_ref, base_sha, mode, requested_mode, status, owner_agent_id, owner_policy_name, owner_automation_run_id, steward_agent_id, metadata, created_at, updated_at, ready_at, cleaned_at)
+    VALUES ($id, $repoRoot, $sourceCwd, $worktreePath, $branch, $baseRef, $baseSha, $mode, $requestedMode, $status, $ownerAgentId, $ownerPolicyName, $ownerAutomationRunId, $stewardAgentId, $metadata, $createdAt, $updatedAt, $readyAt, $cleanedAt)
+    ON CONFLICT(id) DO UPDATE SET
+      repo_root = excluded.repo_root,
+      source_cwd = excluded.source_cwd,
+      worktree_path = excluded.worktree_path,
+      branch = excluded.branch,
+      base_ref = excluded.base_ref,
+      base_sha = excluded.base_sha,
+      mode = excluded.mode,
+      requested_mode = excluded.requested_mode,
+      status = excluded.status,
+      owner_agent_id = coalesce(excluded.owner_agent_id, workspaces.owner_agent_id),
+      owner_policy_name = coalesce(excluded.owner_policy_name, workspaces.owner_policy_name),
+      owner_automation_run_id = coalesce(excluded.owner_automation_run_id, workspaces.owner_automation_run_id),
+      steward_agent_id = coalesce(excluded.steward_agent_id, workspaces.steward_agent_id),
+      metadata = excluded.metadata,
+      updated_at = excluded.updated_at,
+      ready_at = coalesce(excluded.ready_at, workspaces.ready_at),
+      cleaned_at = coalesce(excluded.cleaned_at, workspaces.cleaned_at)
+  `).run({
+    $id: input.id,
+    $repoRoot: input.repoRoot,
+    $sourceCwd: input.sourceCwd,
+    $worktreePath: input.worktreePath,
+    $branch: input.branch ?? null,
+    $baseRef: input.baseRef ?? null,
+    $baseSha: input.baseSha ?? null,
+    $mode: input.mode,
+    $requestedMode: input.requestedMode ?? null,
+    $status: input.status,
+    $ownerAgentId: input.ownerAgentId ?? null,
+    $ownerPolicyName: input.ownerPolicyName ?? null,
+    $ownerAutomationRunId: input.ownerAutomationRunId ?? null,
+    $stewardAgentId: input.stewardAgentId ?? null,
+    $metadata: JSON.stringify(input.metadata ?? {}),
+    $createdAt: input.createdAt ?? now,
+    $updatedAt: input.updatedAt ?? now,
+    $readyAt: input.readyAt ?? null,
+    $cleanedAt: input.cleanedAt ?? null,
+  });
+  electWorkspaceStewards(input.repoRoot);
+  return getWorkspace(input.id)!;
+}
+
+export function getWorkspace(id: string): WorkspaceRecord | null {
+  const row = getDb().query("SELECT * FROM workspaces WHERE id = ?").get(id) as any;
+  return row ? rowToWorkspace(row) : null;
+}
+
+export function listWorkspaces(filter: { repoRoot?: string; ownerAgentId?: string; status?: WorkspaceStatus } = {}): WorkspaceRecord[] {
+  const where: string[] = [];
+  const params: string[] = [];
+  if (filter.repoRoot) { where.push("repo_root = ?"); params.push(filter.repoRoot); }
+  if (filter.ownerAgentId) { where.push("owner_agent_id = ?"); params.push(filter.ownerAgentId); }
+  if (filter.status) { where.push("status = ?"); params.push(filter.status); }
+  const sql = `SELECT * FROM workspaces${where.length ? " WHERE " + where.join(" AND ") : ""} ORDER BY updated_at DESC`;
+  return (getDb().query(sql).all(...params) as any[]).map(rowToWorkspace);
+}
+
+export function deleteWorkspace(id: string): boolean {
+  return getDb().query("DELETE FROM workspaces WHERE id = ?").run(id).changes > 0;
+}
+
+// The agent's current branch worktree: its most recent live (non-terminal) isolated
+// workspace, or undefined. SINGLE HOME for the agent→workspace link — the MCP
+// owner-resolver and the #236 branch-state badge both go through here. listWorkspaces
+// is ORDER BY updated_at DESC, so `.find` returns the most recently active one.
+export function ownedIsolatedWorkspace(agentId: string): WorkspaceRecord | undefined {
+  return listWorkspaces({ ownerAgentId: agentId }).find(isLiveIsolatedWorkspace);
+}
+
+// Shared-mode rows are pure occupancy markers (no worktree on disk) that only
+// mean something while their owner is online. Deletion is normally driven by
+// the reaper's onAgentDisappeared hook, but agents also leave via clean
+// SessionEnd (setStatus offline) and via pruneOfflineAgents (raw agent delete),
+// neither of which touches workspaces — so orphaned shared rows accumulate and
+// bloat the workspace panel. This sweep is the catch-all: drop any non-cleaned
+// shared row whose owner is missing or offline, regardless of how it leaked.
+export function pruneOrphanedSharedWorkspaces(): string[] {
+  return getDb().transaction(() => {
+    const orphanCondition = `
+      mode = 'shared' AND status != 'cleaned' AND (
+        owner_agent_id IS NULL
+        OR owner_agent_id NOT IN (SELECT id FROM agents)
+        OR owner_agent_id IN (SELECT id FROM agents WHERE status = 'offline')
+      )`;
+    const rows = getDb().query(`SELECT id FROM workspaces WHERE ${orphanCondition}`).all() as Array<{ id: string }>;
+    if (!rows.length) return [];
+    getDb().query(`DELETE FROM workspaces WHERE ${orphanCondition}`).run();
+    return rows.map((r) => r.id);
+  })();
+}
+
+// Late-bound listeners for "a workspace row changed in a way that may change its
+// owner's branch-state badge" (#236). A hook, not a direct emit, because getDb().ts is
+// the lowest layer — sse.ts imports db, so db can't import sse without a cycle. sse
+// registers a listener at startup that re-emits the owner agent's status over SSE.
+export type WorkspaceChangeListener = (workspace: WorkspaceRecord) => void;
+export const workspaceChangeListeners = new Set<WorkspaceChangeListener>();
+export function onWorkspaceChange(listener: WorkspaceChangeListener): void {
+  workspaceChangeListeners.add(listener);
+}
+export function emitWorkspaceChange(workspace: WorkspaceRecord | null | undefined): void {
+  if (!workspace) return;
+  for (const listener of workspaceChangeListeners) {
+    try {
+      listener(workspace);
+    } catch {
+      // A badge-refresh listener must never break a workspace write.
+    }
+  }
+}
+
+export function updateWorkspaceStatus(id: string, status: WorkspaceStatus, metadata: Record<string, unknown> = {}): WorkspaceRecord | null {
+  const existing = getWorkspace(id);
+  if (!existing) return null;
+  const nextMeta = { ...existing.metadata, ...metadata };
+  const now = Date.now();
+  getDb().query(`
+    UPDATE workspaces
+    SET status = ?, metadata = ?, updated_at = ?, ready_at = coalesce(ready_at, ?), cleaned_at = coalesce(cleaned_at, ?)
+    WHERE id = ?
+  `).run(
+    status,
+    JSON.stringify(nextMeta),
+    now,
+    status === "ready" ? now : null,
+    status === "cleaned" ? now : null,
+    id,
+  );
+  electWorkspaceStewards(existing.repoRoot);
+  const updated = getWorkspace(id);
+  // Every status transition can flip the owner's badge (active→ready→steward→merged),
+  // so refresh it regardless of caller. Fires on the row's CURRENT owner.
+  emitWorkspaceChange(updated);
+  return updated;
+}
+
+// Repoint a workspace row at a recycled branch after a land-and-continue merge
+// (#206): the worktree switched to a fresh branch cut from the advanced base, so
+// the row must track the new branch (else the next merge command targets a branch
+// that no longer exists) and the new base sha. No-op if the row is gone.
+export function setWorkspaceBranch(id: string, branch: string, baseSha?: string): WorkspaceRecord | null {
+  const existing = getWorkspace(id);
+  if (!existing) return null;
+  getDb().query(`UPDATE workspaces SET branch = ?, base_sha = coalesce(?, base_sha), updated_at = ? WHERE id = ?`)
+    .run(branch, baseSha ?? null, Date.now(), id);
+  return getWorkspace(id);
+}
+
+// Workspace statuses that count as "live" for stewardship — an agent owning one
+// of these is a candidate steward; the repo is worth coordinating.
+export const STEWARD_LIVE_STATUSES = "'active', 'ready', 'conflict', 'review_requested', 'merge_planned'";
+
+export interface RepoStewardRecord {
+  repoRoot: string;
+  stewardAgentId?: string;
+  lastStewardAgentId?: string;
+  electedAt?: number;
+  updatedAt: number;
+}
+
+export function rowToRepoSteward(row: any): RepoStewardRecord {
+  return {
+    repoRoot: row.repo_root,
+    stewardAgentId: row.steward_agent_id ?? undefined,
+    lastStewardAgentId: row.last_steward_agent_id ?? undefined,
+    electedAt: row.elected_at ?? undefined,
+    updatedAt: row.updated_at,
+  };
+}
+
+export function getRepoSteward(repoRoot: string): RepoStewardRecord | null {
+  const row = getDb().query("SELECT * FROM repo_stewards WHERE repo_root = ?").get(repoRoot) as any;
+  return row ? rowToRepoSteward(row) : null;
+}
+
+export function listRepoStewards(): RepoStewardRecord[] {
+  return (getDb().query("SELECT * FROM repo_stewards ORDER BY updated_at DESC").all() as any[]).map(rowToRepoSteward);
+}
+
+// Persist the elected steward for a repo. The row is never deleted, so a repo's
+// stewardship survives a full all-agents-offline gap (steward goes NULL/dormant,
+// last_steward_agent_id keeps continuity) and resumes on the next agent join.
+export function upsertRepoSteward(repoRoot: string, steward: string | null, now: number): void {
+  getDb().query(`
+    INSERT INTO repo_stewards (repo_root, steward_agent_id, last_steward_agent_id, elected_at, updated_at)
+    VALUES ($repoRoot, $steward, $steward, $electedAt, $now)
+    ON CONFLICT(repo_root) DO UPDATE SET
+      steward_agent_id = $steward,
+      last_steward_agent_id = coalesce($steward, repo_stewards.last_steward_agent_id),
+      elected_at = CASE
+        WHEN $steward IS NOT NULL AND $steward IS NOT repo_stewards.steward_agent_id THEN $now
+        ELSE repo_stewards.elected_at
+      END,
+      updated_at = $now
+  `).run({ $repoRoot: repoRoot, $steward: steward, $electedAt: steward ? now : null, $now: now });
+}
+
+export function electWorkspaceStewards(repoRoot?: string): void {
+  const params: string[] = repoRoot ? [repoRoot] : [];
+  const repoRows = getDb().query(`
+    SELECT DISTINCT repo_root FROM workspaces
+    WHERE status IN (${STEWARD_LIVE_STATUSES})
+    ${repoRoot ? "AND repo_root = ?" : ""}
+  `).all(...params) as Array<{ repo_root: string }>;
+  const now = Date.now();
+  for (const row of repoRows) {
+    // Candidate pool: owners of live workspaces in this repo who are online,
+    // oldest first. A steward must be an online agent actively in the repo — an
+    // offline agent can't coordinate, so it is never elected (the old bug).
+    const pool = (getDb().query(`
+      SELECT w.owner_agent_id AS id, MIN(w.created_at) AS created_at
+      FROM workspaces w JOIN agents a ON a.id = w.owner_agent_id
+      WHERE w.repo_root = ? AND w.owner_agent_id IS NOT NULL
+        AND a.status != 'offline' AND w.status IN (${STEWARD_LIVE_STATUSES})
+      GROUP BY w.owner_agent_id
+      ORDER BY created_at ASC
+    `).all(row.repo_root) as Array<{ id: string }>).map((r) => r.id);
+
+    // Keep the current steward if it is still in the pool (stable election);
+    // otherwise promote the oldest online owner, else go dormant (null).
+    const current = getRepoSteward(row.repo_root)?.stewardAgentId ?? null;
+    const steward = (current && pool.includes(current) ? current : pool[0]) ?? null;
+
+    upsertRepoSteward(row.repo_root, steward, now);
+    // Mirror onto live workspace rows only when the steward actually changed, so
+    // re-elections don't churn updated_at and reset the auto-abandon clock for a
+    // dormant repo (a stranded review_requested must still age out).
+    if (steward !== current) {
+      getDb().query(`UPDATE workspaces SET steward_agent_id = ?, updated_at = ? WHERE repo_root = ? AND status IN (${STEWARD_LIVE_STATUSES})`)
+        .run(steward, now, row.repo_root);
+    }
+  }
+}
+
+// Public re-election trigger that does not change any workspace status — used by
+// maintenance to revive a dormant steward (e.g. on the next agent join) before
+// deciding whether a stranded worktree needs escalation.
+export function reelectRepoSteward(repoRoot: string): void {
+  electWorkspaceStewards(repoRoot);
+}
+
+// Merge a metadata patch into a workspace WITHOUT bumping updated_at or running a
+// steward election. For maintenance bookkeeping (stranded/escalation markers)
+// that must not disturb age-based GC timers. undefined values delete keys.
+export function patchWorkspaceMetadata(id: string, patch: Record<string, unknown>): WorkspaceRecord | null {
+  const existing = getWorkspace(id);
+  if (!existing) return null;
+  const next = { ...existing.metadata };
+  for (const [k, v] of Object.entries(patch)) {
+    if (v === undefined) delete next[k];
+    else next[k] = v;
+  }
+  getDb().query("UPDATE workspaces SET metadata = ? WHERE id = ?").run(JSON.stringify(next), id);
+  return getWorkspace(id);
+}
+
+// Re-elect stewards for every repo where an agent owns a live workspace. Called
+// when an agent (re)registers so a dormant repo regains a steward on rejoin
+// without a full unscoped sweep.
+export function electWorkspaceStewardsForAgent(agentId: string): void {
+  const repos = getDb().query(`
+    SELECT DISTINCT repo_root FROM workspaces
+    WHERE owner_agent_id = ? AND status IN (${STEWARD_LIVE_STATUSES})
+  `).all(agentId) as Array<{ repo_root: string }>;
+  for (const r of repos) electWorkspaceStewards(r.repo_root);
+}
+