agencode 0.1.0__py3-none-any.whl

agencli/core/session.py

@@ -0,0 +1,513 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime, timezone
+import inspect
+from pathlib import Path
+import sqlite3
+from threading import RLock
+from typing import Any
+from uuid import uuid4
+
+from langgraph.checkpoint.memory import InMemorySaver, PersistentDict
+
+try:
+    from langgraph.checkpoint.sqlite import SqliteSaver
+except ImportError:  # pragma: no cover - exercised via fallback tests in this env.
+    SqliteSaver = None
+
+try:
+    import aiosqlite
+    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+except ImportError:  # pragma: no cover - exercised via fallback tests in this env.
+    aiosqlite = None
+    AsyncSqliteSaver = None
+
+
+@dataclass(slots=True)
+class ChatTurn:
+    role: str
+    content: str
+    created_at: str
+    agent_name: str | None = None
+    thread_id: str | None = None
+
+
+@dataclass(slots=True)
+class ChatThreadSummary:
+    agent_name: str | None
+    thread_id: str | None
+    turn_count: int
+    last_role: str
+    last_content: str
+    last_created_at: str
+
+
+_CHECKPOINTERS: dict[Path, Any] = {}
+_ASYNC_CHECKPOINTERS: dict[Path, Any] = {}
+
+
+class PersistentInMemoryCheckpointer(InMemorySaver):
+    """File-backed LangGraph checkpointer using LangGraph's persistent dict helper."""
+
+    def __init__(self, base_path: Path) -> None:
+        super().__init__()
+        self.base_path = base_path
+        self.base_path.parent.mkdir(parents=True, exist_ok=True)
+        self._sync_lock = RLock()
+        self.storage = self._load_dict(
+            self.base_path.with_suffix(".storage.pkl"),
+            lambda: defaultdict(dict),
+        )
+        self.writes = self._load_dict(
+            self.base_path.with_suffix(".writes.pkl"),
+            dict,
+        )
+        self.blobs = self._load_dict(
+            self.base_path.with_suffix(".blobs.pkl"),
+            dict,
+        )
+
+    def put(
+        self,
+        config: dict[str, Any],
+        checkpoint: Any,
+        metadata: Any,
+        new_versions: Any,
+    ) -> dict[str, Any]:
+        with self._sync_lock:
+            updated = super().put(config, checkpoint, metadata, new_versions)
+            self._sync()
+            return updated
+
+    def put_writes(
+        self,
+        config: dict[str, Any],
+        writes: Any,
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        with self._sync_lock:
+            super().put_writes(config, writes, task_id, task_path=task_path)
+            self._sync()
+
+    def delete_thread(self, thread_id: str) -> None:
+        with self._sync_lock:
+            super().delete_thread(thread_id)
+            self._sync()
+
+    def _load_dict(self, path: Path, default_factory: Any) -> PersistentDict:
+        store = PersistentDict(default_factory, filename=str(path))
+        if path.exists():
+            store.load()
+        return store
+
+    def _sync(self) -> None:
+        self.storage.sync()
+        self.writes.sync()
+        self.blobs.sync()
+
+
+def ensure_session_db(sessions_dir: str, session_name: str = "default") -> Path:
+    sessions_path = Path(sessions_dir).expanduser()
+    sessions_path.mkdir(parents=True, exist_ok=True)
+    path = sessions_path / f"{session_name}.sqlite3"
+    connection = sqlite3.connect(path)
+    try:
+        connection.execute("CREATE TABLE IF NOT EXISTS session_meta (key TEXT PRIMARY KEY, value TEXT)")
+        connection.execute(
+            """
+            CREATE TABLE IF NOT EXISTS chat_history (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                role TEXT NOT NULL,
+                content TEXT NOT NULL,
+                created_at TEXT NOT NULL
+            )
+            """
+        )
+        columns = {
+            row[1]
+            for row in connection.execute("PRAGMA table_info(chat_history)").fetchall()
+        }
+        if "agent_name" not in columns:
+            connection.execute("ALTER TABLE chat_history ADD COLUMN agent_name TEXT")
+        if "thread_id" not in columns:
+            connection.execute("ALTER TABLE chat_history ADD COLUMN thread_id TEXT")
+        connection.commit()
+    finally:
+        connection.close()
+    return path
+
+
+def append_chat_turn(
+    sessions_dir: str,
+    role: str,
+    content: str,
+    session_name: str = "default",
+    *,
+    agent_name: str | None = None,
+    thread_id: str | None = None,
+) -> None:
+    path = ensure_session_db(sessions_dir, session_name=session_name)
+    created_at = datetime.now(timezone.utc).isoformat()
+    connection = sqlite3.connect(path)
+    try:
+        connection.execute(
+            """
+            INSERT INTO chat_history(role, content, created_at, agent_name, thread_id)
+            VALUES (?, ?, ?, ?, ?)
+            """,
+            (role, content, created_at, agent_name, thread_id),
+        )
+        connection.commit()
+    finally:
+        connection.close()
+
+
+def delete_chat_thread(
+    sessions_dir: str,
+    session_name: str = "default",
+    *,
+    agent_name: str | None = None,
+    thread_id: str | None = None,
+) -> int:
+    if thread_id is None:
+        return 0
+    path = ensure_session_db(sessions_dir, session_name=session_name)
+    connection = sqlite3.connect(path)
+    try:
+        where_clauses = ["thread_id = ?"]
+        params: list[Any] = [thread_id]
+        if agent_name is not None:
+            where_clauses.append("agent_name = ?")
+            params.append(agent_name)
+        cursor = connection.execute(
+            f"DELETE FROM chat_history WHERE {' AND '.join(where_clauses)}",
+            params,
+        )
+        connection.commit()
+        return int(cursor.rowcount or 0)
+    finally:
+        connection.close()
+
+
+def load_chat_history(
+    sessions_dir: str,
+    session_name: str = "default",
+    limit: int = 50,
+    *,
+    agent_name: str | None = None,
+    thread_id: str | None = None,
+) -> list[ChatTurn]:
+    path = ensure_session_db(sessions_dir, session_name=session_name)
+    connection = sqlite3.connect(path)
+    try:
+        where_clauses: list[str] = []
+        params: list[Any] = []
+        if agent_name is not None:
+            where_clauses.append("agent_name = ?")
+            params.append(agent_name)
+        if thread_id is not None:
+            where_clauses.append("thread_id = ?")
+            params.append(thread_id)
+        where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
+        rows = connection.execute(
+            f"""
+            SELECT role, content, created_at, agent_name, thread_id
+            FROM chat_history
+            {where_sql}
+            ORDER BY id DESC
+            LIMIT ?
+            """,
+            [*params, limit],
+        ).fetchall()
+    finally:
+        connection.close()
+    return [
+        ChatTurn(
+            role=row[0],
+            content=row[1],
+            created_at=row[2],
+            agent_name=row[3],
+            thread_id=row[4],
+        )
+        for row in reversed(rows)
+    ]
+
+
+def list_chat_threads(
+    sessions_dir: str,
+    session_name: str = "default",
+    limit: int = 50,
+    *,
+    agent_name: str | None = None,
+    thread_id: str | None = None,
+    query: str | None = None,
+) -> list[ChatThreadSummary]:
+    path = ensure_session_db(sessions_dir, session_name=session_name)
+    connection = sqlite3.connect(path)
+    try:
+        grouped_where_clauses: list[str] = []
+        grouped_params: list[Any] = []
+        if agent_name is not None:
+            grouped_where_clauses.append("agent_name = ?")
+            grouped_params.append(agent_name)
+        if thread_id is not None:
+            grouped_where_clauses.append("thread_id = ?")
+            grouped_params.append(thread_id)
+        grouped_where_sql = f"WHERE {' AND '.join(grouped_where_clauses)}" if grouped_where_clauses else ""
+
+        latest_where_clauses: list[str] = []
+        latest_params: list[Any] = []
+        if query:
+            pattern = f"%{query}%"
+            latest_where_clauses.append(
+                "(COALESCE(latest.agent_name, '') LIKE ? OR COALESCE(latest.thread_id, '') LIKE ? OR latest.content LIKE ?)"
+            )
+            latest_params.extend([pattern, pattern, pattern])
+        latest_where_sql = f"WHERE {' AND '.join(latest_where_clauses)}" if latest_where_clauses else ""
+        rows = connection.execute(
+            f"""
+            SELECT latest.agent_name, latest.thread_id, grouped.turn_count, latest.role, latest.content, latest.created_at
+            FROM (
+                SELECT agent_name, thread_id, COUNT(*) AS turn_count, MAX(id) AS last_id
+                FROM chat_history
+                {grouped_where_sql}
+                GROUP BY agent_name, thread_id
+            ) AS grouped
+            JOIN chat_history AS latest ON latest.id = grouped.last_id
+            {latest_where_sql}
+            ORDER BY latest.id DESC
+            LIMIT ?
+            """,
+            [*grouped_params, *latest_params, limit],
+        ).fetchall()
+    finally:
+        connection.close()
+    return [
+        ChatThreadSummary(
+            agent_name=row[0],
+            thread_id=row[1],
+            turn_count=row[2],
+            last_role=row[3],
+            last_content=row[4],
+            last_created_at=row[5],
+        )
+        for row in rows
+    ]
+
+
+def get_thread_id(agent_name: str, session_name: str = "default") -> str:
+    return f"{session_name}:{agent_name}"
+
+
+def create_thread_id(agent_name: str, session_name: str = "default") -> str:
+    return f"{get_thread_id(agent_name, session_name=session_name)}:{uuid4().hex[:8]}"
+
+
+def resolve_thread_id(
+    agent_name: str,
+    session_name: str = "default",
+    *,
+    thread_id: str | None = None,
+    new_thread: bool = False,
+) -> str:
+    if thread_id is not None and new_thread:
+        raise ValueError("thread_id and new_thread are mutually exclusive")
+    if thread_id is not None:
+        return thread_id
+    if new_thread:
+        return create_thread_id(agent_name, session_name=session_name)
+    return get_thread_id(agent_name, session_name=session_name)
+
+
+def build_thread_config(
+    agent_name: str,
+    session_name: str = "default",
+    *,
+    thread_id: str | None = None,
+) -> dict[str, dict[str, str]]:
+    return {"configurable": {"thread_id": resolve_thread_id(agent_name, session_name=session_name, thread_id=thread_id)}}
+
+
+def describe_langgraph_checkpointer(
+    sessions_dir: str,
+    session_name: str = "default",
+) -> str:
+    if _should_use_sqlite_checkpointer(sessions_dir, session_name=session_name):
+        return "sqlite"
+    return "persistent-memory"
+
+
+def clear_langgraph_checkpointer_cache() -> None:
+    for checkpointer in _CHECKPOINTERS.values():
+        _close_checkpointer_connection(checkpointer)
+    for checkpointer in _ASYNC_CHECKPOINTERS.values():
+        _close_checkpointer_connection(checkpointer)
+    _CHECKPOINTERS.clear()
+    _ASYNC_CHECKPOINTERS.clear()
+
+
+async def clear_langgraph_checkpointer_cache_async() -> None:
+    for checkpointer in _CHECKPOINTERS.values():
+        _close_checkpointer_connection(checkpointer)
+    for checkpointer in _ASYNC_CHECKPOINTERS.values():
+        await _aclose_checkpointer_connection(checkpointer)
+    _CHECKPOINTERS.clear()
+    _ASYNC_CHECKPOINTERS.clear()
+
+
+def ensure_langgraph_checkpointer(
+    sessions_dir: str,
+    session_name: str = "default",
+) -> Any:
+    if _should_use_sqlite_checkpointer(sessions_dir, session_name=session_name):
+        sqlite_path = _sqlite_checkpointer_path(sessions_dir, session_name=session_name)
+        cached = _CHECKPOINTERS.get(sqlite_path)
+        if cached is not None:
+            return cached
+        checkpointer = SqliteSaver(sqlite3.connect(sqlite_path, check_same_thread=False))
+        _CHECKPOINTERS[sqlite_path] = checkpointer
+        return checkpointer
+
+    base_path = _legacy_checkpointer_base_path(sessions_dir, session_name=session_name)
+    cached = _CHECKPOINTERS.get(base_path)
+    if cached is not None:
+        return cached
+    checkpointer = PersistentInMemoryCheckpointer(base_path)
+    _CHECKPOINTERS[base_path] = checkpointer
+    return checkpointer
+
+
+async def ensure_langgraph_checkpointer_async(
+    sessions_dir: str,
+    session_name: str = "default",
+) -> Any:
+    if _should_use_sqlite_checkpointer(sessions_dir, session_name=session_name) and AsyncSqliteSaver is not None and aiosqlite is not None:
+        sqlite_path = _sqlite_checkpointer_path(sessions_dir, session_name=session_name)
+        cached = _ASYNC_CHECKPOINTERS.get(sqlite_path)
+        if cached is not None:
+            return cached
+        connection = await aiosqlite.connect(sqlite_path)
+        checkpointer = AsyncSqliteSaver(connection)
+        _ASYNC_CHECKPOINTERS[sqlite_path] = checkpointer
+        return checkpointer
+    return ensure_langgraph_checkpointer(sessions_dir, session_name=session_name)
+
+
+def has_thread_checkpoint(
+    sessions_dir: str,
+    agent_name: str,
+    session_name: str = "default",
+    *,
+    thread_id: str | None = None,
+) -> bool:
+    checkpointer = ensure_langgraph_checkpointer(sessions_dir, session_name=session_name)
+    return checkpointer.get_tuple(
+        build_thread_config(agent_name, session_name=session_name, thread_id=thread_id)
+    ) is not None
+
+
+async def has_thread_checkpoint_async(
+    sessions_dir: str,
+    agent_name: str,
+    session_name: str = "default",
+    *,
+    thread_id: str | None = None,
+) -> bool:
+    checkpointer = await ensure_langgraph_checkpointer_async(sessions_dir, session_name=session_name)
+    config = build_thread_config(agent_name, session_name=session_name, thread_id=thread_id)
+    aget_tuple = getattr(checkpointer, "aget_tuple", None)
+    if aget_tuple is not None:
+        return await aget_tuple(config) is not None
+    return checkpointer.get_tuple(config) is not None
+
+
+def delete_langgraph_thread(
+    sessions_dir: str,
+    thread_id: str,
+    session_name: str = "default",
+) -> bool:
+    checkpointer = ensure_langgraph_checkpointer(sessions_dir, session_name=session_name)
+    delete_thread = getattr(checkpointer, "delete_thread", None)
+    if delete_thread is None:
+        return False
+    delete_thread(thread_id)
+    return True
+
+
+async def delete_langgraph_thread_async(
+    sessions_dir: str,
+    thread_id: str,
+    session_name: str = "default",
+) -> bool:
+    checkpointer = await ensure_langgraph_checkpointer_async(sessions_dir, session_name=session_name)
+    adelete_thread = getattr(checkpointer, "adelete_thread", None)
+    if adelete_thread is not None:
+        await adelete_thread(thread_id)
+        return True
+    delete_thread = getattr(checkpointer, "delete_thread", None)
+    if delete_thread is None:
+        return False
+    delete_thread(thread_id)
+    return True
+
+
+def _should_use_sqlite_checkpointer(sessions_dir: str, session_name: str = "default") -> bool:
+    if SqliteSaver is None:
+        return False
+    sqlite_path = _sqlite_checkpointer_path(sessions_dir, session_name=session_name)
+    if sqlite_path.exists():
+        return True
+    return not _has_legacy_checkpoint_artifacts(_legacy_checkpointer_base_path(sessions_dir, session_name=session_name))
+
+
+def _legacy_checkpointer_base_path(sessions_dir: str, session_name: str = "default") -> Path:
+    return Path(sessions_dir).expanduser() / f"{session_name}.langgraph"
+
+
+def _sqlite_checkpointer_path(sessions_dir: str, session_name: str = "default") -> Path:
+    base_dir = Path(sessions_dir).expanduser()
+    base_dir.mkdir(parents=True, exist_ok=True)
+    return base_dir / f"{session_name}.langgraph.sqlite3"
+
+
+def _has_legacy_checkpoint_artifacts(base_path: Path) -> bool:
+    return any(
+        artifact.exists()
+        for artifact in (
+            base_path.with_suffix(".storage.pkl"),
+            base_path.with_suffix(".writes.pkl"),
+            base_path.with_suffix(".blobs.pkl"),
+        )
+    )
+
+
+def _close_checkpointer_connection(checkpointer: Any) -> None:
+    connection = getattr(checkpointer, "connection", None) or getattr(checkpointer, "conn", None)
+    if connection is None:
+        return
+    try:
+        result = connection.close()
+        if inspect.isawaitable(result):
+            try:
+                import asyncio
+
+                asyncio.run(result)
+            except Exception:
+                pass
+    except Exception:
+        pass
+
+
+async def _aclose_checkpointer_connection(checkpointer: Any) -> None:
+    connection = getattr(checkpointer, "connection", None) or getattr(checkpointer, "conn", None)
+    if connection is None:
+        return
+    try:
+        result = connection.close()
+        if inspect.isawaitable(result):
+            await result
+    except Exception:
+        pass

agencli/mcp/__init__.py

@@ -0,0 +1 @@
+"""MCP configuration and client helpers."""

agencli/mcp/client.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Any
+
+from agencli.mcp.config import load_mcp_servers
+
+
+class MCPToolClient:
+    """Thin wrapper around MultiServerMCPClient for AgenCLI."""
+
+    def __init__(self, mcp_config_path: str) -> None:
+        self.mcp_config_path = mcp_config_path
+
+    def load_server_map(self) -> dict[str, dict[str, Any]]:
+        servers = load_mcp_servers(self.mcp_config_path)
+        return {name: server.to_client_dict() for name, server in servers.items()}
+
+    async def get_tools(self, server_names: list[str] | None = None) -> list[Any]:
+        try:
+            from langchain_mcp_adapters.client import MultiServerMCPClient
+        except ImportError as exc:
+            raise RuntimeError(
+                "langchain-mcp-adapters is not installed yet. Run `uv sync` after dependencies are added."
+            ) from exc
+
+        server_map = self.load_server_map()
+        if server_names is not None:
+            server_map = {name: server_map[name] for name in server_names if name in server_map}
+        if not server_map:
+            return []
+
+        client = MultiServerMCPClient(server_map)
+        return await client.get_tools()

agencli/mcp/config.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+import json
+import os
+from pathlib import Path
+import shutil
+from typing import Any
+
+
+@dataclass(slots=True)
+class MCPServerConfig:
+    transport: str
+    command: str | None = None
+    args: list[str] = field(default_factory=list)
+    url: str | None = None
+    env: dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, raw: dict[str, Any]) -> "MCPServerConfig":
+        return cls(
+            transport=raw["transport"],
+            command=raw.get("command"),
+            args=list(raw.get("args", [])),
+            url=raw.get("url"),
+            env=dict(raw.get("env", {})),
+        )
+
+    def to_client_dict(self) -> dict[str, Any]:
+        payload: dict[str, Any] = {"transport": self.transport}
+        if self.command:
+            payload["command"] = _resolve_command(self.command)
+        if self.args:
+            payload["args"] = [_resolve_template(arg) for arg in self.args]
+        if self.url:
+            payload["url"] = _resolve_template(self.url)
+        if self.env:
+            payload["env"] = {name: _resolve_template(value) for name, value in self.env.items()}
+        return payload
+
+
+def _resolve_template(value: str) -> str:
+    for key, env_value in os.environ.items():
+        value = value.replace(f"{{{key}}}", env_value)
+    return value
+
+
+def _resolve_command(command: str) -> str:
+    if os.name != "nt":
+        return command
+    if shutil.which(command):
+        return command
+
+    for suffix in (".cmd", ".exe", ".bat"):
+        candidate = command if command.lower().endswith(suffix) else f"{command}{suffix}"
+        if shutil.which(candidate):
+            return candidate
+    return command
+
+
+def load_mcp_servers(mcp_config_path: str) -> dict[str, MCPServerConfig]:
+    path = Path(mcp_config_path).expanduser()
+    if not path.exists():
+        return {}
+
+    raw = json.loads(path.read_text(encoding="utf-8"))
+    if "mcpServers" in raw:
+        raw = raw["mcpServers"]
+    return {name: MCPServerConfig.from_dict(server) for name, server in raw.items()}
+
+
+def save_mcp_servers(mcp_config_path: str, servers: dict[str, MCPServerConfig]) -> Path:
+    path = Path(mcp_config_path).expanduser()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {"mcpServers": {name: asdict(server) for name, server in servers.items()}}
+    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    return path
+
+
+def upsert_mcp_server(mcp_config_path: str, name: str, server: MCPServerConfig) -> Path:
+    servers = load_mcp_servers(mcp_config_path)
+    servers[name] = server
+    return save_mcp_servers(mcp_config_path, servers)
+
+
+def default_mcp_servers(workspace_dir: str) -> dict[str, MCPServerConfig]:
+    return {
+        "search": MCPServerConfig(
+            transport="http",
+            url="https://mcp.tavily.com/mcp/?tavilyApiKey={TAVILY_API_KEY}",
+        ),
+    }
+
+
+def bootstrap_default_mcp_servers(mcp_config_path: str, workspace_dir: str) -> Path:
+    servers = load_mcp_servers(mcp_config_path)
+    for name, server in default_mcp_servers(workspace_dir).items():
+        servers.setdefault(name, server)
+    return save_mcp_servers(mcp_config_path, servers)

agencli/providers/__init__.py

@@ -0,0 +1 @@
+"""Model provider utilities."""