fix: end banner starts a new layout; outcome colour

lantern/bench.py

@@ -1,65 +1,104 @@
-"""Benchmark series: play one level across seeds, optionally carrying the journal forward."""
-
-from __future__ import annotations
-
-import time
-import uuid
-
-from .agent import EpisodeConfig, run_episode
-from .env import TASKS, DungeonEnv
-from .journal import merge_lessons
-from .llm import ChatClient
-from .store import EpisodeStore
-
-SEED_BASE = 1000  # benchmark seeds stay clear of the low seeds visitors try in the Space
-
-
-def bench_seeds(n: int) -> list[int]:
-    return [SEED_BASE + i for i in range(n)]
-
-
-def play_series(task_id: str, seeds: list[int], use_lessons: bool, tag: str,
-                client: ChatClient, store: EpisodeStore, source: str = "bench") -> list[dict]:
-    """Play `seeds` in order on one level and save every episode; returns per-episode summaries.
-
-    Each series starts with an empty journal; with `use_lessons`, lessons written after
-    each episode are merged in before the next seed.
-    """
-    config = EpisodeConfig(use_lessons=use_lessons, reflect=use_lessons)
-    lessons: list[str] = []
-    summaries = []
-    for seed in seeds:
-        env = DungeonEnv(TASKS[task_id], seed=seed)
-        try:
-            events = list(run_episode(env, client, lessons, config))
-        except Exception as exc:  # one bad episode must not sink the series
-            events = [{"type": "error", "message": f"{type(exc).__name__}: {exc}"}]
-        finally:
-            env.close()
-        end = next((e for e in events if e["type"] == "end"), {})
-        lessons_in = len(lessons)
-        if use_lessons and end.get("new_lessons"):
-            lessons = merge_lessons(lessons, end["new_lessons"], 12)
-        record = {"id": uuid.uuid4().hex[:12], "task": task_id, "seed": seed, "source": source,
-                  "bench": tag, "model": client.model, "use_lessons": use_lessons,
-                  "finished_at": time.time(), "outcome": end.get("outcome", "error"),
-                  "steps": end.get("steps", 0), "lessons_in": lessons_in, "events": events}
-        store.save(record, wait=True)
-        summaries.append({k: record[k] for k in ("id", "task", "seed", "use_lessons", "outcome", "steps")}
-                         | {"usage": end.get("usage", {})})
-    return summaries
-
-
-def summary_table(results: list[dict], task_ids: list[str], conditions: list[str]) -> str:
-    """Plain-text results table: wins, average steps, and tokens per level and condition."""
-    lines = [f"{'level':14}{'journal':>8}{'won':>8}{'avg steps':>11}{'tokens':>12}"]
-    for t in task_ids:
-        for cond in conditions:
-            rows = [r for r in results if r["task"] == t and r["use_lessons"] == (cond == "on")]
-            if not rows:
-                continue
-            won = sum(r["outcome"] == "success" for r in rows)
-            steps = sum(r["steps"] for r in rows) / len(rows)
-            tok = sum(r["usage"].get("prompt_tokens", 0) + r["usage"].get("completion_tokens", 0) for r in rows)
-            lines.append(f"{t:14}{cond:>8}{won:>4}/{len(rows):<3}{steps:>11.1f}{tok:>12,}")
-    return "\n".join(lines)
+"""Benchmark series: play one level across seeds, optionally carrying the journal forward."""
+
+from __future__ import annotations
+
+import threading
+import time
+import uuid
+
+from .agent import EpisodeConfig, run_episode
+from .env import TASKS, DungeonEnv
+from .journal import merge_lessons
+from .llm import ChatClient
+from .store import EpisodeStore
+
+SEED_BASE = 1000  # benchmark seeds stay clear of the low seeds visitors try in the Space
+
+
+def bench_seeds(n: int) -> list[int]:
+    return [SEED_BASE + i for i in range(n)]
+
+
+def play_series(task_id: str, seeds: list[int], use_lessons: bool, tag: str,
+                client: ChatClient, store: EpisodeStore, source: str = "bench",
+                stop: threading.Event | None = None) -> list[dict]:
+    """Play `seeds` in order on one level and save every episode; returns per-episode summaries.
+
+    Each series starts with an empty journal; with `use_lessons`, lessons written after
+    each episode are merged in before the next seed. Setting `stop` (e.g. when a spending
+    cap is hit) ends the current episode as "stopped" and skips the remaining seeds.
+    """
+    config = EpisodeConfig(use_lessons=use_lessons, reflect=use_lessons)
+    lessons: list[str] = []
+    summaries = []
+    for seed in seeds:
+        if stop is not None and stop.is_set():
+            break
+        env = DungeonEnv(TASKS[task_id], seed=seed)
+        try:
+            events = list(run_episode(env, client, lessons, config, stop))
+        except Exception as exc:  # one bad episode must not sink the series
+            events = [{"type": "error", "message": f"{type(exc).__name__}: {exc}"}]
+        finally:
+            env.close()
+        end = next((e for e in events if e["type"] == "end"), {})
+        lessons_in = len(lessons)
+        if use_lessons and end.get("new_lessons"):
+            lessons = merge_lessons(lessons, end["new_lessons"], 12)
+        record = {"id": uuid.uuid4().hex[:12], "task": task_id, "seed": seed, "source": source,
+                  "bench": tag, "model": client.model, "use_lessons": use_lessons,
+                  "finished_at": time.time(), "outcome": end.get("outcome", "error"),
+                  "steps": end.get("steps", 0), "lessons_in": lessons_in, "events": events}
+        store.save(record, wait=True)
+        summaries.append({k: record[k] for k in ("id", "task", "seed", "use_lessons", "outcome", "steps")}
+                         | {"usage": end.get("usage", {})})
+    return summaries
+
+
+class BudgetedClient:
+    """Wraps a ChatClient and stops all parallel series once shared token spend hits a cap.
+
+    Each series records only its own running total under `key` in a shared mapping
+    (a Modal Dict in production), so concurrent writers never overwrite each other; the
+    cap is checked against the sum over `all_keys` after every call.
+    """
+
+    def __init__(self, client: ChatClient, shared, key: str, all_keys: list[str],
+                 cap_tokens: int, stop: threading.Event):
+        self.client, self.shared, self.key, self.all_keys = client, shared, key, all_keys
+        self.cap_tokens, self.stop = cap_tokens, stop
+        self.used = 0
+
+    @property
+    def model(self) -> str:
+        return self.client.model
+
+    def total(self) -> int:
+        return sum(self.shared.get(k, 0) for k in self.all_keys)
+
+    def chat(self, *args, **kwargs):
+        if self.total() >= self.cap_tokens:
+            self.stop.set()
+        reply = self.client.chat(*args, **kwargs)
+        self.used += reply.usage.get("prompt_tokens", 0) + reply.usage.get("completion_tokens", 0)
+        self.shared[self.key] = self.used
+        if self.total() >= self.cap_tokens:
+            self.stop.set()
+        return reply
+
+
+def summary_table(results: list[dict], task_ids: list[str], conditions: list[str]) -> str:
+    """Plain-text results table: wins, average steps, and tokens per level and condition."""
+    lines = [f"{'level':14}{'journal':>8}{'won':>8}{'avg steps':>11}{'tokens':>12}"]
+    for t in task_ids:
+        for cond in conditions:
+            # Episodes cut off by the spending cap are neither wins nor losses.
+            rows = [r for r in results if r["task"] == t and r["use_lessons"] == (cond == "on")
+                    and r["outcome"] != "stopped"]
+            if not rows:
+                continue
+            won = sum(r["outcome"] == "success" for r in rows)
+            steps = sum(r["steps"] for r in rows) / len(rows)
+            tok = sum(r["usage"].get("prompt_tokens", 0) + r["usage"].get("completion_tokens", 0) for r in rows)
+            lines.append(f"{t:14}{cond:>8}{won:>4}/{len(rows):<3}{steps:>11.1f}{tok:>12,}")
+    return "\n".join(lines)

