#!/usr/bin/env bash
# test-wakeup-mode.sh — print and JSON mode (`pi -p`, `--mode json`) must not
# do the palace start-up work of extensions/pi/mempalace.ts unless
# MEMPALACE_WAKEUP_PRINT=1. That work is the wake-up injection (status, diary,
# tier-1 mailbox) and the mid-session mailbox poll. TUI and RPC sessions keep
# both.
#
# WHY (2026-10-13, tor-ms22). A one-shot run is a script asking one question.
# The wake-up cost it several palace round trips and a large orientation
# message, and in one probe the model read the injected diary's open items as
# its task. Given "Reply with exactly: PONG", it ran about 12 tool calls and
# filed a wrong drawer in the shared palace.
#
# Same method as test-feed-exit.sh: the SHIPPED text is type-stripped and loaded
# under a stub pi. A local JSON-RPC server, in a separate process, logs every
# tools/call. The driver fires session_start, before_agent_start, agent_settled
# and session_shutdown with the ctx in DRIVER_CTX, and reports whether
# before_agent_start returned the wake-up message.
#   tui              wake-up message, status + diary + event_list calls
#   print / json     no message and NO palace call at all
#   print, resumed   (`pi -p -c`; resume sets wokeUp) the mid-session poll
#                    must stay off too: no event_list call
#   tui, resumed     control for the line above: the poll DOES run there,
#                    so a zero is not vacuous
#   print + MEMPALACE_WAKEUP_PRINT=1   wake-up message, as before
#   rpc              wake-up message (long-lived, keeps it)
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SRC="${1:-$REPO_ROOT/extensions/pi/mempalace.ts}"
WORK="$(mktemp -d)"
SERVER_PID=""
cleanup() { [ -n "$SERVER_PID" ] && kill "$SERVER_PID" 2>/dev/null; rm -rf "$WORK"; }
trap cleanup EXIT
[ -r "$SRC" ] || { echo "FAIL: cannot read $SRC" >&2; exit 2; }

cat >"$WORK/extract.mjs" <<'EXTRACT'
import { readFileSync, writeFileSync } from "node:fs";
import { stripTypeScriptTypes } from "node:module";
let src = readFileSync(process.argv[2], "utf8");
const imp = 'import { Type } from "typebox";';
if (!src.includes(imp)) throw new Error(`expected import line not found: ${imp}`);
src = src.replace(imp, 'const Type = { Unsafe: (s) => s, Object: () => ({}) };');
if (!src.includes("before_agent_start")) throw new Error("extension has no before_agent_start handler");
writeFileSync(process.argv[3], stripTypeScriptTypes(src, { mode: "strip" }));
EXTRACT
node --no-warnings "$WORK/extract.mjs" "$SRC" "$WORK/ext.mjs" || exit 2
node --check "$WORK/ext.mjs" || { echo "FAIL: stripped extension does not parse" >&2; exit 2; }
echo "[extract] type-stripped $(basename "$SRC") ($(wc -c <"$WORK/ext.mjs") bytes)"

cat >"$WORK/server.mjs" <<'SERVER'
import { createServer } from "node:http";
import { appendFileSync } from "node:fs";
const tools = ["mempalace_status", "mempalace_diary_read", "mempalace_event_list"]
  .map((name) => ({ name, description: "", inputSchema: { type: "object" } }));
const server = createServer((req, res) => {
  let body = "";
  req.on("data", (c) => (body += c));
  req.on("end", () => {
    const msg = JSON.parse(body);
    if (msg.id === undefined) { res.writeHead(202); res.end(); return; }
    const reply = (result) => {
      res.writeHead(200, { "content-type": "application/json", connection: "close" });
      res.end(JSON.stringify({ jsonrpc: "2.0", id: msg.id, result }));
    };
    if (msg.method === "initialize") return reply({ protocolVersion: "2024-11-05", capabilities: {}, serverInfo: { name: "fake", version: "0" } });
    if (msg.method === "tools/list") return reply({ tools });
    if (msg.method === "tools/call") {
      const name = msg.params?.name;
      appendFileSync(process.env.CALL_LOG, `${name}\n`);
      const text = name === "mempalace_event_list" ? JSON.stringify({ events: [], count: 0 }) : `${name} ok`;
      return reply({ content: [{ type: "text", text }] });
    }
    reply({});
  });
});
server.listen(0, "127.0.0.1", () => console.log(server.address().port));
SERVER
CALL_LOG="$WORK/calls.log" node "$WORK/server.mjs" >"$WORK/port" &
SERVER_PID=$!
for _ in $(seq 50); do [ -s "$WORK/port" ] && break; sleep 0.1; done
PORT="$(cat "$WORK/port")"
[ -n "$PORT" ] || { echo "FAIL: fake palace did not start" >&2; exit 2; }

