Wire specialist waves, Brain consolidation, and Classic-compatible reports so Agent mode can run end-to-end via OpenRouter without changing the default Classic path. Co-authored-by: Cursor <cursoragent@cursor.com>
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
"""Shared contracts and job-local accounting for Agent-mode workers."""
|
|
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Optional, Protocol
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AgentScope:
|
|
"""A bounded work package passed to exactly one specialist agent."""
|
|
|
|
scope_id: str
|
|
payload: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class AgentResult:
|
|
"""Artifacts returned by a specialist for collection by the orchestrator."""
|
|
|
|
scope_id: str
|
|
artifacts: List[Dict[str, Any]] = field(default_factory=list)
|
|
error: str = ""
|
|
|
|
|
|
@dataclass
|
|
class AgentUsage:
|
|
"""Thread-safe usage accounting owned by one Agent pipeline run."""
|
|
|
|
usd: float = 0.0
|
|
calls: int = 0
|
|
cached: int = 0
|
|
by_stage: Dict[str, Dict[str, Any]] = field(default_factory=dict)
|
|
models: Dict[str, set] = field(default_factory=lambda: {
|
|
"vision": set(),
|
|
"text_cloud": set(),
|
|
})
|
|
_lock: threading.Lock = field(default_factory=threading.Lock, repr=False)
|
|
|
|
def record(
|
|
self,
|
|
stage: str,
|
|
model: str,
|
|
usd: float = 0.0,
|
|
cached: bool = False,
|
|
has_images: bool = False,
|
|
) -> None:
|
|
with self._lock:
|
|
bucket = self.by_stage.setdefault(
|
|
stage, {"usd": 0.0, "calls": 0, "cached": 0}
|
|
)
|
|
if cached:
|
|
self.cached += 1
|
|
bucket["cached"] += 1
|
|
else:
|
|
self.calls += 1
|
|
self.usd += usd
|
|
bucket["calls"] += 1
|
|
bucket["usd"] += usd
|
|
family = "vision" if has_images else "text_cloud"
|
|
self.models[family].add(model)
|
|
|
|
def snapshot(self) -> Dict[str, Any]:
|
|
with self._lock:
|
|
return {
|
|
"usd": self.usd,
|
|
"calls": self.calls,
|
|
"cached": self.cached,
|
|
"by_stage": {k: dict(v) for k, v in self.by_stage.items()},
|
|
"models": {
|
|
"vision": sorted(self.models["vision"]),
|
|
"text_local": [],
|
|
"text_cloud": sorted(self.models["text_cloud"]),
|
|
"fallback_count": 0,
|
|
},
|
|
}
|
|
|
|
|
|
class ScopedAgent(Protocol):
|
|
"""Protocol implemented by each future specialist agent."""
|
|
|
|
name: str
|
|
|
|
def run(self, scope: AgentScope) -> AgentResult:
|
|
...
|
|
|
|
|
|
def failure(scope: AgentScope, error: Exception) -> AgentResult:
|
|
"""Convert a worker exception into a non-fatal scoped result."""
|
|
return AgentResult(scope_id=scope.scope_id, error=str(error))
|