Harden Primary decode lifecycle and isolate MLX worker

deploy/install_decode_watchdog_launchd.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+# Install the external Primary decode watchdog as a per-user LaunchAgent.
+set -euo pipefail
+
+: "${KAKEYA_RUNTIME_REPO:?set KAKEYA_RUNTIME_REPO}"
+: "${KAKEYA_RUNTIME_PYTHON:?set KAKEYA_RUNTIME_PYTHON}"
+
+WATCHDOG_LABEL="${KAKEYA_WATCHDOG_LABEL:-ai.kakeya.decode-watchdog}"
+RUNTIME_LABEL="${KAKEYA_RUNTIME_LABEL:-ai.kakeya.grpc-runtime-prefill}"
+STALL_SECONDS="${KAKEYA_DECODE_STALL_SECONDS:-120}"
+INTERVAL_SECONDS="${KAKEYA_WATCHDOG_INTERVAL_SECONDS:-30}"
+LIVENESS_FILE="${KAKEYA_DECODE_LIVENESS_FILE:-$HOME/.kakeya/primary-decode-liveness.json}"
+STATE_FILE="${KAKEYA_WATCHDOG_STATE_FILE:-$HOME/.kakeya/decode-watchdog-state.json}"
+UNHEALTHY_FILE="${KAKEYA_RUNTIME_UNHEALTHY_FILE:-$HOME/.kakeya/primary-runtime-unhealthy.json}"
+PLIST="$HOME/Library/LaunchAgents/$WATCHDOG_LABEL.plist"
+LOG_FILE="$HOME/.kakeya/decode-watchdog.log"
+
+mkdir -p "$(dirname "$PLIST")" "$(dirname "$LOG_FILE")"
+cat > "$PLIST" <<EOF
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN"
+  "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0"><dict>
+  <key>Label</key><string>$WATCHDOG_LABEL</string>
+  <key>ProgramArguments</key><array>
+    <string>$KAKEYA_RUNTIME_PYTHON</string>
+    <string>$KAKEYA_RUNTIME_REPO/scripts/decode_watchdog.py</string>
+    <string>--liveness-file</string><string>$LIVENESS_FILE</string>
+    <string>--state-file</string><string>$STATE_FILE</string>
+    <string>--unhealthy-file</string><string>$UNHEALTHY_FILE</string>
+    <string>--runtime-label</string><string>$RUNTIME_LABEL</string>
+    <string>--stall-seconds</string><string>$STALL_SECONDS</string>
+  </array>
+  <key>StartInterval</key><integer>$INTERVAL_SECONDS</integer>
+  <key>RunAtLoad</key><true/>
+  <key>ProcessType</key><string>Background</string>
+  <key>StandardOutPath</key><string>$LOG_FILE</string>
+  <key>StandardErrorPath</key><string>$LOG_FILE</string>
+</dict></plist>
+EOF
+
+chmod 644 "$PLIST"
+DOMAIN="gui/$(id -u)"
+launchctl bootout "$DOMAIN/$WATCHDOG_LABEL" 2>/dev/null || true
+launchctl bootstrap "$DOMAIN" "$PLIST"
+launchctl kickstart -k "$DOMAIN/$WATCHDOG_LABEL"
+echo "installed $WATCHDOG_LABEL -> $PLIST"

docs/adr/0017-prefill-compute-worker-orchestration.md

@@ -31,6 +31,42 @@ Kakeya adopts three explicit fleet roles:
    user decode.
 3. **PREFILL_CACHE:** loads no model; spends RAM on immutable snapshots only.
 
+### Local decode process boundary
+
+Primary decode may run behind the feature flag `--decode-worker`. In this
+profile the gRPC/router process does not construct an MLX verifier and does not
+load model weights. It starts one child process over a mode-0600 Unix-domain
+socket; that child owns the single loaded model and all session K/V adapters.
+The transport is protocol-v1 length-prefixed JSON plus an optional opaque
+binary payload. It never uses pickle and is not exposed on TCP.
+
+The local protocol has six operations:
+
+- `Init(session_id, token_ids?)` creates/replaces isolated session K/V;
+- `ImportSnapshot(session_id, compatibility, payload)` imports the existing
+  allens portable snapshot format after exact compatibility validation;
+- `Append(session_id, token_ids)` commits an all-accepted block;
+- `GenerateStep(session_id)` atomically selects and commits one greedy token;
+- `Close(session_id)` releases that session's K/V; and
+- `Health()` returns protocol, process, model geometry, session and MLX-memory
+  state.
+
+The router serializes requests because the MVP MLX execution stream is shared.
+No per-token K/V snapshot crosses IPC. The router retains a proof checkpoint:
+the last imported immutable allens snapshot plus only tokens acknowledged
+after that boundary (or full acknowledged history when no snapshot exists).
+An IPC EOF, child exit, or operation timeout hard-kills the child, starts a new
+one, restores that checkpoint, and retries the unacknowledged current
+operation once. Because checkpoint state advances only after a correlated
+response and `GenerateStep` is atomic in the child, an ambiguous crash cannot
+duplicate a committed token. Client cancellation hard-kills an in-flight
+worker forward; the next request follows the same restore path.
+
+This process boundary is initially opt-in. The in-process MLX path remains the
+rollback path until the Mac acceptance gates pass. It is not a remote decode
+service and does not change the rule that no fleet RPC occurs in the token
+loop.
+
 ### Request flow
 
 For a cold append, the primary:
@@ -118,6 +154,10 @@ Blocking tests cover:
 - real MLX local-prefill vs remote-prefill-import continuation logits and
   argmax equivalence;
 - zero prefill/cache RPCs during `Generate`.
+- protocol correlation/version/size validation for local decode IPC;
+- in-process vs worker greedy-token parity;
+- cancellation hard-kill and child crash/timeout recovery from both
+  full-history and allens-snapshot-plus-proof checkpoints.
 
 ## Consequences
 
@@ -136,6 +176,10 @@ Costs:
   primary;
 - snapshots still duplicate bounded state at checkpoint boundaries;
 - MLX worker execution is serialized per loaded verifier in the MVP.
+- worker-mode locally computed prefill checkpoints are not republished by the
+  router, because export is deliberately absent from the token-loop protocol;
+  allens-produced snapshots remain importable and are the preferred durable
+  recovery boundary.
 
 ## Legacy paths
 

docs/ops/primary-decode-hardening.md

@@ -0,0 +1,79 @@
+# Primary decode hardening
+
+The Primary runtime publishes read-only diagnostics on
+`http://127.0.0.1:8091` by default:
+
+- `GET /healthz` combines decode liveness and unified-memory health.
+- `GET /v1/runtime/liveness` reports phase, session, token index, update
+  timestamp, and PID.
+- `GET /v1/runtime/memory` reports MLX active/cache/peak bytes, process RSS,
+  active sessions, and live KV bytes.
+
+The default memory policy warns at 18 GiB, stops admitting new sessions at
+20 GiB, and marks the runtime unhealthy at 21.5 GiB. When the final active
+session is removed under memory pressure, cleanup runs in this order:
+verifier reset, Python garbage collection, then `mlx.core.clear_cache()`.
+Thresholds and the diagnostics port are configurable with the
+`--memory-*-gb` and `--runtime-health-*` server flags.
+
+All removal paths (explicit close, client cancellation, LRU, TTL, INV-1, and
+INV-2) pass through the `SessionStore` removal hook. This releases the slab
+and removes any per-session verifier binding exactly once.
+
+## Decode watchdog
+
+Install the independent per-user LaunchAgent:
+
+```bash
+KAKEYA_RUNTIME_REPO=/path/to/repo \
+KAKEYA_RUNTIME_PYTHON=/path/to/python \
+KAKEYA_RUNTIME_LABEL=ai.kakeya.grpc-runtime-prefill \
+deploy/install_decode_watchdog_launchd.sh
+```
+
+Every 30 seconds the watchdog reads
+`~/.kakeya/primary-decode-liveness.json`. It restarts the configured Primary
+LaunchAgent only after observing the same decode token stale for at least
+120 seconds twice in succession. The watchdog also recycles a runtime that
+has written the 21.5 GiB unhealthy marker. Runtime startup clears a stale
+marker.
+
+`Session.generate(inter_token_timeout_s=...)` provides a client-side
+notification deadline. It cancels the stream and raises
+`InterTokenTimeoutError`; it does not attempt process recovery. Hard recovery
+remains the external watchdog's responsibility.
+
+## Isolated MLX decode worker
+
+Enable the opt-in process boundary by adding these flags to the Primary
+runtime command:
+
+```bash
+--backend mlx \
+--decode-worker \
+--decode-worker-timeout-s 120 \
+--decode-worker-startup-timeout-s 180
+```
+
+The router then loads no MLX model. A private mode-0600 UDS child owns the
+model and per-session K/V. Keep `--decode-worker-socket` unset for an
+automatically unique temporary path; set it only when launchd supervision or
+socket diagnostics require a stable path.
+
+Migration procedure:
+
+1. deploy with the flag absent and confirm the existing in-process smoke;
+2. enable `--decode-worker` on one Primary and verify `/healthz`, append,
+   streaming generation, cancellation, and one injected child kill;
+3. confirm the child PID changes, the current turn resumes from the imported
+   allens snapshot plus acknowledged proof checkpoint, and the router RSS does
+   not contain a second model;
+4. retain flag removal as rollback until the Mac parity, 100-session,
+   fault-injection, latency, and four-hour acceptance suite has passed.
+
+Worker mode can import allens portable snapshots but deliberately does not
+export a K/V snapshot per token. Locally computed prefill boundaries therefore
+are not republished through the router in this release. A cancellation kills
+the in-flight child forward; the affected session is still closed by the gRPC
+cancellation contract, while other retained session checkpoints are restored
+lazily if later used.

