"""The prompt-cache warmer: an idle session whose children still run re-sends
its last request with a tiny output budget before the provider cache would
expire.
Catches: pings that rebuild the request instead of repeating the one the turn
sent (which would compact or summarize), ping request rows that do not carry
the turn's prefix identity, pings that commit or publish anything, warming
that never stops at its budget, warming that continues after the waking work
finished, and a submit that cannot get through while a ping holds the session.
None of this is observable from the transcript alone; the provider and the
request rows are the witnesses.
"""
import json
import threading
import time
import unittest
from harness import Albedo, Provider, exclusive, text
# A prefix big enough to matter: 4096 cached tokens against the 1024 floor.
PARENT_USAGE = {
"input_tokens": 6000,
"output_tokens": 4,
"input_tokens_details": {"cached_tokens": 4096},
}
CHILD_USAGE = {
"input_tokens": 300,
"output_tokens": 4,
"input_tokens_details": {"cached_tokens": 100},
}
TASK = "hold the fort until released"
# A local cache-table override matching the fixture host: the shortest refresh
# TTL the table can say, bought at 1.0x and read at 0.3x, so pings come 0.9s
# apart and at most floor(1.0 / 0.3) - 1 = 2 of them pay for themselves.
TTL_SECONDS = 1
# Longer than a ping interval by more than a loaded machine's jitter: a ping
# that was due would have gone out within it.
QUIET_SECONDS = 2 * TTL_SECONDS
CACHE_TTL = {
"version": 1,
"entries": [
{
"id": "fixture",
"match": {"host": "127.0.0.1"},
"policy": "refresh",
"clock": "request",
"tiers": [{"seconds": TTL_SECONDS, "write": 1.0}],
"read": 0.3,
"evidence": "measured",
}
],
}
def wait_for(predicate, timeout=30):
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
value = predicate()
if value:
return value
time.sleep(0.05)
raise AssertionError("warming never arrived")
class WarmTest(unittest.TestCase):
def setUp(self):
# The child's turn blocks here until the test releases it, so the
# parent sits idle on work that will wake it.
self.gate = threading.Event()
self.provider = Provider(self.reply)
providers = {
"orchestrator": {
"baseUrl": self.provider.url + "/parent/v1",
"apiKey": "key",
"model": "fixture-model",
"protocol": "responses",
},
"scout": {
"baseUrl": self.provider.url + "/child/v1",
"apiKey": "key",
"model": "scout-model",
"protocol": "responses",
},
}
def prepare(app):
(app.home / "cache-ttl.json").write_text(json.dumps(CACHE_TTL))
app.write_extensions(
{"enabled": {"warm": True}, "warm": {"minCachedTokens": 1024}}
)
self.app = Albedo(
self.provider, protocol="responses", providers=providers, prepare=prepare
)
self.app.__enter__()
self.addCleanup(self.provider.close)
self.addCleanup(self.app.__exit__, None, None, None)
# A failed test must not leave the scout's request blocked.
self.addCleanup(self.gate.set)
def reply(self, request):
if TASK in json.dumps(request):
self.gate.wait(timeout=90)
return text("scout finished", usage=CHILD_USAGE)
return text("orchestrator reply", usage=PARENT_USAGE)
def api(self, path, body=None, method=None):
with self.app.api(path, body, method=method) as response:
return json.load(response)
def requests(self, fragment):
return [
record["request"]
for record in self.provider.requests
if fragment in record["path"]
]
def pings(self, session_requests):
return [
request for request in session_requests if "max_output_tokens" in request
]
def rows(self, session):
with self.app.api(f"/sessions/{session}/requests") as response:
return json.load(response)["rows"]
def tree(self, session):
items, after = [], 0
while True:
with self.app.api(
f"/sessions/{session}/tree?after={after}&limit=100"
) as response:
page = json.load(response)
items.extend(page["items"])
if not page["hasMore"]:
return items
after = page["nextCursor"]
def swarm(self, parent=None):
"""A parent whose turn has run while its scout is still working."""
parent = parent or self.app.session()
made = self.api(
f"/sessions/{parent}/children",
{"name": "scout", "task": TASK, "model": "scout/scout-model"},
)
child = made["member"]["session"]
wait_for(lambda: self.requests("/child/") or None)
self.app.prompt(parent, "run the swarm").close()
self.app.idle(parent)
return parent, child
def assert_no_pings(self, parent, child):
time.sleep(QUIET_SECONDS)
self.assertEqual(len(self.requests("/parent/")), 1)
self.assertEqual([row["kind"] for row in self.rows(parent)], ["turn"])
@exclusive
def test_the_warmer_is_off_unless_enabled(self):
settings = json.loads((self.app.home / "extensions.json").read_text())
del settings["enabled"]["warm"]
(self.app.home / "extensions.json").write_text(json.dumps(settings))
self.assert_no_pings(*self.swarm())
@exclusive
def test_a_session_that_disables_the_warmer_is_not_pinged(self):
parent = self.app.session()
extensions = self.api(
f"/sessions/{parent}/extensions", {"name": "warm", "enabled": False}
)
self.assertFalse(
next(item["enabled"] for item in extensions if item["name"] == "warm")
)
self.assert_no_pings(*self.swarm(parent))
@exclusive
def test_pings_repeat_the_last_request_and_stop_at_the_budget(self):
parent, child = self.swarm()
before = self.tree(parent)
def two_pings():
found = self.pings(self.requests("/parent/"))
return found if len(found) == 2 else None
pings = wait_for(two_pings)
turn = self.requests("/parent/")[0]
for ping in pings:
# A ping is the turn's request as sent, with only the output
# budget lowered — the Responses protocol's minimum.
self.assertEqual(ping, {**turn, "max_output_tokens": 16})
# The request rows carry the same prefix identity and cache marks as
# the turn they repeat, and never a transcript row of their own.
rows = self.rows(parent)
self.assertEqual(
[row["kind"] for row in rows], ["turn", "background", "background"]
)
head = rows[0]
for row in rows[1:]:
self.assertEqual(row["headHash"], head["headHash"])
self.assertEqual(row["inputs"], head["inputs"])
self.assertEqual(row["replaced"], head["replaced"])
self.assertEqual(row["projectionHash"], head["projectionHash"])
self.assertEqual(row["cacheMarks"], head["cacheMarks"])
self.assertEqual(row["cachedInputTokens"], 4096)
self.assertIsNone(row["seq"])
self.assertEqual(row["outcome"], "ok")
self.assertEqual(row["profile"], "orchestrator")
self.assertEqual(row["model"], "fixture-model")
# A ping paid 4096 cached tokens: the hit the TTL model predicted.
# Each ping is due an interval after the send before it, timed from
# the moment the warmer sent it, which its request row stamps a little
# later: rows apart from the turn's are bounded from the turn's.
interval = TTL_SECONDS * 900
self.assertGreaterEqual(rows[1]["startedMs"], rows[0]["startedMs"] + interval)
self.assertGreaterEqual(
rows[2]["startedMs"], rows[0]["startedMs"] + 2 * interval
)
# The transcript is untouched by warming.
self.assertEqual(self.tree(parent), before)
# The budget is spent: no third ping past another interval.
time.sleep(QUIET_SECONDS)
self.assertEqual(len(self.pings(self.requests("/parent/"))), 2)
# Once the scout finishes, nothing new is scheduled either.
self.gate.set()
self.app.idle(child)
self.app.idle(parent)
time.sleep(QUIET_SECONDS)
self.assertEqual(len(self.pings(self.requests("/parent/"))), 2)
# The scout's answer woke the parent with an ordinary turn.
turns = [
request
for request in self.requests("/parent/")
if "max_output_tokens" not in request
]
self.assertEqual(len(turns), 2)
self.assertIn("scout finished", json.dumps(turns[1]))
@exclusive
def test_a_ping_restarts_the_cached_counts_fade(self):
parent, child = self.swarm()
def fades():
return [
event["cacheFade"]
for event in self.app.events(parent)
if event["type"] == "usage"
]
# The fixture's clock counts the TTL from the send's start; nothing is
# left after it.
turn = self.rows(parent)[0]
expires = turn["startedMs"] + TTL_SECONDS * 1000
self.assertEqual(fades()[0], [{"at": expires, "cached": 0}])
# A ping resends the prefix, so the provider's clock starts over from
# it: the fade moves to the ping's send, which the session times from
# just before the request row does.
ping = wait_for(
lambda: next(
(row for row in self.rows(parent) if row["kind"] == "background"), None
)
)
moved = wait_for(
lambda: next(
(step["at"] for step in fades()[-1] if step["at"] != expires), None
)
)
self.assertLessEqual(moved, ping["startedMs"] + TTL_SECONDS * 1000)
self.assertGreater(moved, ping["startedMs"] + TTL_SECONDS * 1000 - 500)
self.gate.set()
self.app.idle(child)
self.app.idle(parent)
@exclusive
def test_a_submit_during_warming_answers_and_warming_ends_with_the_children(self):
parent, child = self.swarm()
first = wait_for(lambda: self.pings(self.requests("/parent/")) or None)
turn = self.requests("/parent/")[0]
self.assertEqual(first[0], {**turn, "max_output_tokens": 16})
# A real submit while warming gets through and is answered normally.
self.app.prompt(parent, "real work now").close()
self.app.idle(parent)
self.assertIn(
"orchestrator reply", [item["preview"] for item in self.tree(parent)]
)
# The work that kept the parent warm is done; warming ends with it.
self.gate.set()
self.app.idle(child)
self.app.idle(parent)
settled = len(self.pings(self.requests("/parent/")))
time.sleep(QUIET_SECONDS)
self.assertEqual(len(self.pings(self.requests("/parent/"))), settled)
# Every ping that did go out repeated a turn's request, budget aside.
turns = [
request
for request in self.requests("/parent/")
if "max_output_tokens" not in request
]
for ping in self.pings(self.requests("/parent/")):
self.assertIn(
ping, [{**repeated, "max_output_tokens": 16} for repeated in turns]
)