joakimp/pi-toolkit
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(pi-task): optional model fallback when the primary fails auth before acting

Browse Source 
A tier's model can be unusable for reasons that have nothing to do with the
task, typically an expired AWS SSO session behind an amazon-bedrock profile.
pi reports that as a normal run: rc 0, agent_end, willRetry=false. The failure
shows only as stopReason="error" + errorMessage on the assistant message
(measured 2026-10-12). pi-task turned it into "no fenced json envelope",
which points at the wrong layer.

- Always surface it, first in the problem list:
  "child model error (<provider>/<id>): ...".
- New, optional: pi-task.fallbackProfiles.<effort> = {provider, id[, thinking]}
  in ~/.pi/agent/settings.json. It is tried ONCE, and only when the primary
  ended in an auth-class error, made zero tool calls, left the boundary
  diff empty, and at least 30 s of wall_s remain. A write task that failed
  after acting is never re-run on a half-edited tree.
- A fallback without `thinking` inherits the primary's rather than "off".
  A fallback identical to the primary is ignored, with a warning.
- The key lives in pi-task's own namespace because the fork tool does not
  fall back. pi-fork's parser silently drops unknown profile keys.
- Reporting: "[FALLBACK]" on the verdict line plus a warning line.
  result.json gains fallback{} and attempts[]; metrics.model/provider name
  the model that produced the verdict. Audit keeps raw.primary.ndjson,
  argv.fallback.json and boundary-after.primary.json. The second attempt
  gets only the remaining wall budget, so the task tool's (wall_s+90) s kill
  still holds.
- Unchanged: exit 0 = PASS / 1 = FAIL, and the closing "audit:" line.

selftest: 14 -> 29 checks, two-sided for every gate (fixtures mirror the
measured event shape; the SSO string is verbatim from a real expired
session). Each gate was mutation-checked: removing it turns its line BAD.

End-to-end, emb-7kj4vr4g, with fake AWS keys in the environment:
- the old pi-task gives FAIL "no fenced json envelope";
- the new one, per tier, gives PASS [FALLBACK] with the correct answer and
  real tool calls;
