ata-coder 2.4.2__py3-none-any.whl

ata_coder/agent_compact.py

@@ -0,0 +1,190 @@
+"""Context compaction and token budget management — mixin for CoderAgent."""
+import copy
+import json
+import logging
+
+from .types import Message
+from .clawd_integration import get_clawd
+from .model_router import get_subagent_model
+
+logger = logging.getLogger(__name__)
+
+
+class CompactionMixin:
+    """Context window compaction and token budget management."""
+
+    # ── Compaction token budget ──────────────────────────────────────────
+    RECENT_TOKEN_BUDGET = 80_000   # max tokens to keep from recent messages
+    COMPACT_IF_FEWER_THAN = 6      # skip compaction if fewer than this many msgs
+
+    def _force_truncate(self) -> None:
+        """Drop the oldest non-system messages when we exceed 95% of max tokens.
+
+        Called only as a last resort after compaction has already run.
+        Keeps the system prompt, the last 4 messages, and everything in between
+        gets replaced with a summary marker.
+        """
+        if len(self._state.messages) <= 6:
+            return
+        # Preserve system message if present; otherwise keep first message
+        first = self._state.messages[0]
+        is_system = isinstance(first, dict) and first.get("role") == "system"
+        keep_tail = self._state.messages[-4:]  # last 2 exchanges
+        truncated = ([first] if is_system else []) + [
+            {"role": "user", "content": "[Conversation truncated — token limit reached]"},
+            {"role": "assistant", "content": "Understood. I'll continue with the most recent context."},
+        ] + keep_tail
+        old = len(self._state.messages)
+        self._state.messages = truncated
+        self._cached_system_prompt = None
+        logger.warning("Force-truncated: %d → %d messages", old, len(truncated))
+
+    # ── Compaction token budget ──────────────────────────────────────────
+
+    async def compact(self) -> str:
+        """
+        Compact conversation by summarising old messages.
+
+        Strategy: keep system prompt + recent messages up to
+        RECENT_TOKEN_BUDGET tokens, summarise everything in between using
+        a cheap LLM call.  Falls back to a lightweight extractive summary
+        if the API call fails.
+
+        Uses token budget for recent messages (not a fixed count) so that
+        large file reads don't survive compaction intact.
+        """
+        if len(self._state.messages) <= self.COMPACT_IF_FEWER_THAN:
+            return "Already compact."
+
+        # Clawd: PreCompact
+        get_clawd().compact()
+
+        first = self._state.messages[0]
+        has_system = isinstance(first, dict) and first.get("role") == "system"
+        all_but_system = self._state.messages[1:] if has_system else self._state.messages[:]
+
+        # Walk backwards through recent messages, accumulating up to the budget
+        recent: list[Message] = []
+        recent_tokens = 0
+        for msg in reversed(all_but_system):
+            msg_tokens = self._estimate_message_tokens(msg)
+            if recent_tokens + msg_tokens > self.RECENT_TOKEN_BUDGET and recent:
+                # Stop — we've filled the recent budget
+                break
+            recent.insert(0, msg)
+            recent_tokens += msg_tokens
+
+        # The middle is everything NOT in recent and NOT the system msg
+        kept_count = len(recent)
+        middle = all_but_system[:-kept_count] if kept_count > 0 else all_but_system
+
+        if not middle:
+            return "Already compact (all messages fit in recent budget)."
+
+        # Extract key facts from middle messages for the fallback summary
+        tool_count = sum(1 for m in middle if m.get("tool_calls"))
+        user_msgs = [m.get("content", "")[:200] for m in middle if m.get("role") == "user"]
+        file_ops = self._collect_file_ops(middle)
+
+        summary = await self._summarise_messages(middle, file_ops, user_msgs, tool_count)
+
+        truncated: list[Message] = (
+            [first] if has_system else []
+        ) + [
+            {"role": "user", "content": "[Conversation summary]\\n" + summary},
+            {"role": "assistant", "content": "Understood. I'll continue with the remaining context using the summary above."},
+        ]
+        truncated.extend(recent)
+        old_count = len(self._state.messages)
+        old_tokens = self.get_token_estimate()
+        self._state.messages = truncated
+        self._cached_system_prompt = None  # messages[0] changed — invalidate
+        new_tokens = self.get_token_estimate()
+
+        logger.info("Compacted: %d→%d msgs, ~%d→%d tokens (files: %d, tools: %d, recent_budget: %d)",
+                    old_count, len(truncated), old_tokens, new_tokens,
+                    len(file_ops), tool_count, recent_tokens)
+        return (f"Compacted from {old_count}→{len(truncated)} messages "
+                f"(~{old_tokens:,}→~{new_tokens:,} tokens, {len(file_ops)} files, {tool_count} tool calls).")
+
+    def _estimate_message_tokens(self, msg: Message) -> int:
+        """Rough token estimate for a single message."""
+        content = msg.get("content", "") or ""
+        # Use the LLM client's estimator if available
+        try:
+            return self.llm.count_tokens_approx([msg])
+        except Exception:
+            # CJK-aware fallback
+            import re
+            cjk = len(re.findall(r'[一-鿿　-〿＀-￯]', content))
+            other = len(content) - cjk
+            tokens = (cjk * 2 // 3) + (other // 4)
+            for tc in msg.get("tool_calls", []):
+                try:
+                    tokens += len(json.dumps(tc)) // 4
+                except Exception:
+                    pass
+            return max(1, tokens)
+
+    def _collect_file_ops(self, messages: list[Message]) -> list[str]:
+        """Collect files modified in a message list."""
+        ops: list[str] = []
+        for m in messages:
+            for tc in m.get("tool_calls", []):
+                fn = tc.get("function", {})
+                if fn.get("name") in ("write_file", "edit_file"):
+                    try:
+                        args = json.loads(fn.get("arguments", "{}"))
+                        fp = args.get("file_path", "")
+                        if fp:
+                            ops.append(fp)
+                    except json.JSONDecodeError:
+                        pass
+        return ops
+
+    async def _summarise_messages(self, middle: list[Message], file_ops: list[str],
+                            user_msgs: list[str], tool_count: int) -> str:
+        """Generate a summary of the middle conversation segment.
+
+        Attempts a cheap LLM call first; falls back to a lightweight extractive
+        summary so the user never loses context entirely.
+        """
+        # ── LLM-based summary (best effort) ──────────────────────────────
+        try:
+            summary_prompt = (
+                "Summarise this conversation segment in 3-5 bullet points. "
+                "Focus on: what the user asked, what files were changed, what "
+                "decisions were made, and any unresolved issues. "
+                "Be concise — this summary will replace the full conversation "
+                "history to save context tokens.\n\n"
+                f"Files modified: {', '.join(file_ops) if file_ops else 'none'}\n"
+                f"Tool calls: {tool_count}\n"
+                f"User requests: {'; '.join(user_msgs[:5])}\n"
+            )
+            # Use a separate cheap-model client for summarisation
+            from .llm_client import LLMClient
+            summary_config = copy.deepcopy(self.llm.config)
+            summary_config.model = get_subagent_model()
+            sc = LLMClient(summary_config)
+            try:
+                resp = await sc.chat([{"role": "user", "content": summary_prompt}], tools=[])
+                llm_summary = (resp.get("content") or "").strip()
+                if llm_summary:
+                    # Include extractive data as context for the LLM summary
+                    parts = [llm_summary]
+                    if file_ops:
+                        parts.append(f"\nFiles touched: {', '.join(file_ops[:10])}")
+                    return "\n".join(parts)
+            finally:
+                await sc.close()
+        except Exception:
+            logger.debug("LLM summarisation unavailable, using extractive fallback")
+
+        # ── Extractive fallback ─────────────────────────────────────────
+        parts = [f"Summarised {len(middle)} messages ({tool_count} tool calls)."]
+        if user_msgs:
+            parts.append(f"Topics: {'; '.join(user_msgs[:5])}")
+        if file_ops:
+            parts.append(f"Files modified: {', '.join(file_ops[:10])}")
+        return "\n".join(parts)
+

ata_coder/agent_controller.py

@@ -0,0 +1,218 @@
+# -*- coding: utf-8 -*-
+"""
+Agent Controller — runs CoderAgent as an asyncio task.
+
+With the async event loop, the agent runs as a coroutine on the same
+thread as the UI. asyncio.Task provides built-in cancellation, and
+asyncio.Event replaces threading.Event for coordination.
+
+No more background threads, heartbeat pumpers, or thread supervisors.
+"""
+
+import asyncio
+import logging
+from typing import Any, Optional
+
+from .agent import CoderAgent, CompleteEvent, ErrorEvent
+from .config import AppConfig
+from .agent_subsystems import AgentSubsystems
+from .core.queue import EventQueue
+from .tools import ToolExecutor
+from .sub_agent_manager import SubAgentManager
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["AgentController"]
+
+
+class AgentController:
+    """
+    Wraps CoderAgent for asyncio task execution.
+
+    Owns:
+    - The CoderAgent instance (runs as an asyncio.Task)
+    - Event queue (agent → UI communication)
+    - SubAgentManager (for parallel sub-agent execution)
+
+    Usage:
+        controller = AgentController(config, subsystems, tool_exec)
+        await controller.start()
+        await controller.submit("write a hello world script")
+        async for event in controller.event_queue:
+            ui.on_event(event)
+        await controller.shutdown()
+    """
+
+    def __init__(
+        self,
+        config: AppConfig | None = None,
+        subsystems: AgentSubsystems | None = None,
+        tool_executor: ToolExecutor | None = None,
+    ):
+        self._config = config or AppConfig.load()
+        self._subsystems = subsystems or AgentSubsystems()
+        self._tool_exec = tool_executor or ToolExecutor(self._config.agent)
+
+        # Async event queue (agent → UI)
+        self.event_queue = EventQueue(maxsize=5000)
+
+        # Async coordination primitives
+        self._cancel = asyncio.Event()
+        self._busy = asyncio.Event()
+
+        # Sub-agent manager (created on start)
+        self._sub_agent_mgr = None
+
+        # Agent and its task
+        self._agent: Optional[CoderAgent] = None
+        self._agent_task: Optional[asyncio.Task] = None
+
+    # ── Lifecycle ──────────────────────────────────────────────────────────
+
+    async def start(self) -> None:
+        """Create the agent and prepare for task submission."""
+        if self._agent_task and not self._agent_task.done():
+            return
+
+        # Create SubAgentManager
+        max_sub = getattr(self._config.agent, "max_sub_agents", 5)
+        sub_timeout = getattr(self._config.agent, "sub_agent_timeout", 300.0)
+        self._sub_agent_mgr = SubAgentManager(
+            self._config,
+            max_concurrent=max_sub,
+            default_timeout=sub_timeout,
+        )
+
+        # Create agent
+        self._agent = CoderAgent(
+            config=self._config,
+            tool_executor=self._tool_exec,
+            subsystems=self._subsystems,
+        )
+        # Wire event queue and sub-agent manager
+        self._agent._event_queue = self.event_queue
+        self._agent.set_sub_agent_manager(self._sub_agent_mgr)
+        self._tool_exec.set_sub_agent_manager(self._sub_agent_mgr)
+
+        self._cancel.clear()
+        self._busy.clear()
+
+        logger.info("AgentController started (async)")
+
+    async def shutdown(self) -> None:
+        """Stop the agent and cleanup."""
+        self._cancel.set()
+        # Cancel agent task
+        if self._agent_task and not self._agent_task.done():
+            self._agent_task.cancel()
+            try:
+                await self._agent_task
+            except asyncio.CancelledError:
+                pass
+            except Exception:
+                logger.exception("Agent task shutdown error")
+            self._agent_task = None
+
+        # Cancel all sub-agents
+        if self._sub_agent_mgr:
+            await self._sub_agent_mgr.shutdown()
+            self._sub_agent_mgr = None
+
+        if self._agent:
+            try:
+                await self._agent.shutdown()
+            except Exception:
+                logger.exception("Agent shutdown error")
+            self._agent = None
+
+        logger.info("AgentController shut down")
+
+    # ── Task submission ────────────────────────────────────────────────────
+
+    async def submit(
+        self,
+        task: str,
+        skill_name: str | None = None,
+        explicit_model: str = "",
+        stream: bool = True,
+    ) -> None:
+        """
+        Submit a task for the agent to process.
+
+        Creates an asyncio.Task that runs the agent coroutine.
+        If the agent is not started, starts it automatically.
+        """
+        if not self._agent:
+            await self.start()
+
+        self._busy.set()
+        self._cancel.clear()
+
+        async def _agent_runner():
+            """Run the agent task with proper error handling."""
+            try:
+                logger.info("Agent starting task: %.80s", task)
+                result = await self._agent.run(
+                    task, stream=stream,
+                    skill_name=skill_name,
+                    explicit_model=explicit_model,
+                )
+                logger.info("Agent completed task (len=%d)", len(result))
+            except asyncio.CancelledError:
+                logger.info("Agent task cancelled")
+                raise
+            except Exception as e:
+                logger.exception("Agent task failed")
+                await self.event_queue.put(
+                    ErrorEvent(f"Agent error: {e}")
+                )
+                await self.event_queue.put(
+                    CompleteEvent(
+                        total_tool_calls=(
+                            self._agent._state.tool_call_count
+                            if self._agent else 0
+                        ),
+                        total_time=0,
+                    )
+                )
+            finally:
+                self._busy.clear()
+
+        self._agent_task = asyncio.create_task(_agent_runner())
+
+    async def cancel(self) -> None:
+        """Request cancellation of the current agent run."""
+        self._cancel.set()
+        if self._agent_task and not self._agent_task.done():
+            self._agent_task.cancel()
+        logger.info("Agent cancel requested")
+
+    def is_busy(self) -> bool:
+        """Check if the agent is currently processing a task."""
+        return self._busy.is_set()
+
+    def is_running(self) -> bool:
+        """Check if the agent task is active."""
+        return self._agent_task is not None and not self._agent_task.done()
+
+    # ── Sub-agent management ───────────────────────────────────────────────
+
+    def set_sub_agent_manager(self, mgr: Any) -> None:
+        """Set the SubAgentManager reference."""
+        self._sub_agent_mgr = mgr
+        if self._agent:
+            self._agent.set_sub_agent_manager(mgr)
+
+    # ── Health ─────────────────────────────────────────────────────────────
+
+    @property
+    def agent(self) -> Optional[CoderAgent]:
+        return self._agent
+
+    def health_status(self) -> dict[str, Any]:
+        """Return health status (simplified — no threads to monitor)."""
+        return {
+            "agent": "running" if self.is_running() else "idle",
+            "busy": self._busy.is_set(),
+            "sub_agents": len(self._sub_agent_mgr.list_active()) if self._sub_agent_mgr else 0,
+        }

ata_coder/agent_extension.py

@@ -0,0 +1,69 @@
+"""Extension management (skills, discovery, hook points) — mixin for CoderAgent."""
+import logging
+from typing import Any
+
+from .skill_extension import SkillExtension
+
+logger = logging.getLogger(__name__)
+
+
+class ExtensionMixin:
+    """Extension lifecycle: register skills, discover extensions, register hook points."""
+
+    # ── Extension management ──────────────────────────────────────────────
+
+    def _register_skills_as_extensions(self) -> None:
+        """Register all loaded SkillManager skills as SkillExtension adapters."""
+        if not self.subsys.has_skills:
+            return
+        for skill in self.subsys.skills.list_skills():
+            ext = SkillExtension(skill)
+            if self.ext_mgr.register(ext):
+                logger.debug("Registered skill extension: skill:%s", skill.name)
+            else:
+                logger.debug("Skill extension already registered: skill:%s", skill.name)
+
+    def _discover_extensions(self) -> None:
+        """Discover extensions from configured extension directories."""
+        ext_dirs = getattr(self.config.agent, "extension_dirs", [])
+        if not ext_dirs:
+            return
+        for d in ext_dirs:
+            loaded = self.ext_mgr.discover(d)
+            if loaded:
+                logger.debug("Discovered %d extensions in %s", len(loaded), d)
+
+    def _register_extension_points(self) -> None:
+        """Register hook points extensions can tap into."""
+        self._ep_on_run_start = self.ext_mgr.extension_point(
+            "on_agent_run_start",
+            "Called when agent.run() starts — (task, skill_name)"
+        )
+        self._ep_on_run_complete = self.ext_mgr.extension_point(
+            "on_agent_run_complete",
+            "Called when agent.run() completes — (task, result, tool_call_count)"
+        )
+        self._ep_on_tool_execute = self.ext_mgr.extension_point(
+            "on_tool_execute",
+            "Called before each tool execution — (tool_name, arguments)"
+        )
+        self._ep_on_tool_result = self.ext_mgr.extension_point(
+            "on_tool_result",
+            "Called after each tool result — (tool_name, result)"
+        )
+        self._ep_on_system_prompt = self.ext_mgr.extension_point(
+            "on_system_prompt_build",
+            "Called during system prompt construction — (prompt, task)"
+        )
+        self._ep_on_model_route = self.ext_mgr.extension_point(
+            "on_model_route",
+            "Called after model routing — (task, complexity, model)"
+        )
+        logger.debug(
+            "Registered %d extension points",
+            len(self.ext_mgr.list_extension_points()),
+        )
+
+    def set_sub_agent_manager(self, mgr: Any) -> None:
+        """Set the SubAgentManager for spawn_subagent tool support."""
+        self._sub_agent_mgr = mgr

ata_coder/agent_routing.py

@@ -0,0 +1,105 @@
+"""Model routing and task complexity classification — mixin for CoderAgent."""
+import logging
+
+from .settings import get_settings
+
+logger = logging.getLogger(__name__)
+
+
+class ModelRoutingMixin:
+    """Model routing based on task complexity heuristics and effort level."""
+
+    # ── Model routing ────────────────────────────────────────────────────
+
+    def _route_for_task(self, task: str, explicit_model: str = "") -> None:
+        """Route to the appropriate model based on task complexity + effort level."""
+        if explicit_model:
+            self._route_model(explicit_model)
+            self._routed_complexity = "explicit"
+        else:
+            settings = get_settings()
+            complexity = self._ai_classify(task)
+
+            if complexity == "simple":
+                routed_model = settings.model_haiku
+            elif complexity == "complex":
+                routed_model = settings.model_opus
+            elif complexity == "normal":
+                routed_model = settings.default_model
+            else:
+                logger.warning("Unknown complexity %r", complexity)
+                routed_model = settings.default_model
+
+            # Adjust by effort
+            effort = getattr(self.config, "effort", "medium")
+            if effort == "low":
+                routed_model = settings.model_haiku
+            elif effort in ("xhigh", "max"):
+                routed_model = settings.model_opus
+
+            self._route_model(routed_model)
+            self._routed_complexity = complexity
+
+        logger.info("Model: %s (complexity=%s)", self.current_model, self._routed_complexity)
+
+    def _route_model(self, model: str) -> None:
+        """Switch the LLM client to use a different model at runtime."""
+        if self.llm.config.model == model:
+            return
+        logger.info("Switching model: %s → %s", self.llm.config.model, model)
+        self.llm.config.model = model
+
+    @property
+    def current_model(self) -> str:
+        return self.llm.config.model
+
+    @staticmethod
+    def _ai_classify(task: str) -> str:
+        """
+        Classify task complexity using fast heuristics (NO extra LLM call).
+
+        Returns 'simple', 'complex', or 'normal'.
+
+        Heuristics:
+        - Very short (< 60 chars) → simple
+        - Very long (> 500 chars) → complex
+        - Contains complexity keywords → complex
+        - Contains simple-query keywords → simple
+        - Default → normal
+        """
+        task_lower = task.lower().strip()
+        task_len = len(task)
+
+        # ── Length-based shortcut (configurable thresholds) ──────────────
+        try:
+            s = get_settings()
+            simple_max = s.get("complexity", "simple_max_chars", default=60)
+            complex_min = s.get("complexity", "complex_min_chars", default=500)
+        except Exception:
+            simple_max, complex_min = 60, 500  # fallback defaults
+
+        if task_len <= simple_max:
+            return "simple"
+        if task_len >= complex_min:
+            return "complex"
+
+        # ── Complexity keywords ───────────────────────────────────────
+        complex_keywords = [
+            "implement", "refactor", "architecture", "migrate", "redesign",
+            "optimize", "debug", "fix bug", "restructure", "multi-file",
+            "across multiple", "entire codebase", "from scratch", "set up",
+            "configure", "deploy", "pipeline", "database schema", "api endpoint",
+        ]
+        if any(kw in task_lower for kw in complex_keywords):
+            return "complex"
+
+        # ── Simple-query keywords ─────────────────────────────────────
+        simple_keywords = [
+            "what is", "how do i", "explain", "show me", "find", "search",
+            "list", "tell me about", "describe", "definition of", "example of",
+            "difference between", "why does", "where is",
+        ]
+        if any(kw in task_lower for kw in simple_keywords):
+            return "simple"
+
+        return "normal"

ata_coder/agent_subsystems.py

@@ -0,0 +1,72 @@
+"""
+Agent subsystems container — replaces the loose collection of Optional[X]
+parameters in CoderAgent.__init__ with a single structured dataclass.
+
+Each subsystem is independently initialisable; CoderAgent only needs the
+ones that are actually provided.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .extension import ExtensionManager
+    from .skills import SkillManager
+    from .memory import MemoryStore
+    from .mcp_client import MCPClient
+    from .prompt_template import TemplateManager
+    from .permissions import PermissionStore
+    from .project import ProjectInfo
+    from .session import SessionManager
+
+
+@dataclass
+class AgentSubsystems:
+    """Holds all optional subsystems wired into the agent at construction time.
+
+    Any field left as None means that feature is disabled — the agent skips
+    the corresponding behaviours (no skill detection, no memory recall, etc.).
+    """
+
+    skills: SkillManager | None = None
+    memory: MemoryStore | None = None
+    mcp: MCPClient | None = None
+    templates: TemplateManager | None = None
+    permissions: PermissionStore | None = None
+    project_info: ProjectInfo | None = None
+    sessions: SessionManager | None = None
+    extensions: ExtensionManager | None = None
+
+    @property
+    def has_skills(self) -> bool:
+        return self.skills is not None
+
+    @property
+    def has_memory(self) -> bool:
+        return self.memory is not None
+
+    @property
+    def has_mcp(self) -> bool:
+        return self.mcp is not None
+
+    @property
+    def has_templates(self) -> bool:
+        return self.templates is not None
+
+    @property
+    def has_permissions(self) -> bool:
+        return self.permissions is not None
+
+    @property
+    def has_project_info(self) -> bool:
+        return self.project_info is not None
+
+    @property
+    def has_sessions(self) -> bool:
+        return self.sessions is not None
+
+    @property
+    def has_extensions(self) -> bool:
+        return self.extensions is not None