ollie/scripts/u/cascade

#!/usr/bin/env bash
# u/cascade — script-driven hub-and-spoke cascading orchestrator.
#
# An expensive model decomposes a task into mechanical subtasks, then cheap
# models execute them in parallel. The script (not an agent) coordinates.
#
# Usage: cascade [flags] [prompt...]
#        echo "prompt" | cascade [flags]
#
# Flags:
#   -pb BACKEND   planner backend (default: $OLLIE_CASCADE_PLAN_BACKEND or anthropic)
#   -pm MODEL     planner model (default: $OLLIE_CASCADE_PLAN_MODEL or claude-sonnet-4-20250514)
#   -wb BACKEND   worker backend (default: $OLLIE_CASCADE_BACKEND or ollama)
#   -wm MODEL     worker model (default: $OLLIE_CASCADE_MODEL or qwen2.5-coder:14b)
#   -max-workers N    max concurrent workers (default: 8)
#   -retries N        max attempts per worker (default: 1, no retry)
#   -synthesize       collect worker results and ask the planner to produce a summary
#   -json-output      structured JSON synthesis (implies -synthesize)
set -euo pipefail

OLLIE="${OLLIE:-$HOME/mnt/ollie}"
B="$OLLIE/s/b"

PB="${OLLIE_CASCADE_PLAN_BACKEND:-openrouter}"
PM="${OLLIE_CASCADE_PLAN_MODEL:-deepseek/deepseek-v4-pro}"
WB="${OLLIE_CASCADE_BACKEND:-openrouter}"
WM="${OLLIE_CASCADE_MODEL:-deepseek/deepseek-v4-flash}"

die() { printf 'cascade: %s\n' "$*" >&2; exit 1; }

synthesize=0
json_output=0
max_workers=8
retries=1

while [ $# -gt 0 ]; do
    case $1 in
        -pb) PB=$2; shift ;;
        -pm) PM=$2; shift ;;
        -wb) WB=$2; shift ;;
        -wm) WM=$2; shift ;;
        -max-workers) max_workers=$2; shift ;;
        -retries) retries=$2; shift ;;
        -synthesize) synthesize=1 ;;
        -json-output) json_output=1; synthesize=1 ;;
        --)  shift; break ;;
        -*)  die "unknown flag: $1" ;;
        *)   break ;;
    esac
    shift
done

if [ $# -gt 0 ]; then
    task="$*"
elif ! [ -t 0 ]; then
    task=$(cat)
else
    die "prompt required"
fi
[ -z "$task" ] && die "empty prompt"

case "$max_workers" in
    ''|*[!0-9]*|0) die "-max-workers must be a positive integer" ;;
esac
case "$retries" in
    ''|*[!0-9]*|0) die "-retries must be a positive integer" ;;
esac

