ata-coder 2.4.2__py3-none-any.whl

ata_coder/server_session.py

@@ -0,0 +1,159 @@
+"""Thread-safe session store for the API server."""
+import asyncio
+import logging
+import os
+import threading
+import time
+import uuid
+from datetime import datetime, timezone
+
+from .config import AppConfig
+from .tools import ToolExecutor
+from .agent import CoderAgent
+from .agent_subsystems import AgentSubsystems
+from .permissions import PermissionStore, PermissionMode
+
+logger = logging.getLogger(__name__)
+
+
+class SessionStore:
+    """Thread-safe session store for the API server."""
+
+    # Auto-cleanup: sessions idle for this many seconds are evicted.
+    SESSION_TTL_SECONDS = 1800  # 30 minutes
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._sessions: dict[str, CoderAgent] = {}
+        self._metadata: dict[str, dict] = {}
+
+    def _evict_expired(self):
+        """Remove sessions that have been idle longer than SESSION_TTL_SECONDS."""
+        now = time.time()
+        expired: list[str] = []
+        with self._lock:
+            for sid, meta in self._metadata.items():
+                created_str = meta.get("created", "")
+                last_active_str = meta.get("last_active", created_str)
+                try:
+                    t = datetime.strptime(last_active_str, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc).timestamp()
+                except (ValueError, OverflowError):
+                    t = 0  # unparseable — evict (conservative: assume very old)
+                if now - t > self.SESSION_TTL_SECONDS:
+                    expired.append(sid)
+            for sid in expired:
+                agent = self._sessions.pop(sid, None)
+                if agent:
+                    try:
+                        asyncio.run(agent.shutdown())
+                    except Exception:
+                        pass
+                self._metadata.pop(sid, None)
+        if expired:
+            logger.info("Evicted %d expired session(s)", len(expired))
+
+    def create(self, config: AppConfig, skill: str = "") -> tuple[str, CoderAgent]:
+        sid = str(uuid.uuid4())[:12]
+
+        # Guard: if a session with this ID already exists, clean it up first
+        with self._lock:
+            existing = self._sessions.pop(sid, None)
+            self._metadata.pop(sid, None)
+        if existing:
+            try:
+                asyncio.run(existing.shutdown())
+            except Exception:
+                pass
+
+        tool_exec = ToolExecutor(config.agent)
+
+        # Build AgentSubsystems container (replaces loose kwargs)
+        perms = PermissionStore(config.agent.workspace_dir)
+
+        # API mode: default to allow (caller implements own auth)
+        if os.environ.get("ATA_CODER_ALLOW_ALL", "").lower() in ("1", "true", "yes"):
+            perms.set_category_rule("shell", PermissionMode.ALLOW)
+            perms.set_category_rule("write", PermissionMode.ALLOW)
+            logger.warning(
+                "⚠️  ATA_CODER_ALLOW_ALL=1 — ALL shell & write operations "
+                "will be silently allowed without permission prompts. "
+                "This is intended for headless/automated environments only."
+            )
+        else:
+            perms.set_prompt_callback(lambda n, a, c: True)
+
+        from .skills import get_skill_manager
+        from .memory import get_memory_store
+        from .project import ProjectDetector
+
+        skill_mgr = get_skill_manager()
+        if skill:
+            skill_mgr.activate(skill)
+
+        subsystems = AgentSubsystems(
+            skills=skill_mgr,
+            memory=get_memory_store(),
+            permissions=perms,
+            project_info=ProjectDetector(config.agent.workspace_dir).detect(),
+        )
+
+        agent = CoderAgent(
+            config=config,
+            tool_executor=tool_exec,
+            subsystems=subsystems,
+        )
+
+        now_ts = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+        with self._lock:
+            self._sessions[sid] = agent
+            self._metadata[sid] = {
+                "created": now_ts,
+                "last_active": now_ts,
+                "messages": 0,
+                "tool_calls": 0,
+                "skill": skill,
+                "model": config.llm.model,
+            }
+        return sid, agent
+
+    def get(self, sid: str) -> CoderAgent | None:
+        self._evict_expired()
+        with self._lock:
+            agent = self._sessions.get(sid)
+            if agent and sid in self._metadata:
+                self._metadata[sid]["last_active"] = time.strftime("%Y-%m-%dT%H:%M:%SZ")
+            return agent
+
+    def get_or_create(self, sid: str | None, config: AppConfig, skill: str = "") -> tuple[str, CoderAgent]:
+        if sid:
+            existing = self.get(sid)
+            if existing:
+                return sid, existing
+        return self.create(config, skill)
+
+    def update_meta(self, sid: str, **kwargs):
+        with self._lock:
+            if sid in self._metadata:
+                self._metadata[sid].update(kwargs)
+
+    def list_sessions(self) -> list[dict]:
+        with self._lock:
+            return [
+                {"session_id": sid, **meta}
+                for sid, meta in self._metadata.items()
+            ]
+
+    def get_meta(self, sid: str) -> dict | None:
+        with self._lock:
+            return self._metadata.get(sid)
+
+    def delete(self, sid: str) -> bool:
+        with self._lock:
+            agent = self._sessions.pop(sid, None)
+            if agent:
+                try:
+                    asyncio.run(agent.shutdown())
+                except Exception:
+                    pass
+            self._metadata.pop(sid, None)
+            return agent is not None

