orca/src/shared/setup-agent-sequencing.ts

import { encodePowerShellCommand } from './powershell-command-encoding'
import {
  nativeWindowsPathToPosixShellPath,
  resolveSetupRunnerCommand,
  type SetupRunnerCommandPlatform,
  type SetupRunnerCommandShell,
  type SetupRunnerShell
} from './setup-runner-command'
import { createNonSecureContextUuid } from './non-secure-context-uuid'

const DEFAULT_WAIT_TIMEOUT_SECONDS = 2 * 60 * 60
// Exported so the gate and its tests share one definition.
export const SETUP_COMPLETE_MESSAGE = 'Setup finished; starting agent.'
export const SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV = 'ORCA_SEQUENCED_STARTUP_COMMAND'
export const SETUP_AGENT_SEQUENCE_STARTUP_SCRIPT_ENV = 'ORCA_SEQUENCED_STARTUP_SCRIPT'

export type SequencedSetupAgentCommands = {
  setupCommand: string
  startupCommand: string
  startupEnv?: Record<string, string>
}

export function resolveSetupAgentSequenceLaunchCommand(
  env: Record<string, string | undefined>,
  fallbackCommand: string | undefined
): string | undefined {
  const sequencedStartup = env[SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV]?.trim()
  return sequencedStartup || fallbackCommand
}

export function createSetupAgentSequenceNonce(): string {
  return createNonSecureContextUuid()
}

export function createSequencedSetupAgentCommands(args: {
  runnerScriptPath: string
  startupCommand: string
  platform: SetupRunnerCommandPlatform
  shell?: SetupRunnerShell
  nonce?: string
  waitTimeoutSeconds?: number
}): SequencedSetupAgentCommands {
  const nonce = args.nonce ?? createSetupAgentSequenceNonce()
  const resolution = resolveSetupRunnerCommand(args.runnerScriptPath, args.platform, args.shell)
  // Why: the gate is typed into the terminal pane and `startupCommand` is already quoted for that
  // pane, so a batch runner launched from a Git Bash pane still needs the bash gate — PowerShell's
  // `Invoke-Expression` cannot parse the POSIX `'\''` escaping the pane's quoting produces. The
  // runner itself still launches through `resolution.command`, never through bash.
  const posixGateForWindowsRunner = resolution.shell === 'windows' && args.shell?.family === 'posix'
  const markerBasePath = posixGateForWindowsRunner
    ? nativeWindowsPathToPosixShellPath(resolution.runnerScriptPathForShell)
    : resolution.runnerScriptPathForShell
  // Why: overlapping gated launches of the same setup runner must not race on
  // a shared completion marker.
  const markerPath = `${markerBasePath}.${nonce}.done`
  const waitTimeoutSeconds = args.waitTimeoutSeconds ?? DEFAULT_WAIT_TIMEOUT_SECONDS

  if (resolution.shell === 'windows' && !posixGateForWindowsRunner) {
    return {
      setupCommand: buildWindowsSetupCommand(
        resolution.runnerScriptPathForShell,
        markerPath,
        nonce
      ),
      startupCommand: buildWindowsStartupCommand(markerPath, nonce, waitTimeoutSeconds),
      startupEnv: {
        [SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV]: args.startupCommand
      }
    }
  }

  const startupScript = buildPosixStartupScript(
    args.startupCommand,
    markerPath,
    nonce,
    waitTimeoutSeconds
  )
  return {
    setupCommand: buildPosixSetupCommand(resolution.command, markerPath, nonce),
    // Why: long worktree paths can push the gate past a PTY's canonical input cap and drop its submit byte.
    startupCommand: `bash -lc 'eval "$${SETUP_AGENT_SEQUENCE_STARTUP_SCRIPT_ENV}"'`,
    startupEnv: {
      [SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV]: args.startupCommand,
      [SETUP_AGENT_SEQUENCE_STARTUP_SCRIPT_ENV]: startupScript
    }
  }
}

function buildPosixSetupCommand(setupCommand: string, markerPath: string, nonce: string): string {
  const marker = quotePosixArg(markerPath)
  const tmp = quotePosixArg(`${markerPath}.tmp`)
  const nonceValue = quotePosixArg(nonce)

  const script = [
    `rm -f ${marker} ${tmp} 2>/dev/null`,
    `( ${setupCommand} )`,
    'status=$?',
    `printf '%s:%s\\n' ${nonceValue} "$status" > ${tmp}`,
    `mv -f ${tmp} ${marker}`,
    'exit "$status"'
  ].join('; ')

  return `bash -lc ${quotePosixArg(script)}`
}

