#!/bin/bash
# Run one pi subagent, headless, to completion.
#
#   run-agent.sh [options] <workdir> <provider> <model> <name>
#
# Options:
#   --brief FILE      task brief to hand the subagent (default: <workdir>/brief.md)
#   --timeout SECS    hard wall-clock limit (default: 1800)
#   --thinking LEVEL  off|minimal|low|medium|high|xhigh|max (default: high)
#   --tools LIST      pi tool allowlist, e.g. read,grep,find,ls for a subagent
#                     that must not run or change anything (default: pi's own)
#   --cwd DIR         directory the subagent runs in (default: the current one)
#   --context-files   let pi load AGENTS.md/CLAUDE.md from --cwd (default: off)
#   --deliver TEXT    what the final answer must be (default: "your complete
#                     result, as markdown")
#
# Sessions -- off by default; a run without --session cannot be continued:
#   --session         save the conversation as session <name>, under
#                     <workdir>/sessions/, so it can be followed up or forked
#   --followup TEXT   send TEXT as the next message in session <name>, instead
#                     of the brief. <provider> <model> may differ from the
#                     earlier run's: the new model reads the whole history
#   --fork FROM       with --followup: start session <name> as a copy of
#                     session FROM and send TEXT there. FROM is left as it was
#
# A follow-up or fork inherits --tools, --thinking, --cwd and --context-files
# from the run it continues, unless given again: a read-only subagent stays
# read-only. The session is opened by file path, never by id -- by id, pi
# looks only among sessions started in the current directory, and where it
# finds none it starts a new, empty one with that id, exit status 0, and only
# a warning on stderr to say so.
#
# Writes, under <workdir>/out/:
#   <name>.md         the subagent's final answer: text of its last message
#   <name>.cost.json  tokens, seconds, exit status, USD cost from the provider
#   <name>.jsonl      the event stream without per-token deltas, for debugging
#   <name>.err        pi's stderr
#   <name>.running    exists only while the run is in flight; holds the PID
#
# A previous run under the same name is never overwritten: its files are moved
# to <name>.attempt<N>.* first, so a retry keeps the bill of the failure -- or
# to <name>.round<N>.* when the new run is a follow-up of it. A retry also
# moves the old session aside, so it starts fresh instead of resuming it.
#
# Exit status is pi's, except that a run pi calls a success is turned into 1
# when its answer is empty or its last message ended in a provider error: pi
# exits 0 for both (an invalid API key, for one), and the cost file records
# the stop reason and error. 124 means the timeout fired -- the run is a
# write-off, but it still cost money; the cost file is written either way.
#
# Six mechanics here were each established the hard way. Do not "simplify"
# them away:
#
#  1. `--offline` is required. pi's startup model-catalog fetch intermittently
#     hangs forever; --offline skips it. Symptom without it: the process
#     produces no output at all, not even the session line, and never exits.
#
#  2. pi's stdout must be redirected to a FILE, never piped. A pipeline
#     deadlocks: `timeout` kills pi, but the downstream reader never sees EOF
#     and hangs past its own timeout. This is why filtering happens after the
#     run, not inline.
#
#  3. The raw stream is filtered and deleted at once. Older pi re-emitted the
#     entire message on every token delta, so it grew quadratically -- 790 MB
#     and 3.1 GB on real runs. pi 1.x documents message_update as delta-only,
#     but the stream is still mostly deltas nobody reads. Budget generous
#     scratch per concurrent subagent until that is confirmed on the version
#     in use.
#
#  4. Costs must be summed over every assistant message, not read off the last
#     one. Each message carries only its own usage.
#
#  5. pi's stdin must be /dev/null. In -p mode pi also reads stdin when it is
#     not a terminal, and waits for EOF. Started from a shell whose stdin never
#     closes -- an agent's tool shell, a backgrounded job -- it hangs forever
#     with no output at all, which looks exactly like the hang in (1).
#
#  6. Byte counts are stripped of whitespace before jq reads them. macOS `wc -c`
#     pads its output ("       6"), jq's tonumber rejects that, and the whole
#     cost file silently came out empty.
#
# `timeout` is GNU coreutils and macOS has none unless coreutils is installed,
# where it may be `gtimeout`. Without either, perl's alarm does the same job: it
# survives exec, and SIGALRM ends pi. That exit (142) is reported as 124.
#
# Extensions and skills are always off (-ne -ns): an extension can block on UI
# that a headless run has no way to answer, and the brief should say everything
# the subagent needs to know.

set -u

