Files
tai/tai_pkg/agents.py
T
retoor 8df5cf7a45 Split monolithic tai.py into tai_pkg package, keep tai-orig.py as original
tai.py (8643 lines) is now a thin facade over tai_pkg/ (33 feature modules, one class or feature per file). tai-orig.py preserves the pre-split single-file version byte for byte (sha256 verified against snapshot); import tai and ./tai.py behavior unchanged.

Equivalence proof, all run on this tree immediately pre-commit: differential battery tai-orig.py vs facade 3718 checks 0 failures (surface, pure fuzz, tool dispatch, store/audit, release ops, fake-HOME install roundtrip, CLI runs); test_tai.py 261 OK; test_seal.py 16 OK; --version/--help byte-identical; dir(tai) zero removals.

Fix included: opencode wire-protocol schema echo leaked into replies; opencode_parse now scans with JSONDecoder.raw_decode capped at 200 scans (13 OpencodeParseTests). Test count 261 = 237 HEAD + 8 pre-existing uncommitted ToolArgumentsTests (now committed with this change, disclosed) + 13 OpencodeParseTests + 3 EmbeddedStringsTests. Patch-target-only test migration (24 hunks), zero assertion changes. tai_pkg clean-regenerated after mid-verification tamper was found and reverted.
2026-10-10 19:25:02 +02:00

99 lines
3.8 KiB
Python

"""Background subagent registry and lifecycle."""
import threading
import time
from tai_pkg.consts import Consts
from tai_pkg.store import Store
from tai_pkg.text import Text
__all__ = ["AgentRegistry", "REGISTRY", "default_runner"]
class AgentRegistry:
"""Background subagent registry."""
def __init__(self):
self.agents = {}
self.lock = threading.Lock()
self.next_id = [1]
def spawn_agent(self, task, profile, timeout, config, seal, depth=0, runner=None):
with self.lock:
agent_id = self.next_id[0]
self.next_id[0] += 1
record = {"id": agent_id, "task": task, "profile": profile, "status": "running", "result": "", "started": time.time(), "ended": None, "thread": None, "spilled": None}
self.agents[agent_id] = record
def target():
try:
if runner is None:
result = default_runner(config, seal, profile, depth, task, timeout)
else:
result = runner(task, profile, timeout)
status = "timeout" if "[time limit reached]" in result else "done"
except Exception as exc:
result = "worker failed: " + Text.short_error(exc)
status = "error"
with self.lock:
record["result"] = result
record["status"] = status
record["ended"] = time.time()
worker_thread = threading.Thread(target=target, daemon=True)
with self.lock:
record["thread"] = worker_thread
worker_thread.start()
return agent_id
def poll_agent(self, agent_id, wait=0, profile=None):
with self.lock:
record = self.agents.get(agent_id)
if profile is not None and record is not None and record.get("profile") != profile:
record = None
worker_thread = record["thread"] if record else None
if record is None:
return "missing", "no agent %d" % agent_id
if worker_thread is not None:
worker_thread.join(timeout=max(0, wait))
with self.lock:
status = record["status"]
result = record["result"]
elapsed = (record["ended"] or time.time()) - record["started"]
if status == "running":
return status, "agent %d still running (%ds elapsed)" % (agent_id, elapsed)
if status == "error":
return status, result
return status, Text.truncate(result, 6000, 2000) or "(empty result)"
def list_agents(self, profile=None):
with self.lock:
return [{"id": record["id"], "task": record["task"], "status": record["status"], "elapsed": int((record["ended"] or time.time()) - record["started"])} for record in self.agents.values() if profile is None or record.get("profile") == profile]
def clear_agents(self, profile=None):
with self.lock:
finished = [key for key, record in self.agents.items() if record["status"] != "running" and (profile is None or record.get("profile") == profile)]
for key in finished:
del self.agents[key]
return len(finished)
def running_agents(self):
with self.lock:
return sum(1 for record in self.agents.values() if record["status"] == "running")
REGISTRY = AgentRegistry()
def default_runner(config, seal, profile, depth, task, timeout, max_steps=Consts.WORKER_STEPS):
from tai_pkg.agent import Agent # deferred: breaks tools/agents/agent cycle
store = Store(config, seal)
try:
worker = Agent(config, store, persist=False, quiet=True, depth=depth)
if profile != worker.profile:
worker.switch_profile(profile, silent=True)
worker.deadline = time.time() + timeout
return worker.run_turn(task, capture=True, max_steps=max_steps)
finally:
store.close()