- with real credentials: PASS on the primary, fallback configured but unused;
- with no fallback configured: FAIL, model error listed first.
This commit is contained in:
Joakim Persson
2026-10-12 11:07:32 +02:00
parent 9c87ee843f
commit 3760d90998
3 changed files with 313 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bin/pi-task
+259 -42
View File
@@ -15,9 +15,16 @@ Honest limits, stated so nobody mistakes this for a sandbox:
"read_only" is therefore VERIFIED POST-HOC by a boundary diff, not enforced.
* cost is checked after the fact; pi has no spend ceiling to pass in.
Model fallback (optional): settings.json `pi-task.fallbackProfiles.<effort>` =
{provider, id[, thinking]}. Tried once, and ONLY when the primary failed with an
auth-class error before running any tool or changing any watched root, so a
write task is never re-run on a half-edited tree. The report says so loudly and
result.json records both attempts. The fork tool does NOT honour this key.
Usage:
pi-task run SPEC.json [--dry-run] [--audit-dir DIR]
pi-task schema
pi-task selftest
"""
from __future__ import annotations
@@ -211,6 +218,157 @@ def validate_envelope(text: str):
return env, problems
# ------------------------------------------------------------ model fallback
# A tier's primary model can be unusable for reasons that have nothing to do
# with the task: an expired SSO session, a revoked key. pi reports that as a
# perfectly normal run — rc=0, agent_end, willRetry=false — with the failure
# visible ONLY as stopReason="error" on the assistant message (measured
# 2026-10-12, Bedrock with an invalid key). Before this, pi-task turned it into
# the misleading "no fenced json envelope".
#
# The retry is deliberately narrow. It happens only when the child provably did
# NOTHING: an auth-class error, zero tool results, and an unchanged boundary. A
# child that failed after acting may have half-edited a writable root, and
# re-running a write task on a dirty tree turns one clean failure into an
# unauditable mess. So any doubt means no retry, and the original failure
# stands.
AUTH_ERROR = re.compile(
r"token is expired|aws sso login|expiredtoken|unrecognizedclientexception"
r"|security token included in the request is (?:invalid|expired)"
r"|could not load credentials|unable to locate credentials|credentialsprovidererror"
r"|credentials? (?:have|has) expired|\b401\b|unauthori[sz]ed"
r"|invalid[ _-]?(?:x-)?api[ _-]?key|no api key|authentication (?:failed|error)",
re.IGNORECASE)
MIN_FALLBACK_WALL_S = 30
def child_error(events: list):
"""errorMessage of the LAST assistant message if it ended in error, else None."""
last = None
for ev in events:
m = ev.get("message") or {}
if ev.get("type") == "message_end" and m.get("role") == "assistant":
last = m
if last and last.get("stopReason") == "error":
return last.get("errorMessage") or "stopReason=error with no errorMessage"
return None
def fallback_decision(events: list, delta: list, remaining_s: float, has_fallback: bool):
"""(retry?, reason). Pure, so selftest can exercise both sides of every gate."""
if not has_fallback:
return False, "no fallback profile for this effort"
err = child_error(events)
if not err:
return False, "primary did not end in a model error"
if not AUTH_ERROR.search(err):
return False, f"primary error is not auth-class: {err[:200]}"
tools = sum(len(ev.get("toolResults") or []) for ev in events if ev.get("type") == "turn_end")
if tools:
return False, f"primary had already run {tools} tool call(s); not safe to re-run"
if delta:
return False, "primary changed a watched root; not safe to re-run"
if remaining_s < MIN_FALLBACK_WALL_S:
return False, f"only {remaining_s:.0f}s of wall budget left"
return True, f"primary auth failure before any action: {err[:200]}"
def _profile(raw, inherit_thinking=None):
if not isinstance(raw, dict) or not raw.get("provider") or not raw.get("id"):
return None
out = {"provider": raw["provider"], "id": raw["id"]}
thinking = raw.get("thinking", inherit_thinking)
if thinking is not None:
out["thinking"] = thinking
return out
def resolve_profiles(settings: dict, effort: str):
"""(primary|None, fallback|None, warnings[]).
The primary comes from pi-fork.effortProfiles, shared with the fork tool on
purpose. The fallback lives in pi-task's OWN namespace, pi-task.fallbackProfiles,
because only pi-task honours it. Putting it inside the fork profile would
suggest fork falls back too — it does not; pi-fork's parser drops unknown keys
without a word.
A fallback that omits `thinking` inherits the primary's. Defaulting to "off",
as the primary path does, would let a tier lose its reasoning silently."""
primary = ((settings.get("pi-fork") or {}).get("effortProfiles") or {}).get(effort)
raw = ((settings.get("pi-task") or {}).get("fallbackProfiles") or {}).get(effort)
warnings = []
if raw is None or not primary:
return primary, None, warnings
fb = _profile(raw, inherit_thinking=primary.get("thinking"))
if fb is None:
warnings.append(f"pi-task.fallbackProfiles.{effort} needs provider and id; ignored")
elif (fb["provider"], fb["id"]) == (primary.get("provider"), primary.get("id")):
warnings.append(f"pi-task.fallbackProfiles.{effort} is the primary model itself; ignored")
fb = None
return primary, fb, warnings
def _selftest_fallback() -> list:
"""Fallback decision + profile resolution, two-sided.
The event shapes are not invented: they mirror `pi -p --mode json` measured
2026-10-12 against Bedrock with an invalid key — rc=0, a normal agent_end
with willRetry=false, and the failure present ONLY as stopReason="error" +
errorMessage on the assistant message. The SSO string is the verbatim text
a user saw from an expired `aws sso` session the same day."""
bad_key = "UnrecognizedClientException: The security token included in the request is invalid."
sso = ("Token is expired. To refresh this SSO session run 'aws sso login' "
"with the corresponding profile.")
def run(err=None, tools=0, stop=None):
stop = stop or ("error" if err else "stop")
msg = {"role": "assistant", "stopReason": stop, "content": [],
**({"errorMessage": err} if err else {})}
evs = [{"type": "agent_start"}, {"type": "turn_start"},
{"type": "message_end", "message": msg}]
evs += [{"type": "turn_end", "toolResults": [{}] * tools}]
evs += [{"type": "agent_end", "willRetry": False, "messages": [msg]}]
return evs
def fb(events, delta=(), remaining=600, has_fb=True):
return fallback_decision(events, list(delta), remaining, has_fb)[0]
prim = {"provider": "p", "id": "big", "thinking": "high"}
alt = {"provider": "q", "id": "alt", "thinking": "medium"}
s = lambda fbp: {"pi-fork": {"effortProfiles": {"deep": prim}},
"pi-task": {"fallbackProfiles": {"deep": fbp}}}
p0, f_ok, _ = resolve_profiles(s(alt), "deep")
_, f_inh, _ = resolve_profiles(s({"provider": "q", "id": "alt"}), "deep")
_, f_same, w_same = resolve_profiles(s(dict(prim)), "deep")
_, f_bad, w_bad = resolve_profiles(s({"provider": "q"}), "deep")
_, f_none, w_none = resolve_profiles({"pi-fork": {"effortProfiles": {"deep": prim}}}, "deep")
return [
("fallback: POSITIVE bad-key auth error, 0 tools, clean tree => retry",
fb(run(bad_key))),
("fallback: POSITIVE expired-SSO error (verbatim) => retry", fb(run(sso))),
("fallback: auth error AFTER a tool call => NO retry", not fb(run(bad_key, tools=1))),
("fallback: auth error but a root CHANGED => NO retry",
not fb(run(bad_key), delta=["/r: working tree changed"])),
("fallback: non-auth error (throttling) => NO retry",
not fb(run("ThrottlingException: Too many requests, please wait"))),
("fallback: successful run => NO retry", not fb(run())),
("fallback: no fallback configured => NO retry", not fb(run(bad_key), has_fb=False)),
("fallback: under 30 s of wall budget left => NO retry",
not fb(run(bad_key), remaining=10)),
("child_error: surfaces the hidden stopReason=error message",
child_error(run(bad_key)) == bad_key and child_error(run()) is None),
("profiles: primary resolved from pi-fork.effortProfiles", p0 == prim),
("profiles: fallback resolved from pi-task.fallbackProfiles", f_ok == alt),
("profiles: fallback WITHOUT thinking inherits the primary's, not 'off'",
(f_inh or {}).get("thinking") == "high"),
("profiles: fallback identical to primary is dropped, with a warning",
f_same is None and bool(w_same)),
("profiles: malformed fallback (no id) is dropped, with a warning",
f_bad is None and bool(w_bad)),
("profiles: no fallback configured => None, silently", f_none is None and not w_none),
]
def cmd_selftest(_args):
"""Two-sided: known-good must pass, each known-bad must fail, and the
boundary detector must both see a change and see no change. Aborts loudly
@@ -281,6 +439,7 @@ def cmd_selftest(_args):
checks.append(("boundary/git: a COMMIT (HEAD move) IS detected",
any("HEAD" in x for x in boundary_delta(gb, boundary([str(g)])))))
checks += _selftest_fallback()
for name, cond in checks:
fails += 0 if cond else 1
print(f" {'ok ' if cond else 'BAD'} {name}")
@@ -293,54 +452,25 @@ def cmd_selftest(_args):
# ------------------------------------------------------------------------- main
def cmd_run(args):
spec_path = Path(args.spec)
spec = json.loads(spec_path.read_text())
missing = [k for k in REQUIRED_SPEC if k not in spec]
if missing:
die(f"spec missing required keys: {missing}")
profiles = (json.loads(SETTINGS.read_text()).get("pi-fork") or {}).get("effortProfiles") or {}
prof = profiles.get(spec["effort"])
if not prof:
die(f"effort '{spec['effort']}' not in {SETTINGS}: pi-fork.effortProfiles ({list(profiles)})")
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
audit = Path(args.audit_dir) if args.audit_dir else AUDIT_ROOT / f"{stamp}-{spec['id']}"
audit.mkdir(parents=True, exist_ok=True)
prompt = build_prompt(spec)
roots = spec.get("roots") or []
budget = spec.get("budget") or {}
wall = int(budget.get("wall_s", 600))
argv = ["pi", "-p", "--mode", "json",
def _argv(prof: dict, session_id: str, session_dir: Path, prompt: str) -> list:
return ["pi", "-p", "--mode", "json",
"--provider", prof["provider"], "--model", prof["id"],
"--thinking", prof.get("thinking", "off"),
"--no-extensions",
"--session-id", f"pitask-{spec['id']}-{stamp}",
"--session-dir", str(audit / "sess"),
"--session-id", session_id,
"--session-dir", str(session_dir),
prompt]
(audit / "spec.json").write_text(json.dumps(spec, indent=2))
(audit / "prompt.txt").write_text(prompt)
(audit / "argv.json").write_text(json.dumps(argv[:-1] + ["<prompt in prompt.txt>"], indent=2))
if args.dry_run:
print(prompt); print(f"\n--- dry run, nothing launched. audit: {audit}"); return 0
before = boundary(roots)
(audit / "boundary-before.json").write_text(json.dumps(before, indent=2))
def run_attempt(argv: list, timeout: float) -> dict:
"""Run one child and parse its NDJSON: final assistant text, cost, tool calls."""
t0 = time.time()
try:
p = subprocess.run(argv, capture_output=True, text=True, timeout=wall)
p = subprocess.run(argv, capture_output=True, text=True, timeout=timeout)
raw, rc = p.stdout, p.returncode
except subprocess.TimeoutExpired as e:
raw, rc = (e.stdout or ""), 124
elapsed = time.time() - t0
(audit / "raw.ndjson").write_text(raw if isinstance(raw, str) else raw.decode())
after = boundary(roots)
(audit / "boundary-after.json").write_text(json.dumps(after, indent=2))
# ---- parse the NDJSON: final assistant text, cost, tool calls
raw = raw if isinstance(raw, str) else raw.decode()
events, text, cost, tools = [], "", 0.0, 0
for line in raw.splitlines():
try:
@@ -356,9 +486,89 @@ def cmd_run(args):
if ev.get("type") == "message_end" and (ev.get("message") or {}).get("role") == "assistant":
text = "".join(c.get("text", "") for c in ev["message"].get("content", [])
if c.get("type") == "text") or text
return {"raw": raw, "rc": rc, "events": events, "text": text, "cost": cost,
"tools": tools, "elapsed": time.time() - t0, "error": child_error(events)}
def cmd_run(args):
spec_path = Path(args.spec)
spec = json.loads(spec_path.read_text())
missing = [k for k in REQUIRED_SPEC if k not in spec]
if missing:
die(f"spec missing required keys: {missing}")
settings = json.loads(SETTINGS.read_text())
prof, fb_prof, fb_warnings = resolve_profiles(settings, spec["effort"])
if not prof:
profiles = (settings.get("pi-fork") or {}).get("effortProfiles") or {}
die(f"effort '{spec['effort']}' not in {SETTINGS}: pi-fork.effortProfiles ({list(profiles)})")
for w in fb_warnings:
print(f"pi-task: warning: {w}", file=sys.stderr)
stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
audit = Path(args.audit_dir) if args.audit_dir else AUDIT_ROOT / f"{stamp}-{spec['id']}"
audit.mkdir(parents=True, exist_ok=True)
prompt = build_prompt(spec)
roots = spec.get("roots") or []
budget = spec.get("budget") or {}
wall = int(budget.get("wall_s", 600))
session_id = f"pitask-{spec['id']}-{stamp}"
argv = _argv(prof, session_id, audit / "sess", prompt)
fb_argv = _argv(fb_prof, f"{session_id}-fallback", audit / "sess-fallback", prompt) if fb_prof else None
(audit / "spec.json").write_text(json.dumps(spec, indent=2))
(audit / "prompt.txt").write_text(prompt)
(audit / "argv.json").write_text(json.dumps(argv[:-1] + ["<prompt in prompt.txt>"], indent=2))
if fb_argv:
(audit / "argv.fallback.json").write_text(
json.dumps(fb_argv[:-1] + ["<prompt in prompt.txt>"], indent=2))
if args.dry_run:
print(prompt)
if fb_prof:
print(f"\n--- fallback if the primary fails auth before acting: "
f"{fb_prof['provider']}/{fb_prof['id']}")
print(f"\n--- dry run, nothing launched. audit: {audit}"); return 0
before = boundary(roots)
(audit / "boundary-before.json").write_text(json.dumps(before, indent=2))
used, att = prof, run_attempt(argv, wall)
attempts = [{"provider": prof["provider"], "id": prof["id"], "thinking": prof.get("thinking", "off"),
"rc": att["rc"], "elapsed_s": round(att["elapsed"], 1), "tool_results": att["tools"],
"cost_usd": round(att["cost"], 6), "error": att["error"]}]
after = boundary(roots)
fallback = {"configured": bool(fb_prof), "used": False, "reason": None}
if fb_prof:
retry, why = fallback_decision(att["events"], boundary_delta(before, after),
wall - att["elapsed"], True)
fallback["reason"] = why
if retry:
# Keep the failed primary's evidence; raw.ndjson is always the run that produced the verdict.
(audit / "raw.primary.ndjson").write_text(att["raw"])
(audit / "boundary-after.primary.json").write_text(json.dumps(after, indent=2))
used, att = fb_prof, run_attempt(fb_argv, wall - att["elapsed"])
attempts.append({"provider": fb_prof["provider"], "id": fb_prof["id"],
"thinking": fb_prof.get("thinking", "off"), "rc": att["rc"],
"elapsed_s": round(att["elapsed"], 1), "tool_results": att["tools"],
"cost_usd": round(att["cost"], 6), "error": att["error"]})
fallback.update({"used": True, "from": f"{prof['provider']}/{prof['id']}",
"to": f"{fb_prof['provider']}/{fb_prof['id']}"})
after = boundary(roots)
raw, rc, events, text = att["raw"], att["rc"], att["events"], att["text"]
cost = sum(a["cost_usd"] for a in attempts) # a failed primary is billed too, if at all
tools = att["tools"]
elapsed = sum(a["elapsed_s"] for a in attempts)
(audit / "raw.ndjson").write_text(raw)
(audit / "boundary-after.json").write_text(json.dumps(after, indent=2))
# ---- validate the envelope: unparseable == FAILED, no matter how fluent
env, problems = validate_envelope(text)
if att["error"]:
# pi exits 0 on a provider error; without this line the only symptom
# would be "no fenced json envelope", which points at the wrong layer.
# First: it is the cause, and the missing envelope merely its symptom.
problems.insert(0, f"child model error ({used['provider']}/{used['id']}): {att['error'][:300]}")
if rc != 0:
problems.append(f"child exited rc={rc}" + (" (wall-clock budget)" if rc == 124 else ""))
if not any(e.get("type") == "agent_end" for e in events):
@@ -380,16 +590,21 @@ def cmd_run(args):
verdict = "PASS" if not problems else "FAIL"
result = {"verdict": verdict, "problems": problems, "envelope": env,
"metrics": {"cost_usd": round(cost, 6), "elapsed_s": round(elapsed, 1),
"tool_results": tools, "rc": rc, "model": prof["id"],
"effort": spec["effort"]},
"tool_results": tools, "rc": rc, "model": used["id"],
"provider": used["provider"], "effort": spec["effort"]},
"fallback": fallback, "attempts": attempts,
"boundary_delta": delta, "boundary_violations": violations,
"audit_dir": str(audit),
"raw_text_if_envelope_missing": None if env else text[-2000:]}
(audit / "result.json").write_text(json.dumps(result, indent=2))
# ---- report to the parent: verdict first, then what to spot-check
print(f"pi-task {verdict} id={spec['id']} {prof['id']} "
f"${cost:.4f} {elapsed:.0f}s tools={tools} rc={rc}")
print(f"pi-task {verdict} id={spec['id']} {used['id']}"
+ (" [FALLBACK]" if fallback["used"] else "")
+ f" ${cost:.4f} {elapsed:.0f}s tools={tools} rc={rc}")
if fallback["used"]:
print(f" ⚠ fallback: {fallback['from']} -> {fallback['to']} — {fallback['reason']}")
print(" This result is from the FALLBACK model; weigh it accordingly.")
for p_ in problems:
print(f" ! {p_}")
if env:
@@ -410,7 +625,9 @@ def cmd_schema(_args):
"id": "short-slug-used-in-session-id-and-audit-path",
"goal": "verbatim task statement",
"deliverable": "the exact shape of answer wanted",
"effort": "fast | balanced | deep (resolved via ~/.pi/agent/settings.json pi-fork.effortProfiles)",
"effort": "fast | balanced | deep (resolved via ~/.pi/agent/settings.json pi-fork.effortProfiles; "
"optional per-effort fallback in pi-task.fallbackProfiles, used ONLY when the primary "
"fails auth before any tool call or file change)",
"read_only": True,
"roots": ["/workspace/repo-the-task-touches"],
"write_allowed": ["subset of roots the child may CHANGE; only meaningful when "