brief=""; limit=""; thinking=""; tools=""; cwd=""; context=""
deliver="your complete result, as markdown"
session=""; followup=""; fork_from=""; set_tools=""; set_context=""
while [[ $# -gt 0 ]]
do
  case "$1" in
    --brief) brief="$2"; shift 2 ;;
    --timeout) limit="$2"; shift 2 ;;
    --thinking) thinking="$2"; shift 2 ;;
    --tools) tools="$2"; set_tools=y; shift 2 ;;
    --cwd) cwd="$2"; shift 2 ;;
    --context-files) context=y; set_context=y; shift ;;
    --deliver) deliver="$2"; shift 2 ;;
    --session) session=y; shift ;;
    --followup) followup="$2"; shift 2 ;;
    --fork) fork_from="$2"; shift 2 ;;
    -h|--help) sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    --) shift; break ;;
    -*) echo "run-agent.sh: unknown option $1" >&2; exit 2 ;;
    *) break ;;
  esac
done

P="${1:?workdir}"; provider="${2:?provider}"; model="${3:?model}"; name="${4:?name}"
P="$(cd "$P" && pwd)" || exit 2
out="$P/out"
sessions="$P/sessions"
mkdir -p "$out"

die() { echo "run-agent.sh: $*" >&2; exit 2; }

