存储抽象

# 完整示例：自定义 SessionStore + 运行一次任务
import asyncio
import uuid
from datetime import datetime, timezone

from cody.core.session import Message, Session
from cody.sdk import Cody


class PostgresStyleSessionStore:
    """模拟 PG：内存实现；生产环境把 _sessions/_rows 换成 SQL。"""

    def __init__(self, dsn: str) -> None:
        self._dsn = dsn
        self._sessions: dict[str, Session] = {}
        # 每条消息带单调整数 id，语义对齐 SQLite messages.id
        self._rows: dict[str, list[tuple[int, Message]]] = {}

    def close(self) -> None:
        pass

    def create_session(
        self, title: str = "New session", model: str = "", workdir: str = ""
    ) -> Session:
        now = datetime.now(timezone.utc).isoformat()
        sid = uuid.uuid4().hex[:12]
        s = Session(
            id=sid,
            title=title,
            messages=[],
            model=model,
            workdir=workdir,
            created_at=now,
            updated_at=now,
        )
        self._sessions[sid] = s
        self._rows[sid] = []
        return s

    def add_message(
        self,
        session_id: str,
        role: str,
        content: str,
        images=None,
    ) -> Message:
        msg = Message(role=role, content=content, images=list(images or []))
        rows = self._rows.setdefault(session_id, [])
        new_id = (max((r[0] for r in rows), default=0) + 1)
        rows.append((new_id, msg))
        s = self._sessions[session_id]
        s.updated_at = datetime.now(timezone.utc).isoformat()
        return msg

    def get_session(self, session_id: str) -> Session | None:
        s = self._sessions.get(session_id)
        if s is None:
            return None
        msgs = [m for _, m in self._rows.get(session_id, [])]
        return Session(
            id=s.id,
            title=s.title,
            messages=msgs,
            model=s.model,
            workdir=s.workdir,
            created_at=s.created_at,
            updated_at=s.updated_at,
            message_count=None,
            compacted_summary=s.compacted_summary,
            compacted_up_to=s.compacted_up_to,
        )

    def list_sessions(self, limit: int = 20) -> list[Session]:
        items = sorted(
            self._sessions.values(),
            key=lambda x: x.updated_at,
            reverse=True,
        )[:limit]
        out = []
        for s in items:
            cnt = len(self._rows.get(s.id, []))
            out.append(
                Session(
                    id=s.id,
                    title=s.title,
                    messages=[],
                    model=s.model,
                    workdir=s.workdir,
                    created_at=s.created_at,
                    updated_at=s.updated_at,
                    message_count=cnt,
                )
            )
        return out

    def delete_session(self, session_id: str) -> bool:
        ok = session_id in self._sessions
        self._sessions.pop(session_id, None)
        self._rows.pop(session_id, None)
        return ok

    def get_latest_session(self, workdir: str | None = None) -> Session | None:
        cand = [
            s
            for s in self._sessions.values()
            if workdir is None or s.workdir == workdir
        ]
        if not cand:
            return None
        latest = max(cand, key=lambda x: x.updated_at)
        return self.get_session(latest.id)

    def get_message_count(self, session_id: str) -> int:
        return len(self._rows.get(session_id, []))

    def update_title(self, session_id: str, title: str) -> None:
        if s := self._sessions.get(session_id):
            s.title = title

    def save_compaction(
        self, session_id: str, summary: str, up_to_message_id: int
    ) -> None:
        if s := self._sessions.get(session_id):
            s.compacted_summary = summary
            s.compacted_up_to = up_to_message_id
            s.updated_at = datetime.now(timezone.utc).isoformat()

    def get_messages_after(
        self, session_id: str, after_id: int
    ) -> list[Message]:
        return [m for mid, m in self._rows.get(session_id, []) if mid > after_id]

    def get_last_message_id(self, session_id: str) -> int | None:
        rows = self._rows.get(session_id, [])
        if not rows:
            return None
        return max(mid for mid, _ in rows)


async def main() -> None:
    store = PostgresStyleSessionStore("postgresql://user:pass@localhost/cody")
    client = (
        Cody()
        .workdir(".")
        .session_store(store)
        .build()
    )
    async with client:
        r = await client.run("用一句话说明当前目录下有哪些文件")
        print(r.output)


if __name__ == "__main__":
    asyncio.run(main())

from cody.core.storage import (
    SessionStoreProtocol,
    AuditLoggerProtocol,
    FileHistoryProtocol,
)

client = (
    Cody()
    .session_store(my_session_store)
    .audit_logger(my_audit_logger)
    .file_history(my_file_history)
    .build()
)

client = Cody().stateless().build()
client = Cody().stateless().audit_logger(real_logger).build()  # 仅覆盖审计

import asyncio
from cody.sdk import (
    CodyError, CodyModelError, CodyToolError, CodyPermissionError,
    CodyNotFoundError, CodyRateLimitError, CodyConfigError,
    CodyTimeoutError, CodyConnectionError, CodySessionError,
)

try:
    result = await client.run("任务")
except CodyRateLimitError as e:
    await asyncio.sleep(e.retry_after)
    result = await client.run("任务")
except CodyToolError as e:
    print(f"工具 {e.details['tool_name']} 失败: {e.message}")
except CodyPermissionError as e:
    print(f"权限拒绝: {e.message}")
except CodyError as e:
    print(f"[{e.code}] {e.message}")