orca/src/relay/relay-agent-hook-runtime.ts

import { AGENT_HOOK_INFER_INTERRUPT_METHOD } from '../shared/agent-hook-interrupt-reconciliation'
import { homedir } from 'node:os'
import type { RelayDispatcher } from './dispatcher'
import type { PtyEnvAugmenter, PtyHandler } from './pty-handler'
import { RelayAgentHookServer } from './agent-hook-server'
import { endpointDirForRelaySocket } from './agent-hook-endpoint-coordinates'
import { PluginOverlayManager } from './plugin-overlay'
import { installOpenCodePluginInCanonicalConfig } from './opencode-canonical-config'
import {
  AGENT_HOOK_INSTALL_PLUGINS_METHOD,
  AGENT_HOOK_REQUEST_REPLAY_METHOD
} from '../shared/agent-hook-relay'
import { publishAgentHookEnvelope } from './agent-hook-envelope-publication'
import { assertPluginSourceUnderByteCap } from './plugin-source-limit'
import {
  resolveOpenCodeSourceConfigDir,
  resolvePiSourceAgentDir,
  resolveOmpConfigDirName
} from './plugin-overlay-env'
import {
  detectExplicitPiAgentKindFromCommand,
  isPiCompatibleAgentType
} from '../shared/pi-agent-kind'
import { resolveSetupAgentSequenceLaunchCommand } from '../shared/setup-agent-sequencing'
import { selectOpenCodeHookAgent } from '../shared/opencode-launch-command'
import { relayLogLine } from './relay-diagnostic-log'
import { restoreOrStripOverlayEnv } from '../shared/agent-overlay-env'
import { registerManagedHookInstaller } from './managed-hook-installer'
import { readSessionShellStartupEnvVar } from '../main/pty/shell-startup-env'

export class RelayAgentHookRuntime {
  private readonly hookServer: RelayAgentHookServer
  private readonly pluginOverlay = new PluginOverlayManager()

  constructor(
    private readonly dispatcher: RelayDispatcher,
    private readonly ptyHandler: PtyHandler,
    sockPath: string,
    endpointDir?: string
  ) {
    this.hookServer = new RelayAgentHookServer({
      endpointDir: endpointDir ?? endpointDirForRelaySocket(sockPath),
      forward: (envelope) => publishAgentHookEnvelope(dispatcher, envelope),
      forwardUnavailable: (envelope) => publishAgentHookEnvelope(dispatcher, envelope),
      // Why: the PTY handler is the only component that knows which panes still have a client
      // surface, so it — not the client — decides whether a hook post describes a live pane.
      isPaneSurfaceRetired: (paneKey) => ptyHandler.isPaneSurfaceRetired(paneKey),
      getAgentLaunchToken: (paneKey) => ptyHandler.getAgentLaunchToken(paneKey),
      getTmuxManagedPty: async (paneKey) => ptyHandler.getTmuxManagedPty(paneKey)
    })
  }

  async start(): Promise<void> {
    try {
      await this.hookServer.start({ publishEndpoint: false })
    } catch (error) {
      relayLogLine(
        `[relay] agent-hook server failed to start: ${error instanceof Error ? error.message : String(error)}`
      )
    }
    this.registerPtyEnvironment()
    this.registerHandlers()
  }

  publishEndpointFile(): void {
    this.hookServer.publishEndpointFile()
  }

  stop(): void {
    this.hookServer.stop()
  }

  private registerPtyEnvironment(): void {
    this.ptyHandler.setClaudeTerminalEvidenceListener((paneKey, evidence) => {
      this.hookServer.claudeTerminalInterrupts.observe(paneKey, evidence)
    })
    this.ptyHandler.setAgentPresenceTrigger((paneKey) => {
      void this.hookServer.checkAgentPresence(paneKey)
    })
    this.ptyHandler.addEnvAugmenter(() => this.hookServer.buildPtyEnv())
    this.ptyHandler.addEnvAugmenter((context) => this.buildPluginEnvironment(context))
    this.ptyHandler.setExitListener(({ paneKey, id }) => {
      if (paneKey) {
        this.hookServer.clearPaneState(paneKey)
      }
      this.pluginOverlay.clearOverlay(paneKey ?? id)
    })
    // Why: the exit listener above only fires on proof of process death, which a shell that
    // survives teardown never produces. Drop the pane's cached status the moment its tab goes, so a
    // reconnecting client cannot be handed a replay of an agent nobody owns.
    this.ptyHandler.setSurfaceRetiredListener(({ paneKey }) => {
      this.hookServer.clearPaneState(paneKey)
    })
  }