# The file of session $1 in $sessions, if there is one. pi names it
# <timestamp>_<id>.jsonl; the id in its header line is what decides.
session_file() {
  local f
  for f in "$sessions"/*_"$1".jsonl
  do
    [[ -e "$f" ]] || continue
    if [[ "$(head -1 "$f" | jq -r '.id // empty' 2>/dev/null)" == "$1" ]]
    then
      echo "$f"; return 0
    fi
  done
  return 1
}

is_running() {
  local pid
  pid="$(cat "$out/$1.running" 2>/dev/null)" && [[ -n "$pid" ]] && kill -0 "$pid" 2>/dev/null
}

# Move run $name's files to $name.$1<N>.*; $1 is "attempt" or "round".
rotate() {
  local kind="$1" n=1 ext f
  [[ -e "$out/$name.cost.json" || -e "$out/$name.md" || -e "$out/$name.err" ]] || return 0
  while [[ -e "$out/$name.$kind$n.cost.json" || -e "$out/$name.$kind$n.err" ]]
  do
    n=$((n + 1))
  done
  for ext in md cost.json jsonl err
  do
    [[ -e "$out/$name.$ext" ]] && mv "$out/$name.$ext" "$out/$name.$kind$n.$ext"
  done
  if [[ -e "$out/$name.$kind$n.cost.json" ]]
  then
    f="$(mktemp)"
    jq --arg a "$name.$kind$n" '.name = $a' "$out/$name.$kind$n.cost.json" > "$f" \
      && mv "$f" "$out/$name.$kind$n.cost.json"
  fi
  # A fresh run must not resume the old conversation: pi would, given its id.
  if [[ "$kind" == attempt ]] && f="$(session_file "$name")"
  then
    mv "$f" "$f.attempt$n"
  fi
}

is_running "$name" && die "$name is still running"
[[ -n "$fork_from" && -z "$followup" ]] && die "--fork needs --followup"
if [[ -n "$session$followup" ]] && ! [[ "$name" =~ ^[A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?$ ]]
then
  die "$name cannot be a session id: letters, digits, . _ - only, starting and ending with a letter or digit"
fi

mode=fresh; src=""; src_name=""
if [[ -n "$followup" ]]
then
  if [[ -n "$fork_from" ]]
  then
    mode=fork; src_name="$fork_from"
    is_running "$fork_from" && die "$fork_from is still running; fork it once it is done"
    if session_file "$name" > /dev/null || [[ -e "$out/$name.cost.json" ]]
    then
      die "$name already exists; fork into a new name"
    fi
  else
    mode=followup; src_name="$name"
  fi
  src="$(session_file "$src_name")" \
    || die "no session $src_name in $sessions -- only a run started with --session can be continued"
  # Inherit what the earlier run was given, unless overridden.
  prev="$out/$src_name.cost.json"
  if [[ -e "$prev" ]]
  then
    [[ -n "$set_tools" ]] || tools="$(jq -r '.tools // empty' "$prev")"
    [[ -n "$set_context" ]] || context="$(jq -r 'if .contextFiles then "y" else empty end' "$prev")"
    [[ -n "$thinking" ]] || thinking="$(jq -r '.thinking // empty' "$prev")"
    [[ -n "$cwd" ]] || cwd="$(jq -r '.cwd // empty' "$prev")"
  fi
  [[ -n "$cwd" ]] || cwd="$(head -1 "$src" | jq -r '.cwd // empty')"
else
  brief="${brief:-$P/brief.md}"
  [[ "$brief" == /* ]] || brief="$PWD/$brief"
  [[ -s "$brief" ]] || die "no brief at $brief"
fi
thinking="${thinking:-high}"; limit="${limit:-1800}"
cwd="$(cd "${cwd:-.}" && pwd)" || die "no directory $cwd"

if [[ $mode == followup ]]; then rotate round; else rotate attempt; fi

args=(--provider "$provider" --model "$model"
      -p -ne -ns --offline
      --thinking "$thinking" --mode json)
case $mode in
  fresh)
    if [[ -n "$session" ]]
    then
      mkdir -p "$sessions"
      args+=(--session-dir "$sessions" --session-id "$name")
    else
      args+=(--no-session)
    fi
    prompt="Read the file $brief and do exactly what it says. Deliver $deliver as your final answer." ;;
  followup)
    args+=(--session "$src")
    prompt="$followup" ;;
  fork)
    args+=(--fork "$src" --session-dir "$sessions" --session-id "$name")
    prompt="$followup" ;;
esac
[[ -n "$context" ]] || args+=(-nc)
[[ -n "$tools" ]] && args+=(--tools "$tools")

if command -v timeout > /dev/null; then limiter=(timeout "$limit")
elif command -v gtimeout > /dev/null; then limiter=(gtimeout "$limit")
else limiter=(perl -e 'alarm shift; exec @ARGV or die "exec: $!\n"' "$limit")
fi

echo $$ > "$out/$name.running"
trap 'rm -f "$out/$name.running"' EXIT
start=$SECONDS

(
  cd "$cwd" || exit 2
  exec "${limiter[@]}" pi "${args[@]}" "$prompt" \
    < /dev/null > "$out/$name.raw" 2> "$out/$name.err"
)
rc=$?
[[ $rc -eq 142 && "${limiter[0]}" == perl ]] && rc=124
elapsed=$((SECONDS - start))

raw_bytes=$(wc -c < "$out/$name.raw" | tr -d '[:space:]')
grep -v '"type":"message_update"' "$out/$name.raw" > "$out/$name.jsonl"
rm -f "$out/$name.raw"

# The deliverable: text blocks of the last assistant message.
jq -rs '
  [.[] | select(.type=="agent_end")] | last | .messages
  | [.[] | select(.role=="assistant")] | last
  | [.content[] | select(.type=="text") | .text] | join("")
' "$out/$name.jsonl" > "$out/$name.md" 2>/dev/null

# pi exits 0 when the provider refused the request, or the model said nothing.
stop=$(jq -rs '[.[] | select(.type=="message_end" and .message.role=="assistant")]
  | last | .message.stopReason // "none"' "$out/$name.jsonl" 2>/dev/null)
if [[ $rc -eq 0 ]] && { [[ "$stop" == error || "$stop" == none ]] || ! grep -q '[^[:space:]]' "$out/$name.md"; }
then
  rc=1
fi

session_path=""
if [[ $mode == followup ]]; then session_path="$src"
elif [[ $mode == fork || -n "$session" ]]; then session_path="$(session_file "$name")"
fi

jq -rs --arg name "$name" --arg model "$model" --arg provider "$provider" \
  --arg rc "$rc" --arg elapsed "$elapsed" --arg raw "$raw_bytes" \
  --arg mode "$mode" --arg from "$src_name" --arg session "$session_path" \
  --arg tools "$tools" --arg thinking "$thinking" --arg cwd "$cwd" \
  --arg context "$context" '
  [.[] | select(.type=="message_end" and .message.role=="assistant") | .message] as $msgs
  | ($msgs | last) as $last
  | ($msgs | map(.usage))
  | {name: $name, provider: $provider, model: $model, exit: ($rc|tonumber),
     stopReason: ($last.stopReason // null),
     error: ($last.errorMessage // null | if . then .[0:500] else . end),
     seconds: ($elapsed|tonumber), turns: length,
     input: (map(.input) | add // 0), output: (map(.output) | add // 0),
     cacheRead: (map(.cacheRead) | add // 0),
     cacheWrite: (map(.cacheWrite) | add // 0),
     rawStreamBytes: ($raw|tonumber),
     cost: (map(.cost.total) | add // 0),
     mode: $mode, from: (if $from == "" then null else $from end),
     session: (if $session == "" then null else $session end),
     tools: (if $tools == "" then null else $tools end),
     thinking: $thinking, cwd: $cwd, contextFiles: ($context != "")}
' "$out/$name.jsonl" > "$out/$name.cost.json" 2>/dev/null

printf '%-20s exit:%-4s stop:%-8s %5ss  answer:%7s bytes  cost:$%s\n' \
  "$name" "$rc" "$stop" "$elapsed" "$(wc -c < "$out/$name.md" | tr -d '[:space:]')" \
  "$(jq -r '.cost' "$out/$name.cost.json" 2>/dev/null)"
if [[ $rc -ne 0 ]]
then
  jq -r '.error // empty' "$out/$name.cost.json" 2>/dev/null | head -c 300; echo
fi
exit $rc