# Step 1: Ask the planner to decompose into subtasks (JSON array).
plan=$("$B" -backend "$PB" -model "$PM" -json \
    "Decompose this task into independent, mechanical subtasks that a weak model can execute without judgment. Each subtask must be fully self-contained (include all needed context). Output a JSON array of objects with \"id\" (short string) and \"instruction\" (complete task description). Keep it minimal — only decompose what benefits from parallelism.

TASK:
$task")

# Validate we got JSON with correct schema.
printf '%s' "$plan" | jq empty 2>/dev/null || die "planner did not return valid JSON: $plan"
printf '%s' "$plan" | jq -e 'type=="array" and length>0 and all(.[]; (.id|type=="string" and length>0) and (.instruction|type=="string" and length>0))' >/dev/null \
    || die "plan schema invalid (need array of {id: string, instruction: string})"

# Step 2: Spawn workers in parallel, collect session paths.
count=$(printf '%s' "$plan" | jq 'length')
[ "$count" -eq 0 ] && die "planner returned empty plan"

pids=()
workers=()
base="cascade-$(date +%s%N)"
tmpdir=$(mktemp -d)
trap 'rm -rf "$tmpdir"' EXIT

for i in $(seq 0 $((count - 1))); do
    id=$(printf '%s' "$plan" | jq -r ".[$i].id")
    instruction=$(printf '%s' "$plan" | jq -r ".[$i].instruction")
    sid="${base}-${id}"
    workers+=("$sid")

    # Throttle: wait for a slot if at max concurrency.
    while [ "$(jobs -rp | wc -l)" -ge "$max_workers" ]; do sleep 0.5; done

    (
        attempt=0
        while [ "$attempt" -lt "$retries" ]; do
            attempt=$((attempt + 1))
            name="${sid}$( [ "$attempt" -gt 1 ] && echo "-r${attempt}" || true)"
            "$B" -name "$name" -backend "$WB" -model "$WM" -keep \
                "You are a worker agent. Execute the following task exactly. Output only the result.

$instruction" > "${tmpdir}/${sid}.out" 2>&1 && touch "${tmpdir}/${sid}.ok" && break
            # Backoff before retry.
            sleep "$((attempt * 2))"
        done
        echo "$attempt" > "${tmpdir}/${sid}.attempts"
        echo "$name" > "${tmpdir}/${sid}.session"
    ) &
    pids+=($!)
done

# Step 3: Wait for all workers.
for pid in "${pids[@]}"; do wait "$pid" || true; done

# Step 4: Gather results as structured JSON array.
MAX_OUTPUT=${OLLIE_CASCADE_MAX_OUTPUT:-8000}  # max chars per worker output
failures=0
results_json="["
first=1
for sid in "${workers[@]}"; do
    if [ -f "${tmpdir}/${sid}.ok" ]; then
        status="done"
    else
        status="failed"
        failures=$((failures + 1))
    fi
    out=$(cat "${tmpdir}/${sid}.out" 2>/dev/null || echo "(no output)")
    attempts=$(cat "${tmpdir}/${sid}.attempts" 2>/dev/null || echo 1)
    session_name=$(cat "${tmpdir}/${sid}.session" 2>/dev/null || echo "$sid")
    if [ "${#out}" -gt "$MAX_OUTPUT" ]; then
        out="${out:0:$MAX_OUTPUT}
...(truncated, ${#out} chars total)"
    fi
    [ "$first" -eq 1 ] && first=0 || results_json="${results_json},"
    results_json="${results_json}$(jq -n --arg id "$sid" --arg status "$status" --arg output "$out" --argjson attempts "$attempts" --arg session "$session_name" \
        '{id: $id, status: $status, attempts: $attempts, session: $session, output: $output}')"
done
results_json="${results_json}]"

# Step 5: Synthesize or just print results.
# Exit code contract:
#   No synthesis: exit 1 if any worker failed (caller must handle partial results).
#   Synthesis: exit 0 — the planner accounts for failures in its answer.
if [ "$synthesize" -eq 0 ]; then
    printf '%s\n' "$results_json" | jq .
    [ "$failures" -gt 0 ] && exit 1
else
    synth_flags=(-backend "$PB" -model "$PM")
    if [ "$json_output" -eq 1 ]; then
        synth_flags+=(-json)
        synth_prompt="Synthesize these worker results into a structured JSON response for the original task. Note: workers with status \"failed\" did not complete successfully — include them in errors.

Output schema: {\"summary\": \"<one-line summary>\", \"answer\": \"<full synthesized answer>\", \"worker_results\": [{\"id\": \"<worker-id>\", \"status\": \"done|failed\", \"attempts\": <int>, \"session\": \"<session-name>\", \"output\": \"<worker output>\"}], \"errors\": [\"<description of any failures>\"]}

ORIGINAL TASK:
$task

WORKER RESULTS (JSON):
$results_json"
    else
        synth_prompt="Synthesize these worker results into a final, coherent answer to the original task. Note: workers with status \"failed\" did not complete successfully — account for this in your synthesis.

ORIGINAL TASK:
$task

WORKER RESULTS (JSON):
$results_json"
    fi
    "$B" "${synth_flags[@]}" "$synth_prompt"
fi