summaryrefslogtreecommitdiff
path: root/skills/pi-subagent/run-agent.sh
blob: aaf67707761273c247a2c67f31600f1dd50f6e15 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
#!/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