ata-coder 2.4.2__py3-none-any.whl

ata_coder/sub_agent.py

@@ -0,0 +1,273 @@
+# -*- coding: utf-8 -*-
+"""
+Sub-Agent — independent agent with isolated context window.
+
+Each SubAgent runs as an asyncio.Task and has:
+- Independent LLM client (separate httpx session)
+- Independent message history (no context leakage)
+- Independent tool executor (own file cache)
+- Cancel support via asyncio.Task.cancel()
+"""
+
+import asyncio
+import json
+import logging
+import os
+import uuid
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+from .config import AppConfig, LLMConfig
+from .llm_client import LLMClient
+from .tools import ToolExecutor, TOOL_DEFINITIONS, ToolResult
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["SubAgent", "SubAgentResult"]
+
+
+@dataclass
+class SubAgentResult:
+    """Result returned when a sub-agent completes."""
+    agent_id: str
+    result: Optional[str] = None
+    error: Optional[str] = None
+    messages: list[dict] = field(default_factory=list)
+    success: bool = True
+    tool_call_count: int = 0
+
+
+class SubAgent:
+    """
+    Independent sub-agent with isolated context window.
+
+    Usage:
+        sub = SubAgent(config=config, skill_prompt="You are a debugger...")
+        await sub.run("Find the bug in this code: ...")
+        result = await sub.wait(timeout=300)
+        print(result.result)
+    """
+
+    def __init__(
+        self,
+        config: AppConfig,
+        skill_prompt: str = "",
+        model: Optional[str] = None,
+        tools: Optional[list[dict]] = None,
+        event_callback: Optional[Callable[[str, str, ToolResult], None]] = None,
+        agent_id: Optional[str] = None,
+    ):
+        self.id = agent_id or f"sub_{uuid.uuid4().hex[:8]}"
+
+        # Independent LLM client — support both OpenAI and Anthropic formats
+        llm_config = LLMConfig(
+            api_key=config.llm.api_key,
+            base_url=config.llm.base_url,
+            model=model or config.llm.model,
+            temperature=config.llm.temperature,
+            max_tokens=config.llm.max_tokens,
+            thinking_strength=config.llm.thinking_strength,
+            use_anthropic=config.llm.use_anthropic,
+        )
+        if llm_config.use_anthropic:
+            from .anthropic_client import AnthropicClient
+            self._llm = AnthropicClient(llm_config)
+            self._use_anthropic = True
+        else:
+            self._llm = LLMClient(llm_config)
+            self._use_anthropic = False
+
+        # Independent tool executor
+        self._tools = ToolExecutor(config.agent)
+        self._tool_defs = tools or list(TOOL_DEFINITIONS)
+        self._llm.register_tools(self._tool_defs)
+
+        self._messages: list[dict] = []
+        self._result: Optional[str] = None
+        self._error: Optional[str] = None
+        self._tool_call_count = 0
+        self._status = "idle"  # idle | running | done | failed | cancelled
+        self._event_callback = event_callback
+        self._skill_prompt = skill_prompt
+        self._task: Optional[asyncio.Task] = None
+        self._done = asyncio.Event()
+
+    # ── Public API ───────────────────────────────────────────────────────
+
+    async def run(self, task: str) -> None:
+        """Start sub-agent execution as an asyncio.Task."""
+        if self._status == "running":
+            raise RuntimeError(f"Sub-agent {self.id} is already running")
+        self._done.clear()
+        self._status = "running"
+        self._task = asyncio.create_task(self._run(task))
+
+    async def wait(self, timeout: Optional[float] = 300.0) -> SubAgentResult:
+        """Wait until the sub-agent completes or times out."""
+        if self._status == "idle":
+            return SubAgentResult(agent_id=self.id, result=None,
+                                  error="Sub-agent was never started", success=False)
+        if not self._task:
+            return SubAgentResult(agent_id=self.id, result=self._result,
+                                  error=self._error, success=(self._status == "done"),
+                                  messages=list(self._messages),
+                                  tool_call_count=self._tool_call_count)
+        try:
+            await asyncio.wait_for(self._done.wait(), timeout=timeout)
+        except asyncio.TimeoutError:
+            await self.cancel()
+            return SubAgentResult(
+                agent_id=self.id,
+                result=self._result,
+                error=f"Timeout after {timeout}s",
+                success=False,
+                tool_call_count=self._tool_call_count,
+            )
+        return SubAgentResult(
+            agent_id=self.id,
+            result=self._result,
+            error=self._error,
+            messages=list(self._messages),
+            success=(self._status == "done"),
+            tool_call_count=self._tool_call_count,
+        )
+
+    async def cancel(self) -> None:
+        """Cancel the sub-agent execution."""
+        if self._task and not self._task.done():
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+        if self._status in ("running", "idle"):
+            self._status = "cancelled"
+        self._done.set()
+
+    def is_running(self) -> bool:
+        return self._status == "running"
+
+    def is_done(self) -> bool:
+        return self._status in ("done", "failed", "cancelled")
+
+    @property
+    def status(self) -> str:
+        return self._status
+
+    @property
+    def result(self) -> Optional[str]:
+        return self._result
+
+    @property
+    def messages(self) -> list[dict]:
+        return list(self._messages)
+
+    @property
+    def tool_call_count(self) -> int:
+        return self._tool_call_count
+
+    # ── Internal ─────────────────────────────────────────────────────────
+
+    async def _run(self, task: str) -> None:
+        """Internal async target."""
+        try:
+            self._messages = [
+                {"role": "system", "content": self._build_prompt()},
+                {"role": "user", "content": task},
+            ]
+            self._result = await self._loop()
+            if self._status == "running":
+                self._status = "done"
+        except asyncio.CancelledError:
+            self._status = "cancelled"
+            raise
+        except Exception as e:
+            logger.exception("Sub-agent %s failed", self.id)
+            self._error = str(e)
+            self._status = "failed"
+        finally:
+            try:
+                await self._llm.close()
+            except Exception:
+                logger.debug(
+                    "Error closing LLM for sub-agent %s", self.id, exc_info=True
+                )
+            self._done.set()
+
+    async def _loop(self) -> str:
+        """Internal tool-call loop (similar to CoderAgent.run() but simplified)."""
+        SAFETY_LIMIT = 50
+        last_text = ""
+
+        while self._status == "running":
+            if self._tool_call_count >= SAFETY_LIMIT:
+                logger.warning("Sub-agent %s reached safety limit", self.id)
+                break
+
+            # Anthropic client takes system prompt separately
+            if self._use_anthropic:
+                sys_msg = ""
+                for m in self._messages:
+                    if m.get("role") == "system":
+                        sys_msg = m.get("content", "")
+                        break
+                response = await self._llm.chat(
+                    self._messages, system_prompt=sys_msg, tools=self._tool_defs
+                )
+            else:
+                response = await self._llm.chat(self._messages, tools=self._tool_defs)
+            tool_calls = response.get("tool_calls", [])
+            text = response.get("content", "")
+
+            if text:
+                last_text = text
+
+            if not tool_calls:
+                return text or last_text or "Done."
+
+            # Execute tools serially
+            for tc in tool_calls:
+                if self._status != "running":
+                    return last_text
+                self._tool_call_count += 1
+                tool_name = tc["function"]["name"]
+                try:
+                    arguments = json.loads(tc["function"]["arguments"])
+                except json.JSONDecodeError:
+                    arguments = {}
+
+                result = await self._tools.execute(tool_name, arguments)
+
+                self._messages.append({
+                    "role": "assistant",
+                    "content": text or None,
+                    "tool_calls": [tc],
+                })
+                self._messages.append({
+                    "role": "tool",
+                    "tool_call_id": tc["id"],
+                    "content": result.to_message(),
+                })
+
+                if self._event_callback:
+                    try:
+                        self._event_callback(self.id, tool_name, result)
+                    except Exception:
+                        logger.debug(
+                            "Event callback failed for sub-agent %s", self.id,
+                            exc_info=True,
+                        )
+
+        return last_text
+
+    def _build_prompt(self) -> str:
+        """Build the sub-agent's system prompt."""
+        parts = []
+        if self._skill_prompt:
+            parts.append(self._skill_prompt)
+        parts.append(
+            "You are a sub-agent working on a delegated task. "
+            "Complete your task independently and return a clear result. "
+            "Do not ask follow-up questions — the main agent cannot see them."
+        )
+        return "\n\n".join(parts)