  private async buildPluginEnvironment(
    context: Parameters<PtyEnvAugmenter>[0]
  ): Promise<Record<string, string>> {
    const env: Record<string, string> = {}
    const overlayId = context.paneKey ?? context.id
    const launchCommandHint = resolveSetupAgentSequenceLaunchCommand(context.env, context.command)
    const opencodeAgent = selectOpenCodeHookAgent(context.launchAgent, launchCommandHint, (agent) =>
      this.pluginOverlay.hasOpenCodeSource(agent)
    )
    restoreOrStripOverlayEnv(
      context.env,
      {
        primary: 'OPENCODE_CONFIG_DIR',
        overlay: 'ORCA_OPENCODE_CONFIG_DIR',
        source: 'ORCA_OPENCODE_SOURCE_CONFIG_DIR',
        preserveExplicitPrimary: true
      },
      {}
    )
    delete context.env.ORCA_OPENCODE_AGENT
    if (opencodeAgent) {
      env.ORCA_OPENCODE_AGENT = opencodeAgent
      const sourceDir = resolveOpenCodeSourceConfigDir(context.env, context.shell)
      const inheritedRelayOverlay = sourceDir
        ? this.pluginOverlay.isRelayOverlayPath(sourceDir)
        : false
      if (sourceDir && !inheritedRelayOverlay) {
        const dir = this.pluginOverlay.materializeOpenCode(overlayId, sourceDir, opencodeAgent)
        if (dir) {
          env.OPENCODE_CONFIG_DIR = dir
          env.ORCA_OPENCODE_CONFIG_DIR = dir
          env.ORCA_OPENCODE_SOURCE_CONFIG_DIR = sourceDir
        }
      } else {
        this.pluginOverlay.installOpenCodePlugin(opencodeAgent, {
          ...context.env,
          XDG_CONFIG_HOME:
            readSessionShellStartupEnvVar('XDG_CONFIG_HOME', context.env, context.shell) ??
            context.env.XDG_CONFIG_HOME
        })
      }
    }
    const explicitKind = isPiCompatibleAgentType(context.launchAgent)
      ? context.launchAgent
      : context.launchAgent === undefined
        ? detectExplicitPiAgentKindFromCommand(launchCommandHint)
        : null
    const kind = explicitKind ?? 'pi'
    const hasLaunchCommand =
      typeof launchCommandHint === 'string' && launchCommandHint.trim().length > 0
    if (kind === 'omp' || !hasLaunchCommand) {
      env.ORCA_OMP_FRESH_CONFIG = this.pluginOverlay.materializeOmpFreshConfig()
    }
    if (!this.pluginOverlay.hasPiSource()) {
      return env
    }
    if (kind === 'pi') {
      const sourceDir = resolvePiSourceAgentDir(context.env, context.shell, 'pi')
      const result = this.pluginOverlay.materializePi(overlayId, sourceDir, 'pi', {
        materializeDefaultHome: explicitKind === 'pi'
      })
      if (result?.sourceAgentDir) {
        env.ORCA_PI_SOURCE_AGENT_DIR = result.sourceAgentDir
      }
    }
    if (kind === 'omp' || !hasLaunchCommand) {
      const sourceDir =
        kind === 'omp'
          ? resolvePiSourceAgentDir(context.env, context.shell, 'omp')
          : context.env.ORCA_OMP_SOURCE_AGENT_DIR
      const configDirName = await resolveOmpConfigDirName(context.env, context.shell)
      if (configDirName !== undefined) {
        env.PI_CONFIG_DIR = configDirName
      }
      const result = this.pluginOverlay.materializePi(overlayId, sourceDir, 'omp', {
        materializeDefaultHome: explicitKind === 'omp',
        configDirName
      })
      if (result?.statusExtensionPath) {
        env.ORCA_OMP_STATUS_EXTENSION = result.statusExtensionPath
      }
      if (result?.sourceAgentDir) {
        env.ORCA_OMP_SOURCE_AGENT_DIR = result.sourceAgentDir
      }
    }
    if (kind === 'prime-agent') {
      const sourceDir = resolvePiSourceAgentDir(context.env, context.shell, 'prime-agent')
      const result = this.pluginOverlay.materializePi(overlayId, sourceDir, 'prime-agent', {
        materializeDefaultHome: explicitKind === 'prime-agent'
      })
      if (result?.sourceAgentDir) {
        env.ORCA_PRIME_AGENT_SOURCE_AGENT_DIR = result.sourceAgentDir
      }
    }
    return env
  }