web/app.js

@@ -272,7 +272,8 @@ function renderEnding() {
   if (!end || !atEnd || state.replayTimer) { box.hidden = true; return; }
   const [cls, label] = OUTCOMES[end.outcome] || ["", end.outcome];
   $("ending-title").textContent = label;
-  $("ending-title").className = cls;
+  $("ending-title").className = `ending-title ${cls}`;
+  $("new-run").disabled = !!state.liveBlocked || !!state.live;
   $("ending-detail").textContent = `${end.steps} steps, ${end.usage.calls} Beam calls.`;
   box.hidden = false;
 }
@@ -436,7 +437,11 @@ $("controls").addEventListener("submit", startRun);
 $("stop").addEventListener("click", stopRun);
 $("follow").addEventListener("click", () => select(state.turns.length - 1, true));
 $("again").addEventListener("click", () => { if (state.episodeId) openReplay(state.episodeId); });
-$("new-run").addEventListener("click", () => { window.scrollTo({ top: 0 }); $("task").focus(); });
+// Start right away on the same level with the next layout, so it's a fresh dungeon.
+$("new-run").addEventListener("click", () => {
+  $("seed").value = String((Number($("seed").value) || 0) + 1);
+  $("controls").requestSubmit();
+});
 $("task").addEventListener("change", () => { if (!state.live && !startEvent()) showGoal(); });
 $("timeline").addEventListener("keydown", (e) => {
   if (e.key === "ArrowLeft") { stopReplay(); select(state.selected - 1, false); e.preventDefault(); }

web/index.html

@@ -46,7 +46,7 @@
           <p class="ending-detail" id="ending-detail"></p>
           <div class="ending-actions">
             <button class="secondary" id="again" type="button">Watch again</button>
-            <button class="primary" id="new-run" type="button">New run</button>
+            <button class="primary" id="new-run" type="button">Try a new layout</button>
           </div>
         </div>
         <p class="empty" id="empty">No runs yet. Pick a level and press <em>Watch Beam explore!</em></p>