docs/primary_decode_acceptance.md

@@ -0,0 +1,112 @@
+# Primary decode Mac acceptance
+
+`scripts/bench_agentic/primary_decode_acceptance.py` is the blocking Mac
+acceptance runner for the isolated Primary decode worker. It emits JSON
+conforming to `schemas/primary_decode_acceptance.schema.json` and JUnit XML.
+The full mode includes the four-hour mixed-prompt run; it is only started by
+an explicit `--mode all` or `--mode endurance`.
+
+Example full command:
+
+```bash
+PYTHONPATH=.:sdks/python python3 \
+  scripts/bench_agentic/primary_decode_acceptance.py \
+  --mode all \
+  --grpc-address 127.0.0.1:50051 \
+  --tokenizer-id Qwen/Qwen3-0.6B \
+  --worker-control-command "/path/to/decode-worker-acceptance-adapter" \
+  --latency-baseline results/platform-tests/bench_mlx_verifier_1779507043.json \
+  --output results/platform-tests/primary_decode_acceptance.json \
+  --junit-output results/platform-tests/primary_decode_acceptance.xml
+```
+
+Start the runtime with its test-only, mode-0600 UDS enabled:
+
+```bash
+PYTHONPATH=.:sdks/python python3 scripts/start_grpc_runtime_server.py \
+  --backend mlx \
+  --verifier-id Qwen/Qwen3-0.6B \
+  --bind 127.0.0.1:50051 \
+  --capacity 128 --sink 4 --window 64 \
+  --decode-worker \
+  --decode-worker-timeout-s 110 \
+  --decode-worker-acceptance-socket /tmp/kakeya-decode-acceptance.sock
+```
+
+The corresponding adapter command is:
+
+```bash
+PYTHONPATH=. python3 \
+  scripts/bench_agentic/decode_worker_acceptance_adapter.py \
+  --socket /tmp/kakeya-decode-acceptance.sock
+```
+
+Run `footprint`, `disconnect`, `hang`, `kv-restore`, and `latency` separately
+for pre-CI/hardware checks. Reserve `--mode all` for the blocking release run:
+it deliberately includes the four-hour endurance workload.
+
+Short development runs can override `--session-count` and
+`--endurance-duration-s`, but the gates remain fixed at 100 sessions and
+14,400 seconds. A shortened run therefore emits useful diagnostics while
+remaining failed and cannot be mistaken for release evidence.
+
+## Decode-worker adapter contract
+
+The runtime provides
+`scripts/bench_agentic/decode_worker_acceptance_adapter.py`, passed through
+`--worker-control-command`. The harness starts it once per operation, writes
+one JSON object to stdin, and expects one JSON object on stdout:
+
+```json
+{
+  "schema_version": 1,
+  "operation": "snapshot",
+  "payload": {}
+}
+```
+
+```json
+{
+  "schema_version": 1,
+  "operation": "snapshot",
+  "ok": true,
+  "data": {}
+}
+```
+
+Nonzero exit status, malformed JSON, `ok: false`, or mismatched operation is
+reported as a gate error. The adapter must be test-only/local-only and must
+not expose fault injection on a network listener.
+
+Required operations:
+
+- `snapshot`: returns integer `runtime_pid`, `worker_pid`,
+  `worker_restart_count`, `process_footprint_bytes`, `active_sessions`, and
+  `active_generations`. The footprint must cover the runtime plus owned decode
+  worker, not only the router process.
+- `inject_hang`: accepts payload `phase: "next_forward"` and
+  `expected_worker_pid`; atomically arms exactly the next MLX forward in that
+  worker and returns `accepted: true`. The injected forward must remain hung
+  until the normal watchdog/recycle path kills the worker.
+- `kv_restore_parity`: accepts `prompt_token_ids` and performs a normal
+  decode, worker recycle, restore from the persisted Allens KV snapshot plus
+  proof checkpoint, and repeated decode. It returns
+  `baseline_first_token_id`, `restored_first_token_id`,
+  `baseline_logits_sha256`, `restored_logits_sha256`, and the literal
+  `restore_source: "allens_kv+proof_checkpoint"`. SHA-256 is over a
+  canonical, dtype-preserving byte representation of the last-token logits.
+
+The gRPC/router branch must additionally guarantee that closing the client
+channel cancels the in-flight Generate and removes its session. The harness
+starts timing only after `snapshot.active_generations` becomes nonzero, then
+requires both active counters to become zero within five seconds.
+
+## Existing benchmark reuse
+
+Mixed-prompt endurance records are aggregated by
+`inference_engine.bench.session_long_run.aggregate_run`. Latency baseline
+input accepts either a previous acceptance report's top-level `latency`
+summary or the existing `bench_mlx_verifier.py` JSON. For the latter, the
+reported MLX generation mean per token is used as the reference for both
+p50 and p95 and is identified as such in `baseline_source`; a prior
+acceptance report is preferred once available.