operator-agent 0.1.0__py3-none-any.whl

operator_agent/core.py

@@ -0,0 +1,376 @@
+"""Core runtime: state management, process spawning, and request handling."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import logging
+import os
+import tempfile
+from asyncio.subprocess import Process
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from .config import CONFIG_DIR, STATE_FILE
+from .providers import BaseProvider, get_provider
+
+if TYPE_CHECKING:
+    from .transports import TransportContext
+
+log = logging.getLogger(__name__)
+
+
+def split_text(text: str, limit: int = 4096) -> list[str]:
+    """Split text at line boundaries for message size limits."""
+    if len(text) <= limit:
+        return [text]
+
+    chunks: list[str] = []
+    current = ""
+    for line in text.split("\n"):
+        if len(current) + len(line) + 1 > limit:
+            if current:
+                chunks.append(current)
+            remainder = line
+            while len(remainder) > limit:
+                chunks.append(remainder[:limit])
+                remainder = remainder[limit:]
+            current = remainder
+        else:
+            current = current + "\n" + line if current else line
+    if current:
+        chunks.append(current)
+    return chunks
+
+
+class Runtime:
+    """Central runtime holding all state and process management."""
+
+    def __init__(self, config: dict):
+        self.config = config
+        self.working_dir: str = config.get("working_dir", str(Path.cwd()))
+        self.models: dict[str, list[str]] = {
+            name: pcfg.get("models", [])
+            for name, pcfg in config.get("providers", {}).items()
+        }
+        self.active_provider_by_chat: dict[int, str] = {}
+        self.active_model_by_chat_provider: dict[tuple[int, str], str] = {}
+        self.session_by_chat_provider: dict[tuple[int, str], str] = {}
+        self.running_process_by_chat: dict[int, Process] = {}
+        self.running_task_by_chat: dict[int, asyncio.Task] = {}
+        self.chat_lock_by_chat: dict[int, asyncio.Lock] = {}
+
+    def init_config_dir(self):
+        """Ensure config directory exists."""
+        os.makedirs(CONFIG_DIR, exist_ok=True)
+
+    # --- State persistence ---
+
+    def save_state(self):
+        """Atomically persist session state to disk."""
+        data = {
+            "active_provider_by_chat": {
+                str(k): v for k, v in self.active_provider_by_chat.items()
+            },
+            "active_model_by_chat_provider": {
+                f"{chat_id}:{provider}": model
+                for (chat_id, provider), model in self.active_model_by_chat_provider.items()
+            },
+            "session_by_chat_provider": {
+                f"{chat_id}:{provider}": sid
+                for (chat_id, provider), sid in self.session_by_chat_provider.items()
+            },
+        }
+        try:
+            fd, tmp = tempfile.mkstemp(dir=CONFIG_DIR, suffix=".tmp")
+            with os.fdopen(fd, "w") as f:
+                json.dump(data, f)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp, STATE_FILE)
+            log.debug(
+                "State saved: %d providers, %d sessions",
+                len(self.active_provider_by_chat),
+                len(self.session_by_chat_provider),
+            )
+        except Exception:
+            log.exception("Failed to save state")
+
+    def load_state(self):
+        """Load persisted state from disk into memory maps."""
+        if not os.path.exists(STATE_FILE):
+            return
+        try:
+            with open(STATE_FILE) as f:
+                data = json.load(f)
+            for k, v in data.get("active_provider_by_chat", {}).items():
+                self.active_provider_by_chat[int(k)] = v
+            for k, v in data.get("active_model_by_chat_provider", {}).items():
+                chat_id_str, provider = k.split(":", 1)
+                self.active_model_by_chat_provider[(int(chat_id_str), provider)] = v
+            for k, v in data.get("session_by_chat_provider", {}).items():
+                chat_id_str, provider = k.split(":", 1)
+                self.session_by_chat_provider[(int(chat_id_str), provider)] = v
+            log.info(
+                "Loaded state: %d providers, %d sessions, %d model overrides",
+                len(self.active_provider_by_chat),
+                len(self.session_by_chat_provider),
+                len(self.active_model_by_chat_provider),
+            )
+        except Exception:
+            log.exception("Failed to load state, starting fresh")
+
+    # --- Accessors ---
+
+    def get_active_provider(self, chat_id: int) -> str:
+        """Return active provider for chat, defaulting to claude."""
+        return self.active_provider_by_chat.get(chat_id, "claude")
+
+    def get_active_model(self, chat_id: int, provider: str) -> str:
+        """Return active model for chat+provider, defaulting to first in list."""
+        stored = self.active_model_by_chat_provider.get((chat_id, provider))
+        if stored and stored in self.models.get(provider, []):
+            return stored
+        models = self.models.get(provider, [])
+        return models[0] if models else "default"
+
+    def get_chat_lock(self, chat_id: int) -> asyncio.Lock:
+        """Get or create a per-chat lock to serialize requests."""
+        lock = self.chat_lock_by_chat.get(chat_id)
+        if lock is None:
+            lock = asyncio.Lock()
+            self.chat_lock_by_chat[chat_id] = lock
+        return lock
+
+    # --- Provider management ---
+
+    def _get_provider_path(self, provider_name: str) -> str:
+        providers_cfg = self.config.get("providers", {})
+        provider_cfg = providers_cfg.get(provider_name, {})
+        return provider_cfg.get("path", provider_name)
+
+    def make_provider(self, provider_name: str) -> BaseProvider:
+        """Create a fresh provider instance."""
+        path = self._get_provider_path(provider_name)
+        return get_provider(provider_name, path)
+
+    # --- Process control ---
+
+    async def stop_chat(self, chat_id: int) -> tuple[bool, str | None]:
+        """Stop running process/task for a chat. Returns (had_something, error_msg)."""
+        process = self.running_process_by_chat.get(chat_id)
+        task = self.running_task_by_chat.get(chat_id)
+
+        if process is None and task is None:
+            return False, None
+
+        try:
+            if process and process.returncode is None:
+                process.terminate()
+                try:
+                    await asyncio.wait_for(process.wait(), timeout=0.5)
+                except TimeoutError:
+                    if process.returncode is None:
+                        process.kill()
+
+            if task and not task.done():
+                task.cancel()
+
+            return True, None
+        except Exception as exc:
+            return True, str(exc)
+
+    # --- Core streaming ---
+
+    async def run_provider(
+        self, provider: BaseProvider, prompt: str, chat_id: int, model: str
+    ):
+        """Spawn a provider subprocess and yield StreamEvents."""
+        session_id = self.session_by_chat_provider.get((chat_id, provider.name))
+        cmd = provider.build_command(prompt, model, session_id)
+
+        log.info("[%s] Spawning: %s", provider.name, " ".join(cmd[:6]) + " ...")
+
+        stderr = (
+            asyncio.subprocess.STDOUT
+            if provider.stderr_to_stdout()
+            else asyncio.subprocess.PIPE
+        )
+        kwargs: dict[str, Any] = {
+            "stdout": asyncio.subprocess.PIPE,
+            "stderr": stderr,
+            "cwd": self.working_dir,
+        }
+        limit = provider.stdout_limit()
+        if limit:
+            kwargs["limit"] = limit
+
+        process = await asyncio.create_subprocess_exec(*cmd, **kwargs)
+        log.info("[%s] Process started, pid=%s", provider.name, process.pid)
+        self.running_process_by_chat[chat_id] = process
+
+        event_count = 0
+        try:
+            assert process.stdout is not None
+            async for line in process.stdout:
+                decoded = line.decode(errors="replace").strip()
+                if not decoded:
+                    continue
+
+                raw = provider.parse_line(decoded)
+                if raw is None:
+                    log.debug("[%s] Skipped line: %s", provider.name, decoded[:200])
+                    continue
+
+                event_count += 1
+                log.debug(
+                    "[%s] Event #%d type=%s",
+                    provider.name,
+                    event_count,
+                    raw.get("type", "?"),
+                )
+
+                for stream_event in provider.parse_event(raw):
+                    if stream_event.kind == "session" and stream_event.session_id:
+                        log.info(
+                            "[%s] New session: %s",
+                            provider.name,
+                            stream_event.session_id,
+                        )
+                        self.session_by_chat_provider[
+                            (chat_id, provider.name)
+                        ] = stream_event.session_id
+                        self.save_state()
+                    yield stream_event
+
+        finally:
+            rc = process.returncode
+            log.info(
+                "[%s] Process pid=%s finished, returncode=%s, events=%d",
+                provider.name,
+                process.pid,
+                rc,
+                event_count,
+            )
+            if rc and rc != 0 and not provider.stderr_to_stdout() and process.stderr:
+                stderr_data = await process.stderr.read()
+                if stderr_data:
+                    stderr_text = stderr_data.decode(errors="replace").strip()[:2000]
+                    log.error("[%s] stderr: %s", provider.name, stderr_text)
+            current = self.running_process_by_chat.get(chat_id)
+            if current is process:
+                self.running_process_by_chat.pop(chat_id, None)
+
+    # --- Request handling ---
+
+    async def process_request(
+        self,
+        provider_name: str,
+        prompt: str,
+        chat_id: int,
+        ctx: TransportContext,
+    ):
+        """Run a full provider request with status ticker and response delivery."""
+        provider = self.make_provider(provider_name)
+        model = self.get_active_model(chat_id, provider_name)
+        display_name = provider_name.capitalize()
+
+        log.info(
+            "[%s] Request from chat_id=%s model=%s: %.80s",
+            provider_name,
+            chat_id,
+            model,
+            prompt,
+        )
+
+        prefix_base = f"{display_name}/{model}"
+        status_msg = await ctx.reply_status(f"[{prefix_base} 0s] Working...")
+        state = {"status": "Working...", "elapsed": 0, "provider": prefix_base}
+        stop = asyncio.Event()
+        ticker = asyncio.create_task(self._run_ticker(ctx, status_msg, state, stop))
+
+        response_parts: list[str] = []
+        error_text = ""
+
+        try:
+            async for event in self.run_provider(provider, prompt, chat_id, model):
+                if event.kind == "status":
+                    state["status"] = event.text
+                elif event.kind == "response":
+                    response_parts.append(event.text)
+                elif event.kind == "error":
+                    error_text = event.text
+
+            stop.set()
+            ticker.cancel()
+
+            await ctx.delete_status(status_msg)
+
+            response_text = provider.format_response(response_parts)
+            prefix = f"[{prefix_base} {state['elapsed']}s]"
+
+            if response_text:
+                log.info(
+                    "[%s] Response (%d chars, %ds)",
+                    provider_name,
+                    len(response_text),
+                    state["elapsed"],
+                )
+                chunks = split_text(response_text)
+                chunks[0] = f"{prefix} {chunks[0]}"
+                for chunk in chunks:
+                    await ctx.reply(chunk)
+            elif error_text:
+                log.warning(
+                    "[%s] Failed after %ds: %s",
+                    provider_name,
+                    state["elapsed"],
+                    error_text[:200],
+                )
+                await ctx.reply(f"{prefix} Error: {error_text[:4000]}")
+            else:
+                log.warning(
+                    "[%s] Empty response after %ds", provider_name, state["elapsed"]
+                )
+                await ctx.reply(f"{prefix} (No response)")
+
+        except asyncio.CancelledError:
+            stop.set()
+            ticker.cancel()
+            with contextlib.suppress(Exception):
+                await ctx.edit_status(status_msg, f"[{prefix_base}] Stopped.")
+            raise
+
+        except Exception as exc:
+            stop.set()
+            ticker.cancel()
+            log.error(
+                "Error processing %s request: %s", provider_name, exc, exc_info=True
+            )
+            error_msg = str(exc)
+            prefix = f"[{prefix_base} {state['elapsed']}s]"
+            try:
+                await ctx.edit_status(
+                    status_msg, f"{prefix} Error: {error_msg[:4000]}"
+                )
+            except Exception:
+                for chunk in split_text(f"{prefix} Error: {error_msg}"):
+                    await ctx.reply(chunk)
+
+    async def _run_ticker(
+        self,
+        ctx: TransportContext,
+        status_msg: Any,
+        state: dict,
+        stop: asyncio.Event,
+    ):
+        """Background task that updates [Provider Xs] status."""
+        while not stop.is_set():
+            await asyncio.sleep(1)
+            if stop.is_set():
+                break
+            state["elapsed"] += 1
+            text = f"[{state['provider']} {state['elapsed']}s] {state['status']}"
+            with contextlib.suppress(Exception):
+                await asyncio.wait_for(ctx.edit_status(status_msg, text), timeout=5.0)

operator_agent/providers/__init__.py

@@ -0,0 +1,86 @@
+"""Provider abstraction for CLI agents."""
+
+from __future__ import annotations
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Literal
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class StreamEvent:
+    """Unified event type emitted by all providers."""
+
+    kind: Literal["status", "response", "session", "error"]
+    text: str = ""
+    session_id: str | None = None
+
+
+class BaseProvider(ABC):
+    """Base class for CLI agent providers."""
+
+    name: str
+
+    def __init__(self, path: str):
+        self.path = path
+
+    @abstractmethod
+    def build_command(
+        self, prompt: str, model: str, session_id: str | None = None
+    ) -> list[str]:
+        """Build the CLI command to execute."""
+
+    @abstractmethod
+    def parse_event(self, event: dict) -> list[StreamEvent]:
+        """Parse a raw JSON event into StreamEvents."""
+
+    def parse_line(self, line: str) -> dict | None:
+        """Parse a raw stdout line into a JSON dict. Returns None to skip."""
+        stripped = line.strip()
+        if not stripped:
+            return None
+        try:
+            return json.loads(stripped)
+        except json.JSONDecodeError:
+            return None
+
+    def stderr_to_stdout(self) -> bool:
+        """If True, merge stderr into stdout."""
+        return False
+
+    def stdout_limit(self) -> int | None:
+        """Buffer limit for stdout, or None for default."""
+        return None
+
+    def format_response(self, parts: list[str]) -> str:
+        """Combine response parts into final text. Default: last part wins."""
+        return parts[-1] if parts else ""
+
+    def clear_session(self, session_id: str | None, working_dir: str) -> str:
+        """Clear session data. Returns human-readable summary."""
+        return "session cleared" if session_id else "no session"
+
+
+PROVIDER_NAMES = ["claude", "codex", "gemini"]
+
+
+def get_provider(name: str, path: str) -> BaseProvider:
+    """Create a provider instance by name."""
+    # Lazy imports: subclasses import from this module, so top-level would be circular.
+    from .claude import ClaudeProvider
+    from .codex import CodexProvider
+    from .gemini import GeminiProvider
+
+    classes: dict[str, type[BaseProvider]] = {
+        "claude": ClaudeProvider,
+        "codex": CodexProvider,
+        "gemini": GeminiProvider,
+    }
+    cls = classes.get(name)
+    if cls is None:
+        raise ValueError(f"Unknown provider: {name}")
+    return cls(path)

operator_agent/providers/claude.py

@@ -0,0 +1,108 @@
+"""Claude CLI provider."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+from . import BaseProvider, StreamEvent
+
+
+def _format_tool_status(tool_name: str, tool_input: dict) -> str:
+    """Format tool usage into human-readable status."""
+    match tool_name:
+        case "Read":
+            path = tool_input.get("file_path", "file")
+            return f"Reading {os.path.basename(path)}..."
+        case "Write":
+            path = tool_input.get("file_path", "file")
+            return f"Writing {os.path.basename(path)}..."
+        case "Edit":
+            path = tool_input.get("file_path", "file")
+            return f"Editing {os.path.basename(path)}..."
+        case "Bash":
+            cmd = tool_input.get("command", "")
+            short = cmd[:50] + "..." if len(cmd) > 50 else cmd
+            return f"Running {short}"
+        case "Glob":
+            pattern = tool_input.get("pattern", "")
+            return f"Finding {pattern}..."
+        case "Grep":
+            pattern = tool_input.get("pattern", "")
+            return f"Searching for {pattern}..."
+        case "WebFetch":
+            url = tool_input.get("url", "")
+            return f"Fetching {url[:40]}..."
+        case "WebSearch":
+            query = tool_input.get("query", "")
+            return f"Searching: {query}..."
+        case "Task":
+            return "Running subagent..."
+        case _:
+            return f"Using {tool_name}..."
+
+
+def _get_project_dir(working_dir: str) -> str:
+    """Derive Claude's project session directory from working_dir."""
+    resolved = str(Path(working_dir).resolve())
+    mangled = resolved.replace("/", "-")
+    return os.path.expanduser(f"~/.claude/projects/{mangled}")
+
+
+class ClaudeProvider(BaseProvider):
+    name = "claude"
+
+    def build_command(self, prompt, model, session_id=None):
+        return [
+            self.path,
+            "-p",
+            "--continue",
+            "--output-format",
+            "stream-json",
+            "--verbose",
+            "--dangerously-skip-permissions",
+            "--model",
+            model,
+            prompt,
+        ]
+
+    def parse_event(self, event):
+        events = []
+        event_type = event.get("type", "")
+
+        if event_type == "assistant":
+            message = event.get("message", {})
+            content = message.get("content", [])
+            for block in content:
+                if block.get("type") == "tool_use":
+                    events.append(
+                        StreamEvent(
+                            kind="status",
+                            text=_format_tool_status(
+                                block.get("name", ""), block.get("input", {})
+                            ),
+                        )
+                    )
+                elif block.get("type") == "text":
+                    block_text = block.get("text", "")
+                    if block_text:
+                        events.append(StreamEvent(kind="response", text=block_text))
+
+        elif event_type == "result":
+            result_text = event.get("result", "")
+            if isinstance(result_text, str) and result_text:
+                events.append(StreamEvent(kind="response", text=result_text))
+
+        return events
+
+    def stdout_limit(self):
+        return 10 * 1024 * 1024
+
+    def clear_session(self, session_id, working_dir):
+        project_dir = Path(_get_project_dir(working_dir))
+        removed = 0
+        if project_dir.is_dir():
+            for f in project_dir.glob("*.jsonl"):
+                f.unlink()
+                removed += 1
+        return f"{removed} session file{'s' if removed != 1 else ''} deleted"

