This commit is contained in:
team3
2026-07-01 22:01:32 +02:00
parent fa718b7d6c
commit b5398f73d2
17 changed files with 1580 additions and 522 deletions

View File

@@ -5,6 +5,7 @@ respective provider fails — the other keeps running unchanged.
"""
import asyncio
import heapq
import logging
import os
import re
@@ -22,6 +23,17 @@ from config import (PROVIDERS, DEFAULT_PROVIDER, MAX_CONCURRENT_AGENTS,
log = logging.getLogger("creator.agents")
_active_processes: dict[str, asyncio.subprocess.Process] = {}
_active_started: dict[str, float] = {} # agent_key → wall-clock start (for the live runtime display)
def active_agents(scope_prefix: str | None = None) -> list[dict]:
"""Currently running agents and how long they've been running. Filter by key prefix
(e.g. f"blocks-{topic}-") for one topic. → [{key, runtime}] sorted longest-first."""
now = time.time()
out = [{"key": k, "runtime": round(now - t, 1)}
for k, t in list(_active_started.items())
if k in _active_processes and (not scope_prefix or k.startswith(scope_prefix))]
return sorted(out, key=lambda a: -a["runtime"])
# Cancelled scopes (key prefixes, symmetric to kill_process). An agent whose
# key starts with one of these prefixes aborts BEFORE the spawn — so agents WAITING
@@ -43,26 +55,71 @@ def _scope_cancelled(agent_key: str) -> bool:
# Caps the real CLI processes — independent of the pipeline semaphore in
# generator.py. The acquire happens BEFORE the spawn so that queue wait time
# does not count against the agent timeout.
_batch_sem = asyncio.Semaphore(MAX_CONCURRENT_AGENTS)
class _PrioritySemaphore:
"""asyncio.Semaphore variant: when slots are scarce, the LOWEST priority number is served first
(FIFO within the same priority). Lets earlier pipeline columns grab agents before later ones."""
def __init__(self, value: int):
self._value = value
self._waiters: list = [] # heap of [priority, seq, future]
self._seq = 0
async def acquire(self, priority: int = 100):
if self._value > 0:
self._value -= 1
return
fut = asyncio.get_event_loop().create_future()
entry = [priority, self._seq, fut]
self._seq += 1
heapq.heappush(self._waiters, entry)
try:
await fut # release() hands us the slot directly (no value change)
except BaseException:
entry[2] = None # tombstone so release() skips this dead waiter
if fut.done() and not fut.cancelled():
self.release() # granted just before we were cancelled → pass it on
raise
def release(self):
while self._waiters:
entry = heapq.heappop(self._waiters)
if entry[2] is not None and not entry[2].done():
entry[2].set_result(None) # hand the slot straight to the highest-priority waiter
return
self._value += 1
_batch_sem = _PrioritySemaphore(MAX_CONCURRENT_AGENTS)
_interactive_sem = asyncio.Semaphore(MAX_CONCURRENT_INTERACTIVE)
# Per-topic caps (lazily created): each topic gets its own batch semaphore of size
# MAX_CONCURRENT_AGENTS_PER_TOPIC, nested INSIDE the global _batch_sem.
_topic_sems: dict[str, asyncio.Semaphore] = {}
# Per-topic caps (lazily created): each topic gets its own priority semaphore of size
# MAX_CONCURRENT_AGENTS_PER_TOPIC, nested INSIDE the global _batch_sem. Priority-based too, so the
# per-topic queue can't undo the global priority when one topic is the only load.
_topic_sems: dict[str, _PrioritySemaphore] = {}
# Earlier kanban columns get the scarce global slot first (smaller = higher priority).
_STAGE_PRIORITY = ("research", "verify", "naming", "small", "dep")
def _agent_priority(key: str) -> int:
for i, tag in enumerate(_STAGE_PRIORITY):
if f"-{tag}-" in key or key.endswith(f"-{tag}"):
return i
return len(_STAGE_PRIORITY) # downstream agents (subblocks/facts/…) after the inventory columns
@asynccontextmanager
async def _batch_gate(scope: str | None):
"""Acquire a batch slot: per-topic semaphore FIRST, then the global one. The order matters —
a waiter holds only its (per-topic) slot while queueing for the global cap, so a saturated topic
never blocks other topics on the global semaphore. scope=None → global cap only."""
topic_sem = _topic_sems.setdefault(scope, asyncio.Semaphore(MAX_CONCURRENT_AGENTS_PER_TOPIC)) if scope else None
if topic_sem is None:
async with _batch_sem:
yield
else:
async with topic_sem:
async with _batch_sem:
yield
async def _batch_gate(scope: str | None, priority: int):
"""Per-topic slot FIRST (fair), then the GLOBAL slot by priority (earlier columns win when
agents are scarce). Order matters — a waiter holds only its per-topic slot while queueing globally."""
topic_sem = _topic_sems.setdefault(scope, _PrioritySemaphore(MAX_CONCURRENT_AGENTS_PER_TOPIC)) if scope else None
if topic_sem is not None:
await topic_sem.acquire(priority)
await _batch_sem.acquire(priority)
try:
yield
finally:
_batch_sem.release()
if topic_sem is not None:
topic_sem.release()
# Serialize OpenCode starts: processes starting simultaneously collide on the
# internal session DB ("database is locked", exit after <1s). The short
@@ -122,6 +179,7 @@ def kill_process(agent_key_prefix: str) -> None:
for key, process in list(_active_processes.items()):
if process.returncode is not None: # clean up dead entries while iterating
_active_processes.pop(key, None)
_active_started.pop(key, None)
continue
if key.startswith(agent_key_prefix):
log.debug("kill agent %s", key)
@@ -137,6 +195,7 @@ async def run_agent(
capabilities: str = "none",
lane: str = "batch",
scope: str | None = None,
on_line=None,
) -> tuple[int, str, str]:
if _scope_cancelled(agent_key): # before queueing: don't even enter the queue
return 1, "", "cancelled"
@@ -144,16 +203,16 @@ async def run_agent(
return 1, "", f"Unknown provider: {provider}"
if shutil.which(PROVIDERS[provider]["cli"]) is None:
return 1, "", f"CLI '{PROVIDERS[provider]['cli']}' not installed (provider: {provider})"
gate = _interactive_sem if lane == "interactive" else _batch_gate(scope)
gate = _interactive_sem if lane == "interactive" else _batch_gate(scope, _agent_priority(agent_key))
async with gate:
if _scope_cancelled(agent_key): # after the acquire: cancelled in the queue → no spawn
return 1, "", "cancelled"
if PROVIDERS[provider]["cli"] == "opencode":
return await _run_opencode(agent_key, prompt, timeout, provider, role, capabilities)
return await _run_opencode(agent_key, prompt, timeout, provider, role, capabilities, on_line=on_line)
return await _run_claude_cli(agent_key, prompt, timeout, role, capabilities)
async def _communicate(agent_key: str, cmd: list[str], stdin_data: bytes | None, timeout: int, stagger: bool = False) -> tuple[int, str, str]:
async def _communicate(agent_key: str, cmd: list[str], stdin_data: bytes | None, timeout: int, stagger: bool = False, on_line=None) -> tuple[int, str, str]:
start = time.monotonic()
async def spawn():
@@ -172,12 +231,29 @@ async def _communicate(agent_key: str, cmd: list[str], stdin_data: bytes | None,
else:
process = await spawn()
_active_processes[agent_key] = process
_active_started[agent_key] = time.time()
try:
try:
stdout, stderr = await asyncio.wait_for(
process.communicate(input=stdin_data),
timeout=timeout,
)
if on_line is not None:
# Streaming path: read stdout line by line, hand each raw line to on_line LIVE.
out_chunks: list[str] = []
async def _pump():
async for raw in process.stdout:
s = raw.decode("utf-8", errors="replace")
out_chunks.append(s)
try:
on_line(s)
except Exception:
log.debug("on_line callback failed", exc_info=True)
await asyncio.wait_for(_pump(), timeout=timeout)
await process.wait()
stderr_b = await process.stderr.read()
stdout, stderr = "".join(out_chunks).encode("utf-8"), stderr_b
else:
stdout, stderr = await asyncio.wait_for(
process.communicate(input=stdin_data),
timeout=timeout,
)
except asyncio.TimeoutError:
_kill(process)
try:
@@ -196,6 +272,7 @@ async def _communicate(agent_key: str, cmd: list[str], stdin_data: bytes | None,
# the NEW process from tracking.
if _active_processes.get(agent_key) is process:
del _active_processes[agent_key]
_active_started.pop(agent_key, None)
async def _run_claude_cli(agent_key: str, prompt: str, timeout: int, role: str, capabilities: str) -> tuple[int, str, str]:
@@ -208,7 +285,7 @@ async def _run_claude_cli(agent_key: str, prompt: str, timeout: int, role: str,
return await _communicate(agent_key, cmd, prompt.encode("utf-8"), timeout)
async def _run_opencode(agent_key: str, prompt: str, timeout: int, provider: str, role: str, capabilities: str) -> tuple[int, str, str]:
async def _run_opencode(agent_key: str, prompt: str, timeout: int, provider: str, role: str, capabilities: str, on_line=None) -> tuple[int, str, str]:
cfg = PROVIDERS[provider]
# Prompt via temp file instead of argv (ARG_MAX protection for large project prompts)
with tempfile.NamedTemporaryFile("w", suffix=".md", delete=False, encoding="utf-8", dir=tempfile.gettempdir()) as f:
@@ -224,9 +301,11 @@ async def _run_opencode(agent_key: str, prompt: str, timeout: int, provider: str
"--dangerously-skip-permissions",
"-f", str(prompt_path),
]
if on_line is not None:
cmd += ["--format", "json"] # raw JSON events → parsed live by on_line
try:
rc, stdout, stderr = await _communicate(agent_key, cmd, None, timeout, stagger=True)
return rc, _clean_opencode_output(stdout), stderr
rc, stdout, stderr = await _communicate(agent_key, cmd, None, timeout, stagger=True, on_line=on_line)
return rc, (stdout if on_line is not None else _clean_opencode_output(stdout)), stderr
finally:
prompt_path.unlink(missing_ok=True)