function buildPosixStartupScript(
  startupCommand: string,
  markerPath: string,
  nonce: string,
  waitTimeoutSeconds: number
): string {
  const marker = quotePosixArg(markerPath)
  const tmp = quotePosixArg(`${markerPath}.tmp`)
  const nonceValue = quotePosixArg(nonce)
  const timeout = Math.max(1, Math.floor(waitTimeoutSeconds))
  const startupSuccessCommand = buildPosixStartupSuccessCommand(startupCommand)
  // Why: the PTY launch path feeds this command through an interactive shell,
  // so keeping the wrapper on one line avoids visible `quote>` continuation
  // prompts while still preserving valid `while`/`if` shell syntax.
  const script = [
    `deadline=$((SECONDS + ${timeout}));`,
    'echo "Waiting for setup to finish before starting agent..." >&2;',
    'while :; do',
    `if [ -f ${marker} ]; then`,
    `IFS=: read -r seen status < ${marker} || true;`,
    `if [ "$seen" = ${nonceValue} ]; then`,
    `rm -f ${marker} ${tmp} 2>/dev/null;`,
    // Why: failure and timeout announce themselves; a silent success left
    // "Waiting for setup..." as the pane's last line forever.
    `if [ "$status" = "0" ]; then echo ${quotePosixArg(SETUP_COMPLETE_MESSAGE)} >&2; if [ -n "\${${SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV}:-}" ]; then eval "\$${SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV}"; exit "$?"; else ${startupSuccessCommand}; fi; fi;`,
    'echo "Setup failed; skipping agent startup." >&2;',
    'exit "${status:-1}";',
    'fi;',
    'fi;',
    'if [ "$SECONDS" -ge "$deadline" ]; then',
    'echo "Timed out waiting for setup before starting agent." >&2;',
    'exit 124;',
    'fi;',
    'sleep 1;',
    'done'
  ].join(' ')

  return script
}

function buildPosixStartupSuccessCommand(startupCommand: string): string {
  if (
    hasUnquotedPosixCommandSeparator(startupCommand) ||
    hasLeadingPosixEnvAssignment(startupCommand)
  ) {
    return `eval ${quotePosixArg(startupCommand)}; exit "$?"`
  }
  return `exec ${startupCommand}`
}

function hasLeadingPosixEnvAssignment(command: string): boolean {
  return /^[A-Za-z_][A-Za-z0-9_]*=/.test(command.trimStart())
}

function hasUnquotedPosixCommandSeparator(command: string): boolean {
  let quote: "'" | '"' | null = null
  let escaped = false
  for (const char of command) {
    if (escaped) {
      escaped = false
      continue
    }
    if (char === '\\') {
      escaped = true
      continue
    }
    if (quote) {
      if (char === quote) {
        quote = null
      }
      continue
    }
    if (char === "'" || char === '"') {
      quote = char
      continue
    }
    if (char === ';' || char === '&' || char === '|' || char === '\n' || char === '\r') {
      return true
    }
  }
  return false
}

function buildWindowsSetupCommand(
  runnerScriptPath: string,
  markerPath: string,
  nonce: string
): string {
  // Why: delayed expansion keeps path metacharacters as data when cmd invokes the batch runner.
  const script = [
    `$runner = ${quotePowerShellString(runnerScriptPath)}`,
    `$marker = ${quotePowerShellString(markerPath)}`,
    '$tmp = $marker + ".tmp"',
    `$nonce = ${quotePowerShellString(nonce)}`,
    'Remove-Item -LiteralPath $marker, $tmp -Force -ErrorAction SilentlyContinue',
    '$processInfo = [System.Diagnostics.ProcessStartInfo]::new()',
    '$processInfo.FileName = $env:ComSpec',
    '$processInfo.Arguments = \'/d /s /v:on /c ""!ORCA_SETUP_RUNNER!""\'',
    '$processInfo.UseShellExecute = $false',
    '$processInfo.EnvironmentVariables["ORCA_SETUP_RUNNER"] = $runner',
    '$process = [System.Diagnostics.Process]::Start($processInfo)',
    '$process.WaitForExit()',
    '$setupStatus = $process.ExitCode',
    '$utf8 = [System.Text.UTF8Encoding]::new($false)',
    '[System.IO.File]::WriteAllText($tmp, ($nonce + ":" + $setupStatus + [Environment]::NewLine), $utf8)',
    'Move-Item -LiteralPath $tmp -Destination $marker -Force',
    'exit $setupStatus'
  ].join('; ')

  return encodePowerShellInvocation(script)
}

