update
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user