  private registerHandlers(): void {
    this.dispatcher.onRequest(AGENT_HOOK_INFER_INTERRUPT_METHOD, async (params) => ({
      applied: this.hookServer.inferInterrupt(params)
    }))
    this.dispatcher.onRequest(AGENT_HOOK_REQUEST_REPLAY_METHOD, async () => ({
      replayed: this.hookServer.replayCachedPayloadsForPanes()
    }))
    registerManagedHookInstaller(this.dispatcher)
    this.dispatcher.onRequest(AGENT_HOOK_INSTALL_PLUGINS_METHOD, async (params) => {
      const startupPrompt = params.opencodeStartupPromptSource
      const opencode = params.opencodePluginSource
      const opencode2 = params.opencode2PluginSource
      const pi = params.piExtensionSource
      const omp = params.ompExtensionSource
      const primeAgent = params.primeAgentExtensionSource
      assertPluginSourceUnderByteCap('opencodeStartupPromptSource', startupPrompt)
      assertPluginSourceUnderByteCap('opencodePluginSource', opencode)
      assertPluginSourceUnderByteCap('opencode2PluginSource', opencode2)
      assertPluginSourceUnderByteCap('piExtensionSource', pi)
      assertPluginSourceUnderByteCap('ompExtensionSource', omp)
      assertPluginSourceUnderByteCap('primeAgentExtensionSource', primeAgent)
      this.pluginOverlay.setSources({
        opencodeStartupPromptSource: typeof startupPrompt === 'string' ? startupPrompt : undefined,
        opencodePluginSource: typeof opencode === 'string' ? opencode : undefined,
        opencode2PluginSource: typeof opencode2 === 'string' ? opencode2 : undefined,
        piExtensionSource: typeof pi === 'string' ? pi : undefined,
        ompExtensionSource: typeof omp === 'string' ? omp : undefined,
        primeAgentExtensionSource: typeof primeAgent === 'string' ? primeAgent : undefined
      })
      // Why: a running OpenCode 2 service reloads a changed plugin file, so an Orca upgrade
      // reaches it on connect instead of at the next pane spawn. Never creates an install.
      for (const [agent, source] of [
        ['opencode', opencode],
        ['opencode2', opencode2]
      ] as const) {
        if (typeof source === 'string' && source) {
          installOpenCodePluginInCanonicalConfig(source, agent, process.env, homedir(), true)
        }
      }
      const startupPromptInstalled = this.pluginOverlay.installOpenCodeStartupPromptPlugin(
        process.env
      )
      return {
        installed: {
          opencodeStartupPrompt: startupPromptInstalled,
          opencode: this.pluginOverlay.hasOpenCodeSource(),
          opencode2: this.pluginOverlay.hasOpenCode2Source(),
          pi: this.pluginOverlay.hasPiSource('pi'),
          omp: this.pluginOverlay.hasPiSource('omp'),
          primeAgent: this.pluginOverlay.hasPiSource('prime-agent')
        }
      }
    })
  }
}