ata_coder/server_shell.py

@@ -0,0 +1,129 @@
+"""Persistent PowerShell/bash sessions for interactive terminal in the API server."""
+import logging
+import os
+import queue
+import subprocess
+import threading
+import uuid
+
+logger = logging.getLogger(__name__)
+
+_shell_sessions: dict[str, tuple[subprocess.Popen, "queue.Queue[str]", threading.Lock, str]] = {}
+_shell_lock = threading.Lock()
+
+
+def shell_open(cwd: str, sid: str = "") -> str:
+    """Start a persistent shell process. Returns session ID.
+
+    Uses PowerShell on Windows, bash on Linux/macOS.
+    """
+    sid = sid or uuid.uuid4().hex[:10]
+    out_queue: "queue.Queue[str]" = queue.Queue()
+
+    if os.name == "nt":
+        # Windows: PowerShell with binary pipes to completely bypass
+        # the internal _readerthread (which ignores encoding and uses GBK).
+        # We create our own reader threads that decode UTF-8 manually.
+        proc = subprocess.Popen(
+            ["powershell.exe", "-NoLogo", "-NoExit"],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE, cwd=cwd,
+        )
+        # Switch PS to UTF-8 output (must use bytes since stdin is binary)
+        proc.stdin.write("chcp 65001 >$null\n".encode("utf-8"))
+        proc.stdin.write("[Console]::OutputEncoding = [Text.Encoding]::UTF8\n".encode("utf-8"))
+        proc.stdin.flush()
+        prompt_text = "PS> "
+    else:
+        # Linux/macOS: bash with binary pipes
+        proc = subprocess.Popen(
+            ["bash", "--norc"],
+            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE, cwd=cwd,
+        )
+        prompt_text = "$ "
+
+    def _reader(pipe, label):
+        """Read binary lines from a pipe, decode UTF-8, and enqueue."""
+        for line_bytes in pipe:
+            try:
+                out_queue.put(line_bytes.decode("utf-8", errors="replace"))
+            except Exception:
+                pass
+
+    threading.Thread(target=_reader, args=(proc.stdout, "stdout"), daemon=True).start()
+    threading.Thread(target=_reader, args=(proc.stderr, "stderr"), daemon=True).start()
+
+    # Drain startup banner — wait for initial output then drain
+    try:
+        out_queue.get(timeout=3.0)  # first line proves the process is alive
+    except queue.Empty:
+        pass  # process produced no banner, proceed anyway
+    while not out_queue.empty():
+        try:
+            out_queue.get_nowait()
+        except queue.Empty:
+            break
+
+    with _shell_lock:
+        _shell_sessions[sid] = (proc, out_queue, threading.Lock(), prompt_text)
+    logger.info("Shell opened: %s", sid)
+    return sid
+
+
+def shell_ensure(sid: str, cwd: str):
+    """Get or create a shell session."""
+    with _shell_lock:
+        if sid in _shell_sessions:
+            return _shell_sessions[sid]
+    # Auto-create with the requested SID
+    shell_open(cwd, sid=sid)
+    with _shell_lock:
+        return _shell_sessions.get(sid, (None, None, None, ""))
+
+
+def shell_close(sid: str):
+    """Close a single shell session, terminating the process and releasing pipes."""
+    with _shell_lock:
+        entry = _shell_sessions.pop(sid, None)
+    if entry:
+        proc, _, _, _ = entry
+        try:
+            proc.stdin.write(b"exit\n")
+            proc.stdin.flush()
+            proc.stdin.close()
+        except Exception:
+            pass
+        try:
+            proc.wait(timeout=3)
+        except Exception:
+            try:
+                proc.kill()
+                proc.wait(timeout=1)
+            except Exception:
+                pass
+        # Close remaining pipes to prevent Windows proactor "unclosed transport" warnings
+        if proc.stdout:
+            try:
+                proc.stdout.close()
+            except Exception:
+                pass
+        if proc.stderr:
+            try:
+                proc.stderr.close()
+            except Exception:
+                pass
+        logger.info("Shell closed: %s", sid)
+
+
+def shell_close_all():
+    """Close all active shell sessions. Call during server shutdown."""
+    with _shell_lock:
+        sids = list(_shell_sessions.keys())
+    for sid in sids:
+        shell_close(sid)
+
+
+def get_shell_sessions():
+    """Return the module-level shell sessions dict (for use by handler)."""
+    return _shell_sessions