ata_coder/sub_agent_manager.py

@@ -0,0 +1,203 @@
+# -*- coding: utf-8 -*-
+"""
+SubAgentManager — lifecycle management for concurrent sub-agents.
+
+Handles spawning, collecting, cancelling, and listing sub-agents.
+Uses asyncio.Semaphore for concurrency limits.
+"""
+
+import asyncio
+import logging
+import uuid
+from typing import Callable, Optional
+
+from .config import AppConfig
+from .sub_agent import SubAgent, SubAgentResult
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["SubAgentManager"]
+
+
+class SubAgentManager:
+    """
+    Manages sub-agent lifecycle: spawn, collect, cancel, list.
+
+    Usage:
+        mgr = SubAgentManager(config, max_concurrent=5)
+        aid = await mgr.spawn("Search for all TODO comments", skill_prompt="...")
+        # ... do other work ...
+        result = await mgr.collect(aid, timeout=300)
+        print(result.result)
+    """
+
+    def __init__(
+        self,
+        config: AppConfig,
+        max_concurrent: int = 5,
+        default_timeout: float = 300.0,
+    ):
+        self._config = config
+        self._default_timeout = default_timeout
+        self._agents: dict[str, SubAgent] = {}
+        self._semaphore = asyncio.Semaphore(max_concurrent)
+        self._collected: set[str] = set()  # track agents already released via collect()
+
+    # ── Spawn ──────────────────────────────────────────────────────────
+
+    async def spawn(
+        self,
+        task: str,
+        skill_prompt: str = "",
+        model: Optional[str] = None,
+        tools: Optional[list[dict]] = None,
+        event_callback: Optional[Callable] = None,
+    ) -> str:
+        """
+        Spawn a sub-agent. Returns the agent_id.
+
+        Args:
+            task: The task to delegate (must be self-contained)
+            skill_prompt: Optional skill system prompt for the sub-agent
+            model: Optional model override
+            tools: Optional tool list (defaults to built-in tools)
+            event_callback: Optional callback for tool results
+
+        Returns:
+            agent_id string
+
+        Raises:
+            RuntimeError: if max concurrent agents reached
+        """
+        # Acquire semaphore slot with timeout (avoids TOCTOU + private attr access)
+        try:
+            await asyncio.wait_for(self._semaphore.acquire(), timeout=30.0)
+        except asyncio.TimeoutError:
+            running = sum(1 for a in self._agents.values() if a.is_running())
+            raise RuntimeError(
+                f"Max concurrent sub-agents reached ({running} running). "
+                f"Wait for some to complete or increase the limit."
+            )
+
+        agent_id = f"sub_{uuid.uuid4().hex[:8]}"
+        sub = SubAgent(
+            config=self._config,
+            skill_prompt=skill_prompt,
+            model=model,
+            tools=tools,
+            event_callback=event_callback,
+            agent_id=agent_id,
+        )
+        self._agents[agent_id] = sub
+
+        try:
+            await sub.run(task)
+        except Exception:
+            # Release semaphore on failure — caller may never call collect()
+            self._semaphore.release()
+            raise
+        running = sum(1 for a in self._agents.values() if a.is_running())
+        logger.info("Sub-agent spawned: %s (running=%d)", agent_id, running)
+        return agent_id
+
+    # ── Collect ────────────────────────────────────────────────────────
+
+    async def collect(self, agent_id: str, timeout: Optional[float] = None) -> SubAgentResult:
+        """Wait for and collect a sub-agent's result."""
+        if timeout is None:
+            timeout = self._default_timeout
+        sub = self._agents.get(agent_id)
+        if not sub:
+            return SubAgentResult(
+                agent_id=agent_id, result=None,
+                error=f"Unknown agent: {agent_id}", success=False,
+            )
+        result = await sub.wait(timeout=timeout)
+        # Release semaphore slot (track to avoid double-release in clear_finished)
+        self._semaphore.release()
+        self._collected.add(agent_id)
+        return result
+
+    async def collect_all(self, timeout: Optional[float] = None) -> list[SubAgentResult]:
+        """Collect results from all sub-agents."""
+        results = []
+        for aid in list(self._agents.keys()):
+            results.append(await self.collect(aid, timeout))
+        return results
+
+    # ── Cancel ─────────────────────────────────────────────────────────
+
+    async def cancel(self, agent_id: str) -> bool:
+        """Cancel a specific sub-agent."""
+        sub = self._agents.get(agent_id)
+        if sub and sub.is_running():
+            await sub.cancel()
+            return True
+        return False
+
+    async def cancel_all(self) -> None:
+        """Cancel all running sub-agents."""
+        for sub in list(self._agents.values()):
+            if sub.is_running():
+                await sub.cancel()
+
+    # ── Queries ────────────────────────────────────────────────────────
+
+    def get(self, agent_id: str) -> Optional[SubAgent]:
+        """Get a sub-agent by ID."""
+        return self._agents.get(agent_id)
+
+    def list_all(self) -> list[SubAgent]:
+        """List all sub-agents."""
+        return list(self._agents.values())
+
+    def list_active(self) -> list[SubAgent]:
+        """List only running sub-agents."""
+        return [a for a in self._agents.values() if a.is_running()]
+
+    def list_finished(self) -> list[SubAgent]:
+        """List only completed/failed/cancelled sub-agents."""
+        return [a for a in self._agents.values() if a.is_done()]
+
+    @property
+    def active_count(self) -> int:
+        return sum(1 for a in self._agents.values() if a.is_running())
+
+    @property
+    def total_count(self) -> int:
+        return len(self._agents)
+
+    # ── Cleanup ────────────────────────────────────────────────────────
+
+    def clear_finished(self) -> int:
+        """Remove finished agents from tracking. Returns count removed.
+
+        Releases one semaphore slot per removed agent that was NOT already
+        collected via :meth:`collect` — without this, slots belonging to
+        agents cleaned up without :meth:`collect` would leak permanently.
+        """
+        to_remove = [
+            aid for aid, a in self._agents.items()
+            if a.is_done()
+        ]
+        for aid in to_remove:
+            del self._agents[aid]
+            if aid not in self._collected:
+                self._semaphore.release()
+            else:
+                self._collected.discard(aid)
+        if to_remove:
+            logger.debug("Cleared %d finished sub-agent(s) (slots released)", len(to_remove))
+        return len(to_remove)
+
+    async def shutdown(self) -> None:
+        """Cancel all agents and wait for them to finish."""
+        await self.cancel_all()
+        # Wait for all tasks to finish
+        for agent in list(self._agents.values()):
+            if agent._task and not agent._task.done():
+                try:
+                    await asyncio.wait_for(agent._task, timeout=5.0)
+                except (asyncio.TimeoutError, asyncio.CancelledError):
+                    pass
+        self._agents.clear()

