99 lines
3.8 KiB
Python
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()
|