"""Concurrent child kernel boot stays responsive and forwards all answers."""
import json
import os
import time
import unittest
import urllib.error
from harness import Albedo, Provider, python, text
HEAVY = os.environ.get("ALBEDO_E2E_HEAVY") == "1"
ROOTS = 4 if HEAVY else 2
CHILDREN = 12 if HEAVY else 4
def user_text(request):
return next(
(
item.get("content", "")
for item in reversed(request["messages"])
if item.get("role") == "user"
),
"",
)
class AgentsSwarmTests(unittest.TestCase):
def test_children_boot_without_blocking_actors_or_losing_answers(self):
def script(request):
last = request["messages"][-1]
if last.get("role") == "user" and 'kind="task"' in user_text(request):
return python("print('cell ran')")
return text("done")
provider = Provider(script)
try:
with Albedo(provider) as app:
roots = [app.session() for _ in range(ROOTS)]
children = []
started = time.monotonic()
for root in roots:
for number in range(CHILDREN):
with app.api(
f"/sessions/{root}/children",
{"name": f"sp-{number}", "task": f"count to {number}"},
) as response:
children.append(json.load(response)["session"]["id"])
self.assertLess(
time.monotonic() - started, 30, "spawn must not await kernel boot"
)
slowest = 0.0
deadline = time.monotonic() + 300
while time.monotonic() < deadline:
answers = self.answers(app, roots)
tick = time.monotonic()
for root in roots:
with app.api(f"/agents?session={root}") as response:
self.assertEqual(
len(json.load(response)["nodes"]), CHILDREN + 1
)
for child in children[::6]:
with app.api(f"/sessions/{child}/status") as response:
json.load(response)
slowest = max(slowest, time.monotonic() - tick)
if sum(answers) == ROOTS * CHILDREN:
break
time.sleep(0.5)
self.assertEqual(self.answers(app, roots), [CHILDREN] * ROOTS)
self.assertLess(
slowest, 5, "status and tree calls must stay responsive"
)
# Fast kernels can finish before the first status poll. Their
# completed Python outputs prove every child actually booted.
completed = [
record
for record in provider.requests
if record["request"]["messages"][-1].get("role") == "tool"
and "cell ran"
in record["request"]["messages"][-1].get("content", "")
]
self.assertEqual(len(completed), ROOTS * CHILDREN)
with app.api("/health") as response:
self.assertTrue(json.load(response)["ok"])
finally:
provider.close()
def test_explicit_session_provider_binds_and_unknown_profile_is_rejected(self):
provider = Provider(lambda _request: text("ok"))
try:
with Albedo(provider) as app:
session = app.session()
with app.api("/sessions") as response:
info = next(
item for item in json.load(response) if item["id"] == session
)
self.assertEqual(info["provider"], app.profile)
with self.assertRaises(urllib.error.HTTPError) as rejected:
app.api(
"/sessions",
{"workspace": str(app.workspace), "provider": "not-configured"},
).close()
self.assertEqual(rejected.exception.code, 400)
finally:
provider.close()
@staticmethod
def answers(app, roots):
counts = []
for root in roots:
with app.api(f"/sessions/{root}/preview?limit=200") as response:
items = json.load(response)["items"]
counts.append(
sum(
item["type"] == "user" and "unreviewed" in item["preview"]
for item in items
)
)
return counts