rnxsim 0.0.0 → 0.1.313

package/src/host/agent-host.ts

@@ -0,0 +1,666 @@
+// agent routing extension for SootSimBridgeHost.
+//
+// owns the single FIFO reader per agent session (closes the
+// "single-reader channel" gap that agent-sessions.ts explicitly warns
+// about), and fans events out to every WS subscriber — CLI `sootsim
+// agent watch`, electron main, and browser shells all consume the same
+// stream.
+//
+// message shapes follow the existing ws-bridge convention:
+//   client → daemon: { id, type: 'agent:…', …payload }
+//   daemon → client: { id, result } | { id, error, code? }
+// plus server-initiated pushes (no id):
+//   { type: 'agent:event',         sessionId, event }
+//   { type: 'agent:session-status', session }
+
+import fs from 'node:fs'
+import path from 'node:path'
+import {
+  scanDevServers,
+  type DiscoveredServer,
+} from '../../scripts/dev-server-scanner.ts'
+import {
+  AgentSessionError,
+  endSession as endAgentSession,
+  sendPrompt as sendAgentPrompt,
+  startSession as startAgentSession,
+  subscribeEvents as subscribeAgentEvents,
+  transcriptPath,
+  type Provider,
+} from '../agent-sessions.ts'
+import {
+  deleteProject,
+  findProjectById,
+  findSessionById,
+  getUserDataDir,
+  listProjects,
+  listSessions,
+  recordTurnTelemetry,
+  seedFromDemoAppRegistry,
+  applySessionStatusPatch,
+  updateSessionStatuses,
+  upsertProject,
+  type AgentSession,
+  type AttachedProject,
+} from '../attached-projects.ts'
+import {
+  sendPromptToTeamMachineSession,
+  teamMachineSessionIdFromCliIdentity,
+} from '../team-machine-prompt.ts'
+import type { AgentEvent } from '../agent-events.ts'
+import type { AgentPromptEnvelope } from '../agent-prompt.ts'
+import type { WebSocket } from 'ws'
+
+// ws.readyState constant — avoid importing the module just for this enum.
+const WS_OPEN = 1
+const SESSION_STATUS_PERSIST_DELAY_MS = 25
+
+interface FanoutSubscription {
+  unsubscribe: () => void
+  refCount: number
+}
+
+interface PendingPromptEcho {
+  sentAt: number
+}
+
+export interface AgentHostOptions {
+  /** callback so hosts with a Contrast-env context (electron, dev middleware)
+   *  can inject their own exclude list. defaults to the env-derived list
+   *  the rest of the repo uses, so `sootsim serve` standalone does the
+   *  right thing. */
+  getExcludePorts?: () => number[]
+  /** returns the active CLI lease for a sim. used only for prompt routing to
+   *  an external Team Machine agent that claimed the tab. */
+  resolveCliLease?: (simId: string) => AgentCliLeaseSnapshot | null
+}
+
+export interface AgentCliLeaseSnapshot {
+  kind: 'cli' | 'user-active'
+  cliIdentityKey: string
+  expiresAt: number
+}
+
+/** default excludes keep the dev-server scanner from hitting Contrast's own
+ *  web / zero / r2 ports. mirrored from
+ *  sootsim-engine/src/dev-scan-excludes.ts — duplicated rather than
+ *  imported because sootsim is a dependency of sootsim-engine, not the
+ *  other way around. */
+function defaultExcludePorts(): number[] {
+  return [
+    Number(process.env.VITE_PORT_WEB || process.env.PORT || 3000),
+    Number(process.env.VITE_PORT_ZERO || 7849),
+    Number(process.env.VITE_PORT_R2 || 9500),
+  ].filter((p) => Number.isFinite(p) && p > 0)
+}
+
+export class AgentHost {
+  // live fan-out, keyed by sessionId. refcount tracks how many sockets have
+  // subscribed; the FIFO reader stays open as long as refcount > 0.
+  private subscriptions = new Map<string, FanoutSubscription>()
+  // per-socket subscription set so ws.close can clean up without scanning
+  // every session.
+  private sessionsBySocket = new Map<WebSocket, Set<string>>()
+  // every connected socket (regardless of role) — gets session-status
+  // pushes so a CLI `sootsim agent sessions --watch` can react to state
+  // changes driven by another client.
+  private allSockets = new Set<WebSocket>()
+  // agent wrappers also emit prompt-received from the FIFO reader, but the
+  // shell/daemon already knows the user-facing display text when send-prompt
+  // is accepted. emit the friendly prompt immediately and suppress the raw
+  // wrapper echoes that follow a moment later.
+  private pendingPromptEchoes = new Map<string, PendingPromptEcho[]>()
+  // accepted prompts serialize inside the wrapper, so a second send while the
+  // session is already working becomes queued follow-up work. keep a small
+  // in-memory count so session status stays `working` until that backlog
+  // truly drains.
+  private pendingTurns = new Map<string, number>()
+  // status fan-out is realtime state. keep the current value in memory and
+  // batch its durable store write so an fsync cannot delay the websocket event
+  // it describes. close() flushes the final batch synchronously.
+  private pendingSessionStates = new Map<string, AgentSession>()
+  private pendingSessionPatches = new Map<string, Partial<AgentSession>>()
+  private sessionStatusPersistTimer: ReturnType<typeof setTimeout> | null = null
+  private opts: AgentHostOptions
+
+  constructor(opts: AgentHostOptions = {}) {
+    this.opts = opts
+  }
+
+  registerSocket(ws: WebSocket): void {
+    this.allSockets.add(ws)
+  }
+
+  unregisterSocket(ws: WebSocket): void {
+    const sessions = this.sessionsBySocket.get(ws)
+    if (sessions) {
+      for (const sessionId of sessions) {
+        this.decrementSubscription(sessionId)
+      }
+      this.sessionsBySocket.delete(ws)
+    }
+    this.allSockets.delete(ws)
+  }
+
+  /** handle an agent:* message. returns true iff the message was recognized
+   *  as an agent message (so the caller knows to stop dispatching). */
+  async handleMessage(ws: WebSocket, msg: any): Promise<boolean> {
+    const type = msg?.type
+    if (typeof type !== 'string' || !type.startsWith('agent:')) return false
+    const id = msg.id
+    try {
+      const result = await this.dispatch(ws, type, msg)
+      this.respond(ws, id, result)
+    } catch (err) {
+      if (err instanceof AgentSessionError) {
+        this.respondError(ws, id, err.message, err.code)
+      } else {
+        this.respondError(ws, id, err instanceof Error ? err.message : String(err))
+      }
+    }
+    return true
+  }
+
+  /** run once on daemon boot. idempotent — `seedFromDemoAppRegistry` no-ops
+   *  when the store already has projects. */
+  async seedOnBoot(): Promise<void> {
+    try {
+      await seedFromDemoAppRegistry()
+    } catch (err) {
+      process.stderr.write(
+        `[sootsim-agent] seedFromDemoAppRegistry failed: ${err instanceof Error ? err.message : String(err)}\n`,
+      )
+    }
+  }
+
+  /** terminate every subscription and drop every socket reference. called
+   *  by the bridge host during shutdown. */
+  close(): void {
+    if (this.sessionStatusPersistTimer) {
+      clearTimeout(this.sessionStatusPersistTimer)
+      this.sessionStatusPersistTimer = null
+    }
+    this.flushSessionStatuses()
+    for (const sub of this.subscriptions.values()) {
+      try {
+        sub.unsubscribe()
+      } catch {}
+    }
+    this.subscriptions.clear()
+    this.sessionsBySocket.clear()
+    this.allSockets.clear()
+  }
+
+  // --- dispatch ---
+
+  private async dispatch(ws: WebSocket, type: string, msg: any): Promise<unknown> {
+    switch (type) {
+      case 'agent:list-projects':
+        return listProjects()
+      case 'agent:upsert-project':
+        return upsertProject(msg.input ?? {})
+      case 'agent:delete-project':
+        deleteProject(String(msg.projectId))
+        return { ok: true }
+      case 'agent:auto-attach-for-url':
+        return this.autoAttachForUrl(msg.input ?? {})
+      case 'agent:list-sessions':
+        return listSessions(msg.projectId ? String(msg.projectId) : undefined)
+      case 'agent:start-session':
+        return this.doStartSession(msg.input ?? {})
+      case 'agent:send-claimed-prompt':
+        return this.sendClaimedPrompt(msg)
+      case 'agent:send-prompt': {
+        const sessionId = String(msg.sessionId)
+        const session = findSessionById(sessionId)
+        if (!session) {
+          throw new AgentSessionError('NO_SESSION', `no session: ${sessionId}`)
+        }
+        const prompt = this.normalizePromptEnvelope(msg)
+        await sendAgentPrompt(sessionId, prompt)
+        return this.notePromptAccepted(sessionId, prompt, session.status === 'working')
+      }
+      case 'agent:end-session':
+        this.dropSessionFanout(String(msg.sessionId))
+        await endAgentSession(String(msg.sessionId))
+        const ended = findSessionById(String(msg.sessionId))
+        if (ended) {
+          this.broadcastSessionStatus(ended)
+        }
+        return { ok: true }
+      case 'agent:get-transcript':
+        return this.getTranscript(String(msg.sessionId))
+      case 'agent:get-paths':
+        return this.getPaths()
+      case 'agent:subscribe-events':
+        return this.subscribeSocket(ws, String(msg.sessionId))
+      case 'agent:unsubscribe-events':
+        return this.unsubscribeSocket(ws, String(msg.sessionId))
+      default:
+        throw new AgentSessionError('UNKNOWN_AGENT_MSG', `unknown agent message: ${type}`)
+    }
+  }
+
+  // --- operation impls ---
+
+  private async sendClaimedPrompt(msg: any): Promise<{
+    ok: true
+    routed: 'team-machine'
+    sessionId: string
+  }> {
+    const simId = typeof msg.simId === 'string' ? msg.simId.trim() : ''
+    if (!simId) {
+      throw new AgentSessionError('NO_SIM', 'agent:send-claimed-prompt requires simId')
+    }
+    const lease = this.opts.resolveCliLease?.(simId) ?? null
+    if (!lease || lease.kind !== 'cli' || lease.expiresAt <= Date.now()) {
+      throw new AgentSessionError('NO_CLAIM', `sim ${simId} has no active CLI claim`)
+    }
+    const sessionId = teamMachineSessionIdFromCliIdentity(lease.cliIdentityKey)
+    if (!sessionId) {
+      throw new AgentSessionError(
+        'UNSUPPORTED_CLAIM',
+        `sim ${simId} is claimed by a CLI identity that is not promptable`,
+      )
+    }
+    await sendPromptToTeamMachineSession({
+      sessionId,
+      prompt: this.normalizePromptEnvelope(msg),
+    })
+    return { ok: true, routed: 'team-machine', sessionId }
+  }
+
+  private async doStartSession(input: {
+    projectId: string
+    provider?: Provider
+    codexBin?: string
+    claudeBin?: string
+    freshThread?: boolean
+  }): Promise<{ session: AgentSession; wrapperPid: number }> {
+    const project = findProjectById(input.projectId)
+    if (!project) {
+      throw new AgentSessionError('NO_PROJECT', `no project: ${input.projectId}`)
+    }
+    const result = await startAgentSession(input)
+    this.broadcastSessionStatus(result.session)
+    return result
+  }
+
+  private async autoAttachForUrl(input: {
+    bundleUrl?: string
+    provider?: Provider
+  }): Promise<{ project: AttachedProject | null }> {
+    const bundleUrl = input.bundleUrl ?? ''
+    const targetPort = (() => {
+      try {
+        return new URL(bundleUrl).port || null
+      } catch {
+        return null
+      }
+    })()
+    if (!targetPort) return { project: null }
+    const excludePorts = this.opts.getExcludePorts?.() ?? defaultExcludePorts()
+    const servers = await scanDevServers({ excludePorts })
+    const match = servers.find((s) => String(s.port) === targetPort)
+    if (!match || !match.cwd) return { project: null }
+    const existing = listProjects().find((p) => p.cwd === match.cwd) ?? null
+    const knownBundleUrls = Array.from(
+      new Set([...(existing?.knownBundleUrls ?? []), match.bundleUrl, bundleUrl]),
+    )
+    const project = upsertProject({
+      cwd: match.cwd,
+      name: match.projectName ?? path.basename(match.cwd),
+      preferredProvider: input.provider ?? existing?.preferredProvider,
+      sourceRoots: existing?.sourceRoots ?? [match.cwd],
+      knownBundleUrls,
+      framework: existing?.framework ?? mapFrameworkToProjectFramework(match.framework),
+      bundleId: match.bundleId ?? existing?.bundleId,
+    })
+    return { project }
+  }
+
+  private getTranscript(sessionId: string): string | { error: string; code: string } {
+    const p = transcriptPath(sessionId)
+    if (!fs.existsSync(p)) {
+      return { error: 'transcript not found', code: 'NO_TRANSCRIPT' }
+    }
+    return fs.readFileSync(p, 'utf8')
+  }
+
+  private getPaths() {
+    const dir = getUserDataDir()
+    return {
+      userDataDir: dir,
+      storeFile: path.join(dir, 'attached-projects.json'),
+      sessionsDir: path.join(dir, 'sessions'),
+      transcriptsDir: path.join(dir, 'transcripts'),
+    }
+  }
+
+  // --- subscription management ---
+
+  private subscribeSocket(
+    ws: WebSocket,
+    sessionId: string,
+  ): { ok: true; refCount: number } {
+    let sockets = this.sessionsBySocket.get(ws)
+    if (!sockets) {
+      sockets = new Set()
+      this.sessionsBySocket.set(ws, sockets)
+    }
+    if (sockets.has(sessionId)) {
+      return { ok: true, refCount: this.subscriptions.get(sessionId)?.refCount ?? 1 }
+    }
+    sockets.add(sessionId)
+    const existing = this.subscriptions.get(sessionId)
+    if (existing) {
+      existing.refCount++
+      return { ok: true, refCount: existing.refCount }
+    }
+    const unsubscribe = subscribeAgentEvents(sessionId, (event) => {
+      const coalesced = this.coalescePromptEcho(sessionId, event)
+      if (coalesced) {
+        this.applySessionEvent(sessionId, coalesced)
+        this.fanOutEvent(sessionId, coalesced)
+      }
+      // recordTurnTelemetry mirrors what electron main used to do — keep it
+      // here so cost tracking happens regardless of which client subscribed.
+      if (event.type === 'turn-completed') {
+        const session = findSessionById(sessionId)
+        if (session) {
+          try {
+            recordTurnTelemetry(session.projectId, {
+              usd: event.costUsd,
+              ts: event.ts,
+            })
+          } catch (err) {
+            process.stderr.write(
+              `[sootsim-agent] recordTurnTelemetry failed: ${err instanceof Error ? err.message : String(err)}\n`,
+            )
+          }
+        }
+      }
+    })
+    this.subscriptions.set(sessionId, { unsubscribe, refCount: 1 })
+    return { ok: true, refCount: 1 }
+  }
+
+  private unsubscribeSocket(
+    ws: WebSocket,
+    sessionId: string,
+  ): { ok: true; refCount: number } {
+    const sockets = this.sessionsBySocket.get(ws)
+    if (!sockets || !sockets.has(sessionId)) return { ok: true, refCount: 0 }
+    sockets.delete(sessionId)
+    return this.decrementSubscription(sessionId)
+  }
+
+  private decrementSubscription(sessionId: string): { ok: true; refCount: number } {
+    const existing = this.subscriptions.get(sessionId)
+    if (!existing) return { ok: true, refCount: 0 }
+    existing.refCount--
+    if (existing.refCount <= 0) {
+      try {
+        existing.unsubscribe()
+      } catch {}
+      this.subscriptions.delete(sessionId)
+      return { ok: true, refCount: 0 }
+    }
+    return { ok: true, refCount: existing.refCount }
+  }
+
+  /** end-session tears the FIFO down, so drop our reader before it
+   *  disappears regardless of remaining subscriber refcount. */
+  private dropSessionFanout(sessionId: string): void {
+    const existing = this.subscriptions.get(sessionId)
+    if (existing) {
+      try {
+        existing.unsubscribe()
+      } catch {}
+      this.subscriptions.delete(sessionId)
+    }
+    for (const sockets of this.sessionsBySocket.values()) {
+      sockets.delete(sessionId)
+    }
+    this.clearPromptTracking(sessionId)
+  }
+
+  // --- wire pushes + responses ---
+
+  private normalizePromptEnvelope(msg: any): AgentPromptEnvelope {
+    if (msg?.prompt && typeof msg.prompt === 'object') {
+      const prompt = msg.prompt as Record<string, unknown>
+      return {
+        text: String(prompt.text ?? ''),
+        ...(typeof prompt.displayText === 'string'
+          ? { displayText: prompt.displayText }
+          : {}),
+        ...(typeof prompt.inspectSummary === 'string'
+          ? { inspectSummary: prompt.inspectSummary }
+          : {}),
+        ...(typeof prompt.inspectTrace === 'string'
+          ? { inspectTrace: prompt.inspectTrace }
+          : {}),
+      }
+    }
+    return {
+      text: String(msg?.text ?? ''),
+      ...(typeof msg?.displayText === 'string' ? { displayText: msg.displayText } : {}),
+      ...(typeof msg?.inspectSummary === 'string'
+        ? { inspectSummary: msg.inspectSummary }
+        : {}),
+      ...(typeof msg?.inspectTrace === 'string'
+        ? { inspectTrace: msg.inspectTrace }
+        : {}),
+    }
+  }
+
+  private notePromptAccepted(
+    sessionId: string,
+    prompt: AgentPromptEnvelope,
+    assumeQueued: boolean,
+  ): { ok: true; queued: boolean; pendingTurns: number; queueDepth: number } {
+    const now = Date.now()
+    const echoes = this.pendingPromptEchoes.get(sessionId) ?? []
+    echoes.push({ sentAt: now })
+    this.pendingPromptEchoes.set(sessionId, echoes)
+    const pendingTurns =
+      Math.max(this.pendingTurns.get(sessionId) ?? 0, assumeQueued ? 1 : 0) + 1
+    this.pendingTurns.set(sessionId, pendingTurns)
+    const promptText = prompt.displayText ?? prompt.text
+    this.patchSession(sessionId, {
+      lastPrompt: promptText,
+      status: 'working',
+      needsAttention: false,
+    })
+    this.fanOutEvent(sessionId, {
+      type: 'prompt-received',
+      text: promptText,
+      ...(prompt.inspectSummary ? { inspectSummary: prompt.inspectSummary } : {}),
+      ...(prompt.inspectTrace ? { inspectTrace: prompt.inspectTrace } : {}),
+      ts: now,
+    } as AgentEvent)
+    return {
+      ok: true,
+      queued: pendingTurns > 1,
+      pendingTurns,
+      queueDepth: Math.max(0, pendingTurns - 1),
+    }
+  }
+
+  private applySessionEvent(sessionId: string, event: AgentEvent): void {
+    switch (event.type) {
+      case 'prompt-received':
+      case 'turn-started':
+        this.patchSession(sessionId, {
+          status: 'working',
+          needsAttention: false,
+        })
+        return
+      case 'turn-completed': {
+        const pendingTurns = this.consumeSettledTurn(sessionId)
+        this.patchSession(sessionId, {
+          status: pendingTurns > 0 ? 'working' : 'idle',
+          needsAttention: false,
+          lastTurnFiles: event.filesTouched,
+          currentlyEditing: undefined,
+        })
+        return
+      }
+      case 'approval-needed':
+        this.patchSession(sessionId, {
+          status: 'needs-attention',
+          needsAttention: true,
+        })
+        return
+      case 'error': {
+        const pendingTurns = this.consumeSettledTurn(sessionId)
+        this.patchSession(sessionId, {
+          status: pendingTurns > 0 ? 'working' : 'needs-attention',
+          needsAttention: pendingTurns <= 0,
+          currentlyEditing: undefined,
+        })
+        return
+      }
+      case 'exited':
+        this.clearPromptTracking(sessionId)
+        this.patchSession(sessionId, {
+          status: 'ended',
+          needsAttention: false,
+          wrapperPid: undefined,
+          currentlyEditing: undefined,
+        })
+        return
+      case 'ready':
+      case 'turn-reasoning':
+      case 'turn-message':
+      case 'turn-plan':
+      case 'tool-call':
+      case 'file-edited':
+      case 'file-diff-delta':
+        return
+    }
+  }
+
+  private patchSession(sessionId: string, patch: Partial<AgentSession>): void {
+    const existing =
+      this.pendingSessionStates.get(sessionId) ?? findSessionById(sessionId)
+    if (!existing) return
+    const updated = applySessionStatusPatch(existing, patch)
+    this.pendingSessionStates.set(sessionId, updated)
+    this.pendingSessionPatches.set(sessionId, {
+      ...this.pendingSessionPatches.get(sessionId),
+      ...patch,
+    })
+    this.broadcastSessionStatus(updated)
+    if (this.sessionStatusPersistTimer) return
+    this.sessionStatusPersistTimer = setTimeout(() => {
+      this.sessionStatusPersistTimer = null
+      this.flushSessionStatuses()
+    }, SESSION_STATUS_PERSIST_DELAY_MS)
+  }
+
+  private flushSessionStatuses(): void {
+    if (this.pendingSessionPatches.size === 0) return
+    const pendingStates = this.pendingSessionStates
+    const pendingPatches = this.pendingSessionPatches
+    this.pendingSessionStates = new Map()
+    this.pendingSessionPatches = new Map()
+    try {
+      updateSessionStatuses([...pendingPatches].map(([id, patch]) => ({ id, patch })))
+    } catch (error) {
+      for (const [id, session] of pendingStates) {
+        if (!this.pendingSessionStates.has(id)) {
+          this.pendingSessionStates.set(id, session)
+        }
+      }
+      for (const [id, patch] of pendingPatches) {
+        this.pendingSessionPatches.set(id, {
+          ...patch,
+          ...this.pendingSessionPatches.get(id),
+        })
+      }
+      process.stderr.write(
+        `[sootsim-agent] session status persistence failed: ${error instanceof Error ? error.message : String(error)}\n`,
+      )
+    }
+  }
+
+  private coalescePromptEcho(sessionId: string, event: AgentEvent): AgentEvent | null {
+    if (event.type !== 'prompt-received') return event
+    const pending = this.pendingPromptEchoes.get(sessionId)
+    if (!pending || pending.length === 0) return event
+    while (pending.length > 0 && Date.now() - pending[0]!.sentAt > 15_000) {
+      pending.shift()
+    }
+    if (pending.length === 0) {
+      this.pendingPromptEchoes.delete(sessionId)
+      return event
+    }
+    pending.shift()
+    if (pending.length === 0) {
+      this.pendingPromptEchoes.delete(sessionId)
+    } else {
+      this.pendingPromptEchoes.set(sessionId, pending)
+    }
+    return null
+  }
+
+  private consumeSettledTurn(sessionId: string): number {
+    const pendingTurns = Math.max(0, (this.pendingTurns.get(sessionId) ?? 1) - 1)
+    if (pendingTurns > 0) {
+      this.pendingTurns.set(sessionId, pendingTurns)
+    } else {
+      this.pendingTurns.delete(sessionId)
+    }
+    return pendingTurns
+  }
+
+  private clearPromptTracking(sessionId: string): void {
+    this.pendingPromptEchoes.delete(sessionId)
+    this.pendingTurns.delete(sessionId)
+  }
+
+  private fanOutEvent(sessionId: string, event: AgentEvent): void {
+    const payload = JSON.stringify({ type: 'agent:event', sessionId, event })
+    for (const [ws, sessions] of this.sessionsBySocket) {
+      if (!sessions.has(sessionId)) continue
+      if (ws.readyState !== WS_OPEN) continue
+      try {
+        ws.send(payload)
+      } catch {}
+    }
+  }
+
+  private broadcastSessionStatus(session: AgentSession): void {
+    const payload = JSON.stringify({ type: 'agent:session-status', session })
+    for (const ws of this.allSockets) {
+      if (ws.readyState !== WS_OPEN) continue
+      try {
+        ws.send(payload)
+      } catch {}
+    }
+  }
+
+  private respond(ws: WebSocket, id: unknown, result: unknown): void {
+    if (ws.readyState !== WS_OPEN) return
+    try {
+      ws.send(JSON.stringify({ id, result }))
+    } catch {}
+  }
+
+  private respondError(ws: WebSocket, id: unknown, error: string, code?: string): void {
+    if (ws.readyState !== WS_OPEN) return
+    try {
+      ws.send(JSON.stringify({ id, error, ...(code ? { code } : {}) }))
+    } catch {}
+  }
+}
+
+function mapFrameworkToProjectFramework(
+  fw: DiscoveredServer['framework'],
+): 'expo' | 'one' | 'rock' | 'unknown' {
+  if (fw === 'expo') return 'expo'
+  if (fw === 'one' || fw === 'vxrn') return 'one'
+  return 'unknown'
+}