function buildWindowsStartupCommand(
  markerPath: string,
  nonce: string,
  waitTimeoutSeconds: number
): string {
  const timeout = Math.max(1, Math.floor(waitTimeoutSeconds))
  // Why: native Windows setup runners launch through cmd.exe, but PowerShell
  // gives us safe bounded file polling/parsing without a fragile batch label loop.
  const script = [
    // Why: the startup command is user-authored and may invoke a `.ps1`, which IS
    // execution-policy gated even though `-EncodedCommand` is not. This is the in-payload
    // stand-in for the `-ExecutionPolicy Bypass` switch dropped from the command line
    // (same trade as the agent-hooks launcher). Progress must be silenced first and
    // restored after: Set-ExecutionPolicy autoloads a module whose "Preparing modules for
    // first use." record would otherwise land on the stderr this gate writes to.
    //
    // The failure is reported rather than swallowed. Autoload can fail for reasons that
    // have nothing to do with policy -- a 5.1 install with duplicate extended type data
    // fails every cmdlet in Microsoft.PowerShell.Security -- and the old
    // `-ErrorAction SilentlyContinue` plus empty `catch` hid that completely, leaving the
    // user with an execution-policy refusal from their own script and no trace that the
    // relief had been attempted. `-ErrorAction Stop` is what routes a non-terminating
    // failure into the catch at all. Still never throws: a diagnostic is worth a line of
    // stderr, but not the startup this gate exists to run.
    "$orcaProgress = $ProgressPreference; $ProgressPreference = 'SilentlyContinue'",
    'try { Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass -Force -ErrorAction Stop } ' +
      'catch { [Console]::Error.WriteLine("Orca: could not relax the execution policy for this " + ' +
      '"session (" + $_.FullyQualifiedErrorId + "). A startup command that runs a .ps1 " + ' +
      '"may be blocked.") }',
    '$ProgressPreference = $orcaProgress',
    `$marker = ${quotePowerShellString(markerPath)}`,
    'if ([string]::IsNullOrWhiteSpace($marker)) {',
    '  [Console]::Error.WriteLine("Missing setup marker path.")',
    '  exit 1',
    '}',
    '$tmp = $marker + ".tmp"',
    `$nonce = ${quotePowerShellString(nonce)}`,
    `$deadline = (Get-Date).AddSeconds(${timeout})`,
    '[Console]::Error.WriteLine("Waiting for setup to finish before starting agent...")',
    'while ($true) {',
    '  if (Test-Path -LiteralPath $marker) {',
    '    $content = Get-Content -LiteralPath $marker -TotalCount 1',
    '    if ($content -match "^([0-9A-Za-z_-]+):([0-9]+)$" -and $Matches[1] -eq $nonce) {',
    '      $setupStatus = [int]$Matches[2]',
    '      Remove-Item -LiteralPath $marker, $tmp -Force -ErrorAction SilentlyContinue',
    '      if ($setupStatus -ne 0) {',
    '        [Console]::Error.WriteLine("Setup failed; skipping agent startup.")',
    '        exit $setupStatus',
    '      }',
    `      $startup = $env:${SETUP_AGENT_SEQUENCE_STARTUP_COMMAND_ENV}`,
    '      if ([string]::IsNullOrWhiteSpace($startup)) {',
    '        [Console]::Error.WriteLine("Missing sequenced startup command.")',
    '        exit 1',
    '      }',
    `      [Console]::Error.WriteLine(${quotePowerShellString(SETUP_COMPLETE_MESSAGE)})`,
    '      Invoke-Expression $startup',
    '      if ($global:LASTEXITCODE -ne $null) { exit $global:LASTEXITCODE }',
    '      if (-not $?) { exit 1 }',
    '      exit 0',
    '    }',
    '  }',
    '  if ((Get-Date) -ge $deadline) {',
    '    [Console]::Error.WriteLine("Timed out waiting for setup before starting agent.")',
    '    exit 124',
    '  }',
    '  Start-Sleep -Seconds 1',
    '}'
  ].join('; ')

  return encodePowerShellInvocation(script)
}

// Why: `-EncodedCommand` is not execution-policy gated (only `-File` is), so `-ExecutionPolicy
// Bypass` was a no-op — and it is one of the most heavily EDR-flagged PowerShell tokens. The
// base64 stays: these strings are typed into a terminal pane and re-parsed by its shell.
function encodePowerShellInvocation(script: string): string {
  return `powershell.exe -NoProfile -NonInteractive -EncodedCommand ${encodePowerShellCommand(script)}`
}

function quotePosixArg(value: string): string {
  if (/^[A-Za-z0-9_./:-]+$/.test(value)) {
    return value
  }
  return `'${value.replace(/'/g, `'\\''`)}'`
}

function quotePowerShellString(value: string): string {
  return `'${value.replace(/'/g, "''")}'`
}

export function getSetupAgentSequenceShellForTests(
  runnerScriptPath: string,
  platform: SetupRunnerCommandPlatform
): SetupRunnerCommandShell {
  return resolveSetupRunnerCommand(runnerScriptPath, platform).shell
}