orca/src/main/pi/agent-status-handler-source.ts

import type { PiAgentKind } from '../../shared/pi-agent-kind'
import { getPiAgentStatusUiPromptHandlerSourceLines } from './agent-status-ui-prompt-source'

// Why: keep the generated handler registrations separate from hook transport;
// both are independently sizeable and the installed extension concatenates them.
export function getPiAgentStatusHandlerSourceLines(kind: PiAgentKind): string[] {
  const sessionStartHandler =
    kind !== 'omp'
      ? [
          "  pi.on('session_start', (event, ctx) => {",
          '    updateSessionMetadata(ctx)',
          ...(kind === 'pi' ? ['    piUiPromptDepth = 0'] : []),
          '    // Why: /reload re-registers the active session, but it is not a',
          '    // turn boundary and must not clear the visible status or unread state.',
          "    if (event.reason === 'reload') return",
          "    post('session_start')",
          '  })',
          ''
        ]
      : []

  // Why: OMP can switch sessions in-process, so each latest-only post needs fresh identity.
  const ctxParam = ', ctx'
  const bareCtxParams = '_event, ctx'
  const captureSessionMetadata = ['    updateRuntimeOmpSessionMetadata(ctx)']
  const primeDaemonWorkerGuard =
    kind === 'prime-agent'
      ? [
          '  // Why: Prime loads extensions in both its frontend and event-emitting daemon worker.',
          '  if (!process.env.PRIME_AGENT_INTERNAL_DAEMON_WORKER) return'
        ]
      : []
  const ownerEnv = kind === 'prime-agent' ? 'ORCA_PRIME_AGENT_STATUS_OWNED' : 'ORCA_PI_STATUS_OWNED'

  // Why: OMP suppresses its approval lifecycle unless an extension listens for it,
  // and it is the only signal that the run is parked on a permission prompt rather
  // than still working. Prime has no OMP runtime, so the handlers would be dead there.
  const approvalHandlers =
    kind === 'prime-agent'
      ? []
      : [
          `  pi.on('tool_approval_requested', (event${ctxParam}) => {`,
          ...captureSessionMetadata,
          '    if (!isOmpRuntime()) return',
          "    post('tool_approval_requested', {",
          '      tool_name: event.toolName,',
          '      reason: event.reason,',
          '      approval_mode: event.approvalMode,',
          '    })',
          '  })',
          '',
          `  pi.on('tool_approval_resolved', (event${ctxParam}) => {`,
          ...captureSessionMetadata,
          '    if (!isOmpRuntime()) return',
          "    post('tool_approval_resolved', {",
          '      tool_name: event.toolName,',
          '      approved: event.approved,',
          '    })',
          '  })',
          ''
        ]

  return [
    '// Why: pi assistant messages carry content as an array of parts',
    "// ({ type: 'text', text } / tool_use / tool_result / reasoning). We only",
    "// surface the concatenated text parts as the visible 'last assistant",
    "// message' for the dashboard preview — tool_use / reasoning would be",
    '// noise (the dashboard already shows the active tool name + input).',
    'function extractAssistantText(message: unknown): string {',
    "  if (!message || typeof message !== 'object') return ''",
    '  const content = (message as { content?: unknown }).content',
    "  if (typeof content === 'string') return content",
    "  if (!Array.isArray(content)) return ''",
    "  let out = ''",
    '  for (const part of content) {',
    "    if (part && typeof part === 'object' && (part as { type?: unknown }).type === 'text') {",
    '      const text = (part as { text?: unknown }).text',
    "      if (typeof text === 'string') out += text",
    '    }',
    '  }',
    '  return out',
    '}',
    '',
    "// Why: pi's tool_call event input shape is tool-specific (event.input is",
    '// the raw args object). The agent-hooks server already runs',
    '// deriveToolInputPreview(toolName, input) to render a friendly preview',
    "// for known tool names ('bash' → command, 'read'/'write'/'edit' → path,",
    '// etc.), so we forward the raw object verbatim under the same field',
    '// names Claude uses (tool_name / tool_input) and let the server pick the',
    '// preview. Keeps tool-name knowledge centralized on the receiver side.',
    "// Why: child agents inherit the lead's pane env; only its process may",
    '// register status hooks. PID identity keeps in-process reloads reporting.',
    'export default function (pi): void {',
    ...primeDaemonWorkerGuard,
    `  const ownerPid = process.env.${ownerEnv}`,
    '  const selfPid = String(process.pid)',
    '  if (ownerPid && ownerPid !== selfPid) return',
    `  process.env.${ownerEnv} = selfPid`,
    ...sessionStartHandler,
    `  pi.on('before_agent_start', (event${ctxParam}) => {`,
    ...captureSessionMetadata,
    "    post('before_agent_start', { prompt: event.prompt ?? '' })",
    '  })',
    '',
    `  pi.on('agent_start', (${bareCtxParams}) => {`,
    ...captureSessionMetadata,
    '    clearPendingAgentEndCheck()',
    '    agentEndReported = false',
    // Why: a turn cannot begin under a dialog holding input focus, so this is the one
    // boundary that can recover a modal whose close never arrived.
    ...(kind === 'pi' ? ['    piUiPromptDepth = 0', '    piTurnInFlight = true'] : []),
    "    post('agent_start')",
    '  })',
    '',
    `  pi.on('tool_execution_start', (event${ctxParam}) => {`,
    ...captureSessionMetadata,
    "    post('tool_execution_start', {",
    '      tool_name: event.toolName,',
    '      tool_input: event.args,',
    '    })',
    '  })',
    '',
    `  pi.on('tool_call', (event${ctxParam}) => {`,
    ...captureSessionMetadata,
    "    post('tool_call', {",
    '      tool_name: event.toolName,',
    '      tool_input: event.input,',
    '    })',
    '  })',
    '',
    `  pi.on('tool_execution_end', (event${ctxParam}) => {`,
    ...captureSessionMetadata,
    "    post('tool_execution_end', {",
    '      tool_name: event.toolName,',
    '    })',
    '  })',
    '',
    ...approvalHandlers,
    ...getPiAgentStatusUiPromptHandlerSourceLines(kind),
    "  // Why: capture the assistant's final text on each completed message",
    '  // so the dashboard preview reflects the most recent reply even before',
    '  // agent_end fires. message_end is the right hook because pi guarantees',
    '  // it fires after the message is finalized (post-streaming).',
    `  pi.on('message_end', (event${ctxParam}) => {`,
    ...captureSessionMetadata,
    "    if (event.message?.role !== 'assistant') return",
    '    const text = extractAssistantText(event.message)',
    '    if (!text) return',
    "    post('message_end', { role: 'assistant', text })",
    '  })',
    '',
    '  // Why: modern Pi stays non-idle across retry/compaction/follow-up work,',
    '  // while legacy Pi becomes idle after its final agent_end handlers.',
    '  // OMP instead marks non-terminal agent_end events with willContinue, so it',
    '  // returns before the recheck timer is ever armed.',
    '  const AGENT_END_IDLE_RECHECK_MS = 25',
    '  const AGENT_END_IDLE_RECHECK_MAX_MS = 250',
    '  let agentSettledSupported = false',
    '  let agentEndReported = false',
    '  let agentEndIdleRecheckMs = AGENT_END_IDLE_RECHECK_MS',
    '  let pendingAgentEndCheck: ReturnType<typeof setTimeout> | null = null',
    '  let pendingAgentEndContext: { isIdle: () => boolean } | null = null',
    '',
    '  function clearPendingAgentEndCheck(): void {',
    '    if (pendingAgentEndCheck !== null) clearTimeout(pendingAgentEndCheck)',
    '    pendingAgentEndCheck = null',
    '    pendingAgentEndContext = null',
    '  }',
    '',
    '  // Why: isIdle flips before agent_settled handlers run, so both paths',
    '  // share a per-run guard instead of racing duplicate completion posts.',
    '  function postAgentEndOnce(): void {',
    '    if (agentEndReported) return',
    '    agentEndReported = true',
    // Why: distinct from agentEndReported, which also dedupes the completion post and so
    // starts false on a pane that has not run a turn yet — that pane is idle, not busy.
    ...(kind === 'pi' ? ['    piTurnInFlight = false'] : []),
    "    post('agent_end')",
    '  }',
    '',
    '  function checkPendingAgentEnd(): void {',
    '    pendingAgentEndCheck = null',
    '    const ctx = pendingAgentEndContext',
    '    if (!ctx || agentSettledSupported || agentEndReported) {',
    '      pendingAgentEndContext = null',
    '      return',
    '    }',
    '    try {',
    '      if (ctx.isIdle()) {',
    '        pendingAgentEndContext = null',
    '        postAgentEndOnce()',
    '        return',
    '      }',
    '    } catch {',
    '      pendingAgentEndContext = null',
    '      return',
    '    }',
    '    pendingAgentEndCheck = setTimeout(checkPendingAgentEnd, agentEndIdleRecheckMs)',
    "    if (typeof pendingAgentEndCheck.unref === 'function') pendingAgentEndCheck.unref()",
    '    agentEndIdleRecheckMs = Math.min(agentEndIdleRecheckMs * 2, AGENT_END_IDLE_RECHECK_MAX_MS)',
    '  }',
    '',
    `  pi.on('agent_settled', (${bareCtxParams}) => {`,
    ...captureSessionMetadata,
    '    agentSettledSupported = true',
    '    clearPendingAgentEndCheck()',
    '    postAgentEndOnce()',
    '  })',
    '',
    "  pi.on('agent_end', (event, ctx) => {",
    ...captureSessionMetadata,
    '    if (event?.willContinue === true) {',
    '      clearPendingAgentEndCheck()',
    '      return',
    '    }',
    '    if (isOmpRuntime()) {',
    '      postAgentEndOnce()',
    '      return',
    '    }',
    '    if (agentSettledSupported) return',
    "    if (!ctx || typeof ctx.isIdle !== 'function') {",
    '      postAgentEndOnce()',
    '      return',
    '    }',
    '    clearPendingAgentEndCheck()',
    '    agentEndIdleRecheckMs = AGENT_END_IDLE_RECHECK_MS',
    '    pendingAgentEndContext = ctx',
    '    pendingAgentEndCheck = setTimeout(checkPendingAgentEnd, 0)',
    "    if (typeof pendingAgentEndCheck.unref === 'function') pendingAgentEndCheck.unref()",
    '  })',
    '}',
    ''
  ]
}