operator_agent/providers/codex.py

@@ -0,0 +1,115 @@
+"""Codex CLI provider."""
+
+from __future__ import annotations
+
+import json
+
+from . import BaseProvider, StreamEvent
+
+
+def _format_status(event: dict) -> str | None:
+    """Extract human-readable status from a Codex JSON event."""
+    event_type = event.get("type", "")
+    item = event.get("item", {})
+    item_type = item.get("type", "")
+
+    if event_type == "item.started":
+        if item_type == "command_execution":
+            cmd = item.get("command", "")
+            if "-lc " in cmd:
+                cmd = cmd.split("-lc ", 1)[1].strip("'\"")
+            short = cmd[:50] + "..." if len(cmd) > 50 else cmd
+            return f"Running {short}"
+        if item_type == "reasoning":
+            return "Thinking..."
+        if item_type == "file_changes":
+            return "Editing files..."
+        if item_type == "web_searches":
+            return "Searching the web..."
+        if item_type == "mcp_tool_calls":
+            return "Using tool..."
+
+    if event_type == "item.completed" and item_type == "reasoning":
+        text = item.get("text", "")
+        if text:
+            clean = text.strip("*").strip()
+            short = clean[:40] + "..." if len(clean) > 40 else clean
+            return short
+
+    return None
+
+
+class CodexProvider(BaseProvider):
+    name = "codex"
+
+    def build_command(self, prompt, model, session_id=None):
+        if session_id:
+            return [
+                self.path,
+                "exec",
+                "resume",
+                "--dangerously-bypass-approvals-and-sandbox",
+                "--json",
+                "-m",
+                model,
+                session_id,
+                prompt,
+            ]
+        return [
+            self.path,
+            "exec",
+            "--dangerously-bypass-approvals-and-sandbox",
+            "--json",
+            "-m",
+            model,
+            prompt,
+        ]
+
+    def parse_line(self, line):
+        stripped = line.strip()
+        if not stripped or not stripped.startswith("{"):
+            return None
+        try:
+            return json.loads(stripped)
+        except json.JSONDecodeError:
+            return None
+
+    def parse_event(self, event):
+        events = []
+        event_type = event.get("type", "")
+
+        status = _format_status(event)
+        if status:
+            events.append(StreamEvent(kind="status", text=status))
+
+        if event_type == "turn.started":
+            events.append(StreamEvent(kind="status", text="Working..."))
+
+        elif event_type == "error":
+            msg = event.get("message")
+            if isinstance(msg, str) and msg:
+                events.append(StreamEvent(kind="error", text=msg))
+                if "reconnect" in msg.lower():
+                    events.append(StreamEvent(kind="status", text="Reconnecting..."))
+
+        elif event_type == "item.completed":
+            item = event.get("item", {})
+            if item.get("type") == "agent_message":
+                text_value = item.get("text")
+                if isinstance(text_value, str) and text_value:
+                    events.append(StreamEvent(kind="response", text=text_value))
+
+        elif event_type == "turn.failed":
+            msg = event.get("message", "")
+            if msg:
+                events.append(StreamEvent(kind="error", text=msg))
+
+        elif event_type == "thread.started":
+            thread_id = event.get("thread_id")
+            if isinstance(thread_id, str) and thread_id:
+                events.append(StreamEvent(kind="session", session_id=thread_id))
+
+        return events
+
+    def stderr_to_stdout(self):
+        return True