cat >"$WORK/driver.mjs" <<'DRIVER'
const ctx = { ...JSON.parse(process.env.DRIVER_CTX || "{}"), ui: { notify() {} } };
const handlers = {};
const pi = new Proxy({}, {
  get: (_t, k) => k === "on" ? (ev, fn) => { (handlers[ev] ||= []).push(fn); } : () => {},
});
const { default: ext } = await import("./ext.mjs");
await ext(pi);
const fire = async (ev, event) => {
  const out = [];
  for (const fn of handlers[ev] || []) out.push(await fn({ type: ev, ...event }, ctx));
  return out;
};
await fire("session_start", { reason: process.env.DRIVER_REASON || "startup" });
const results = await fire("before_agent_start", { prompt: "PONG" });
const woke = results.some((r) => r?.message?.customType === "mempalace-wakeup");
console.log(`WAKEUP_MESSAGE=${woke ? "yes" : "no"}`);
await fire("agent_settled", {});
await new Promise((r) => setTimeout(r, 800)); // the mailbox poll is not awaited
await fire("session_shutdown", { reason: "quit" });
DRIVER

failures=0
check() { if [ "$1" = 0 ]; then echo "ok   $2"; else echo "FAIL $2"; failures=$((failures + 1)); fi; }
calls() { grep -c "^$1\$" "$WORK/calls.log" || true; }

run_case() { # $1 ctx JSON, $2 session_start reason, [$3 WAKEUP_PRINT] -> MSG RC ALL STATUS DIARY EVENTS
  : >"$WORK/calls.log"
  set +e
  out="$(cd "$WORK" && env -u MEMPALACE_REMOTE_TOKEN -u MEMPALACE_MAILBOX \
    MEMPALACE_REMOTE_URL="http://127.0.0.1:$PORT/mcp" MEMPALACE_PI_DEVICE=testdev \
    MEMPALACE_FEED=0 DRIVER_CTX="$1" DRIVER_REASON="$2" MEMPALACE_WAKEUP_PRINT="${3:-}" \
    timeout 60 node --no-warnings driver.mjs 2>"$WORK/driver.err")"
  RC=$?
  set -e
  MSG="$(printf '%s\n' "$out" | sed -n 's/^WAKEUP_MESSAGE=//p')"
  ALL="$(wc -l <"$WORK/calls.log" | tr -d ' ')"
  STATUS="$(calls mempalace_status)"; DIARY="$(calls mempalace_diary_read)"; EVENTS="$(calls mempalace_event_list)"
  echo "[$4] rc=$RC message=${MSG:-?} calls=$ALL (status=$STATUS diary=$DIARY event_list=$EVENTS)"
}

run_case '{"mode":"tui","hasUI":true}' startup "" tui
check "$([ "$RC" = 0 ] && [ "$MSG" = yes ] && [ "$STATUS" = 1 ] && [ "$DIARY" = 1 ] && [ "$EVENTS" -ge 1 ] && echo 0 || echo 1)" \
  "tui: wake-up injected, with status, diary and mailbox reads"

for m in print json; do
  run_case "{\"mode\":\"$m\",\"hasUI\":false}" startup "" "$m"
  check "$([ "$RC" = 0 ] && [ "$MSG" = no ] && [ "$ALL" = 0 ] && echo 0 || echo 1)" \
    "$m: no wake-up message and no palace call at all (calls=$ALL)"
done

run_case '{"mode":"print","hasUI":false}' resume "" "print resumed"
check "$([ "$RC" = 0 ] && [ "$MSG" = no ] && [ "$EVENTS" = 0 ] && echo 0 || echo 1)" \
  "print resumed (pi -p -c): the mid-session mailbox poll stays off (event_list=$EVENTS)"

run_case '{"mode":"tui","hasUI":true}' resume "" "tui resumed (control)"
check "$([ "$RC" = 0 ] && [ "$MSG" = no ] && [ "$EVENTS" -ge 1 ] && echo 0 || echo 1)" \
  "control: tui resumed skips the injection but DOES poll (event_list=$EVENTS), so the zero above is not vacuous"

run_case '{"mode":"print","hasUI":false}' startup 1 "print+WAKEUP_PRINT=1"
check "$([ "$RC" = 0 ] && [ "$MSG" = yes ] && [ "$STATUS" = 1 ] && echo 0 || echo 1)" \
  "print + MEMPALACE_WAKEUP_PRINT=1: wake-up injected, as before"

run_case '{"mode":"rpc","hasUI":true}' startup "" rpc
check "$([ "$RC" = 0 ] && [ "$MSG" = yes ] && echo 0 || echo 1)" "rpc: wake-up injected (long-lived)"

if [ "$failures" = 0 ]; then echo "PASS: all assertions held"; else echo "FAIL: $failures assertion(s) failed"; exit 1; fi