ata_coder/system_prompt_builder.py

@@ -0,0 +1,146 @@
+"""
+System prompt builder — composes the full system prompt for the coding agent.
+
+The prompt is structured to put the skill persona FIRST (the agent's identity),
+followed by supporting context in descending order of relevance.
+"""
+
+import logging
+import platform
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class SystemPromptBuilder:
+    """Constructs the full system prompt for the coding agent.
+
+    Structure (in order):
+      1. Skill persona (who the agent IS — dominant)
+      2. Memory (relevant context from past sessions)
+      3. Environment (OS, workspace, model — 2 lines)
+      4. Project (language, framework — only if detected)
+      5. Tools (compact name-only list)
+      6. MCP (external servers — only if connected)
+      7. Rules (5 concise operational rules)
+    """
+
+    def __init__(self, subsystems: Any, workspace_dir: str | Path,
+                 model: str = "", default_prompt: str = ""):
+        self.subsys = subsystems
+        self.workspace = str(workspace_dir)
+        self.model = model
+        self.default_prompt = default_prompt
+
+    def build(self, tool_definitions: list[dict], user_input: str = "") -> str:
+        """Build the complete system prompt string."""
+        parts: list[str] = []
+
+        # 1. Skill persona — ALWAYS first
+        parts.append(self._skill_section())
+
+        # 2. Memory — relevant past context
+        memory = self._memory_section(user_input)
+        if memory:
+            parts.append(memory)
+
+        # 3. Environment — compact
+        parts.append(self._environment_section())
+
+        # 4. Project — only if detected
+        project = self._project_section()
+        if project:
+            parts.append(project)
+
+        # 5. Tools — compact name list
+        parts.append(self._tools_section(tool_definitions))
+
+        # 6. MCP — only if connected
+        mcp = self._mcp_section()
+        if mcp:
+            parts.append(mcp)
+
+        # 7. Rules — always last
+        parts.append(self._rules_section())
+
+        return "\n\n".join(parts)
+
+    # ── Sections ──────────────────────────────────────────────────────────
+
+    def _skill_section(self) -> str:
+        """The agent's persona — from active skills or default prompt."""
+        # Primary: ExtensionManager aggregates all active skill extensions
+        if self.subsys.has_extensions and self.subsys.extensions:
+            aggregated = self.subsys.extensions.aggregate_prompts(base_prompt="")
+            if aggregated.strip():
+                return aggregated
+
+        # Fallback: direct skill manager
+        if self.subsys.has_skills:
+            try:
+                prompt = self.subsys.skills.get_system_prompt()
+                if prompt.strip():
+                    return prompt
+            except Exception:
+                pass
+
+        return self.default_prompt
+
+    def _memory_section(self, user_input: str = "") -> str:
+        """Targeted memory recall (max 5, score >= 3)."""
+        if not self.subsys.has_memory:
+            return ""
+        mem = self.subsys.memory
+        if user_input.strip():
+            ctx = mem.recall_context(user_input, max_memories=5)
+        else:
+            ctx = mem.get_memory_context(max_total=5)
+        return ctx if ctx else ""
+
+    def _environment_section(self) -> str:
+        """Single-line environment summary."""
+        model_part = f" | Model: {self.model}" if self.model else ""
+        return (
+            f"OS: {platform.system()} {platform.release()}  "
+            f"| Python: {platform.python_version()}  "
+            f"| Workspace: {self.workspace}"
+            f"{model_part}"
+        )
+
+    def _project_section(self) -> str:
+        if self.subsys.has_project_info and self.subsys.project_info:
+            return self.subsys.project_info.to_prompt()
+        return ""
+
+    def _tools_section(self, tool_definitions: list[dict]) -> str:
+        """Compact tool list — names only, no descriptions."""
+        names: list[str] = []
+        for t in tool_definitions:
+            fn = t.get("function", t)
+            name = fn.get("name", "")
+            if name:
+                names.append(name)
+        return f"Tools ({len(names)}): " + ", ".join(sorted(names))
+
+    def _mcp_section(self) -> str:
+        if not self.subsys.has_mcp:
+            return ""
+        mcp = self.subsys.mcp
+        if not mcp.connected_servers:
+            return ""
+        servers = ", ".join(mcp.connected_servers)
+        return f"MCP servers ({mcp.tool_count} tools): {servers}"
+
+    def _rules_section(self) -> str:
+        """Concise operational rules — replaces the old verbose ops + formatting sections."""
+        return (
+            "Rules:\n"
+            "- Shell cwd is already the workspace. Use compound commands (cd X && do_Y) if needed.\n"
+            "- Shell command blocked? Use: python -c \"import subprocess; subprocess.run([...], cwd='...')\"\n"
+            "- Read error messages carefully — diagnose before retrying. Auto-correction handles common failures.\n"
+            "- File reads are cached — re-reading the same file returns a [cached] note. Use offset/limit to page.\n"
+            "- Context auto-compacts at ~200k tokens. When you see the compaction notice, trust the summary.\n"
+            "- Respond clearly: use **bold** for key terms, ## headings for structure, code blocks with language tags."
+        )
+