fableforge-agent-runtime 0.1.0__py3-none-any.whl

agent_runtime/server.py

@@ -0,0 +1,119 @@
+"""FastAPI server — HTTP API for agent runtime."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Any
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+
+from .daemon import AgentDaemon
+from .models import (
+    CheckpointInfo,
+    HealthResponse,
+    MemoryRetrieveResponse,
+    MemorySearchResult,
+    MemoryStoreRequest,
+    SessionCreate,
+    SessionResume,
+    SessionState,
+)
+
+
+def create_app(daemon: AgentDaemon) -> FastAPI:
+    app = FastAPI(title="AgentRuntime", version="0.1.0")
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+    sm = daemon.session_manager
+
+    @app.get("/health", response_model=HealthResponse)
+    async def health():
+        return daemon.health()
+
+    @app.post("/sessions", response_model=SessionState, status_code=201)
+    async def create_session(body: SessionCreate):
+        state = sm.create_session(
+            name=body.name,
+            model=body.model,
+            config={
+                "system_prompt": body.system_prompt,
+                "tools": body.tools,
+            },
+        )
+        return state
+
+    @app.get("/sessions", response_model=list[SessionState])
+    async def list_sessions():
+        return sm.list_sessions()
+
+    @app.get("/sessions/{session_id}", response_model=SessionState)
+    async def get_session(session_id: str):
+        try:
+            return sm.get_session_status(session_id)
+        except Exception:
+            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
+
+    @app.post("/sessions/{session_id}/start", response_model=SessionState)
+    async def start_session(session_id: str):
+        try:
+            return await sm.start_session(session_id)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    @app.post("/sessions/{session_id}/pause", response_model=SessionState)
+    async def pause_session(session_id: str):
+        try:
+            return await sm.pause_session(session_id)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    @app.post("/sessions/{session_id}/resume", response_model=SessionState)
+    async def resume_session(session_id: str, body: SessionResume | None = None):
+        checkpoint_id = body.checkpoint_id if body else None
+        try:
+            return await sm.resume_session(session_id, checkpoint_id=checkpoint_id)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    @app.post("/sessions/{session_id}/stop", response_model=SessionState)
+    async def stop_session(session_id: str):
+        try:
+            return await sm.stop_session(session_id)
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    @app.get("/sessions/{session_id}/memory")
+    async def get_memory(session_id: str, key: str | None = None):
+        mem = sm.get_memory_store()
+        if key:
+            entry = mem.retrieve(key)
+            if entry is None:
+                raise HTTPException(status_code=404, detail=f"Memory key {key} not found")
+            return MemoryRetrieveResponse(key=entry.key, value=entry.value, timestamp=entry.timestamp)
+        return {"keys": mem.list_keys()}
+
+    @app.post("/sessions/{session_id}/memory", status_code=201)
+    async def store_memory(session_id: str, body: MemoryStoreRequest):
+        mem = sm.get_memory_store()
+        mem.store(key=body.key, value=body.value)
+        return {"status": "stored", "key": body.key}
+
+    @app.get("/sessions/{session_id}/checkpoints", response_model=list[CheckpointInfo])
+    async def list_checkpoints(session_id: str):
+        serializer = sm.get_state_serializer()
+        return serializer.list_checkpoints(session_id)
+
+    @app.post("/sessions/{session_id}/checkpoints", response_model=CheckpointInfo, status_code=201)
+    async def create_checkpoint(session_id: str, label: str | None = None):
+        serializer = sm.get_state_serializer()
+        try:
+            return serializer.create_checkpoint(session_id, label=label)
+        except ValueError as e:
+            raise HTTPException(status_code=404, detail=str(e))
+
+    return app

agent_runtime/session_manager.py

@@ -0,0 +1,156 @@
+"""Session lifecycle management — create, start, pause, resume, stop."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import Any
+
+from .memory_store import MemoryStore
+from .models import SessionCreate, SessionState, SessionStatus
+from .state_serializer import StateSerializer
+
+logger = logging.getLogger(__name__)
+
+
+class SessionError(Exception):
+    pass
+
+
+class SessionManager:
+    def __init__(
+        self,
+        state_serializer: StateSerializer | None = None,
+        memory_store: MemoryStore | None = None,
+    ) -> None:
+        self._state_serializer = state_serializer or StateSerializer()
+        self._memory_store = memory_store or MemoryStore()
+        self._running_tasks: dict[str, asyncio.Task] = {}
+        self._heartbeat_intervals: dict[str, float] = {}
+        self._auto_checkpoint_interval = 60.0
+
+    def create_session(self, name: str, model: str, config: dict[str, Any] | None = None) -> SessionState:
+        cfg = config or {}
+        state = SessionState(
+            session_id=uuid.uuid4().hex[:16],
+            name=name,
+            model=model,
+            system_prompt=cfg.get("system_prompt", ""),
+            tools=cfg.get("tools", []),
+        )
+        self._state_serializer.save_session(state)
+        logger.info("Created session %s (%s)", state.session_id, name)
+        return state
+
+    async def start_session(self, session_id: str) -> SessionState:
+        state = self._state_serializer.load_session(session_id)
+        if state is None:
+            raise SessionError(f"Session {session_id} not found")
+        if state.status == SessionStatus.RUNNING:
+            raise SessionError(f"Session {session_id} is already running")
+        state.status = SessionStatus.RUNNING
+        state.updated_at = datetime.now(timezone.utc)
+        self._state_serializer.save_session(state)
+        task = asyncio.create_task(self._session_loop(session_id))
+        self._running_tasks[session_id] = task
+        logger.info("Started session %s", session_id)
+        return state
+
+    async def pause_session(self, session_id: str) -> SessionState:
+        state = self._state_serializer.load_session(session_id)
+        if state is None:
+            raise SessionError(f"Session {session_id} not found")
+        if state.status != SessionStatus.RUNNING:
+            raise SessionError(f"Session {session_id} is not running")
+        state.status = SessionStatus.PAUSED
+        state.updated_at = datetime.now(timezone.utc)
+        self._state_serializer.save_session(state)
+        task = self._running_tasks.pop(session_id, None)
+        if task and not task.done():
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+        self._state_serializer.create_checkpoint(session_id, label="pause-checkpoint")
+        logger.info("Paused session %s", session_id)
+        return state
+
+    async def resume_session(self, session_id: str, checkpoint_id: str | None = None) -> SessionState:
+        state = self._state_serializer.load_session(session_id)
+        if state is None:
+            raise SessionError(f"Session {session_id} not found")
+        if checkpoint_id:
+            state = self._state_serializer.resume_from_checkpoint(session_id, checkpoint_id)
+        state.status = SessionStatus.RUNNING
+        state.updated_at = datetime.now(timezone.utc)
+        self._state_serializer.save_session(state)
+        task = asyncio.create_task(self._session_loop(session_id))
+        self._running_tasks[session_id] = task
+        logger.info("Resumed session %s", session_id)
+        return state
+
+    async def stop_session(self, session_id: str) -> SessionState:
+        state = self._state_serializer.load_session(session_id)
+        if state is None:
+            raise SessionError(f"Session {session_id} not found")
+        state.status = SessionStatus.STOPPED
+        state.updated_at = datetime.now(timezone.utc)
+        self._state_serializer.save_session(state)
+        task = self._running_tasks.pop(session_id, None)
+        if task and not task.done():
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+        self._state_serializer.create_checkpoint(session_id, label="stop-checkpoint")
+        logger.info("Stopped session %s", session_id)
+        return state
+
+    def list_sessions(self) -> list[SessionState]:
+        return self._state_serializer.list_sessions()
+
+    def get_session_status(self, session_id: str) -> SessionState:
+        state = self._state_serializer.load_session(session_id)
+        if state is None:
+            raise SessionError(f"Session {session_id} not found")
+        return state
+
+    def get_memory_store(self) -> MemoryStore:
+        return self._memory_store
+
+    def get_state_serializer(self) -> StateSerializer:
+        return self._state_serializer
+
+    async def _session_loop(self, session_id: str) -> None:
+        last_checkpoint = asyncio.get_event_loop().time()
+        try:
+            while True:
+                await asyncio.sleep(1.0)
+                now = asyncio.get_event_loop().time()
+                if now - last_checkpoint >= self._auto_checkpoint_interval:
+                    state = self._state_serializer.load_session(session_id)
+                    if state and state.status == SessionStatus.RUNNING:
+                        self._state_serializer.create_checkpoint(
+                            session_id, label="auto-checkpoint"
+                        )
+                        self._state_serializer.prune_checkpoints(session_id, keep_last=10)
+                    last_checkpoint = now
+        except asyncio.CancelledError:
+            logger.info("Session loop cancelled for %s", session_id)
+
+    async def shutdown(self) -> None:
+        logger.info("Shutting down session manager, preserving %d active sessions", len(self._running_tasks))
+        for session_id, task in list(self._running_tasks.items()):
+            if not task.done():
+                task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+            self._state_serializer.create_checkpoint(session_id, label="shutdown-checkpoint")
+        self._running_tasks.clear()
+        logger.info("All sessions preserved and shut down")

agent_runtime/state_serializer.py

@@ -0,0 +1,222 @@
+"""Serialize and deserialize agent state to/from SQLite."""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from sqlalchemy import (
+    Column,
+    DateTime,
+    Index,
+    LargeBinary,
+    String,
+    Text,
+    create_engine,
+    text,
+)
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+from .models import (
+    CheckpointInfo,
+    Message,
+    SessionState,
+    SessionStatus,
+    ToolCall,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_DB_PATH = Path.home() / ".agent_runtime" / "state.db"
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class SessionRow(Base):
+    __tablename__ = "sessions"
+
+    session_id = Column(String(64), primary_key=True)
+    name = Column(String(256), default="")
+    model = Column(String(128), default="")
+    system_prompt = Column(Text, default="")
+    tools = Column(Text, default="[]")
+    status = Column(String(32), default=SessionStatus.CREATED.value)
+    memory_json = Column(Text, default="{}")
+    messages_json = Column(Text, default="[]")
+    tool_history_json = Column(Text, default="[]")
+    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
+    updated_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
+
+
+class CheckpointRow(Base):
+    __tablename__ = "checkpoints"
+
+    checkpoint_id = Column(String(64), primary_key=True)
+    session_id = Column(String(64), nullable=False, index=True)
+    label = Column(String(256), nullable=True)
+    state_json = Column(Text, nullable=False)
+    created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
+
+    __table_args__ = (Index("ix_checkpoints_session_created", "session_id", "created_at"),)
+
+
+class StateSerializer:
+    def __init__(self, db_path: Path | str | None = None) -> None:
+        if db_path is None:
+            db_path = DEFAULT_DB_PATH
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self.engine = create_engine(f"sqlite:///{self.db_path}", echo=False)
+        Base.metadata.create_all(self.engine)
+        self._session_factory = sessionmaker(bind=self.engine)
+
+    def _get_session(self) -> Session:
+        return self._session_factory()
+
+    def save_session(self, state: SessionState) -> None:
+        with self._get_session() as db:
+            row = db.get(SessionRow, state.session_id)
+            if row is None:
+                row = SessionRow(session_id=state.session_id)
+                db.add(row)
+            row.name = state.name
+            row.model = state.model
+            row.system_prompt = state.system_prompt
+            row.tools = json.dumps(state.tools)
+            row.status = state.status.value
+            row.memory_json = json.dumps(state.memory, default=str)
+            row.messages_json = json.dumps(
+                [m.model_dump(mode="json") for m in state.messages],
+                default=str,
+            )
+            row.tool_history_json = json.dumps(
+                [t.model_dump(mode="json") for t in state.tool_history],
+                default=str,
+            )
+            row.created_at = state.created_at
+            row.updated_at = datetime.now(timezone.utc)
+            db.commit()
+        logger.info("Saved session %s", state.session_id)
+
+    def load_session(self, session_id: str) -> SessionState | None:
+        with self._get_session() as db:
+            row = db.get(SessionRow, session_id)
+            if row is None:
+                return None
+        return self._row_to_state(row)
+
+    def delete_session(self, session_id: str) -> bool:
+        with self._get_session() as db:
+            row = db.get(SessionRow, session_id)
+            if row is None:
+                return False
+            db.delete(row)
+            db.execute(
+                text("DELETE FROM checkpoints WHERE session_id = :sid"),
+                {"sid": session_id},
+            )
+            db.commit()
+        return True
+
+    def list_sessions(self) -> list[SessionState]:
+        with self._get_session() as db:
+            rows = db.query(SessionRow).order_by(SessionRow.created_at).all()
+        return [self._row_to_state(r) for r in rows]
+
+    def create_checkpoint(self, session_id: str, label: str | None = None) -> CheckpointInfo:
+        state = self.load_session(session_id)
+        if state is None:
+            raise ValueError(f"Session {session_id} not found")
+        checkpoint_id = uuid.uuid4().hex[:16]
+        state_json = json.dumps(state.model_dump(mode="json"), default=str)
+        with self._get_session() as db:
+            row = CheckpointRow(
+                checkpoint_id=checkpoint_id,
+                session_id=session_id,
+                label=label,
+                state_json=state_json,
+            )
+            db.add(row)
+            db.commit()
+            created_at = row.created_at
+        logger.info("Created checkpoint %s for session %s", checkpoint_id, session_id)
+        return CheckpointInfo(
+            checkpoint_id=checkpoint_id,
+            session_id=session_id,
+            created_at=created_at,
+            label=label,
+        )
+
+    def list_checkpoints(self, session_id: str) -> list[CheckpointInfo]:
+        with self._get_session() as db:
+            rows = (
+                db.query(CheckpointRow)
+                .filter(CheckpointRow.session_id == session_id)
+                .order_by(CheckpointRow.created_at)
+                .all()
+            )
+        return [
+            CheckpointInfo(
+                checkpoint_id=r.checkpoint_id,
+                session_id=r.session_id,
+                created_at=r.created_at,
+                label=r.label,
+            )
+            for r in rows
+        ]
+
+    def resume_from_checkpoint(self, session_id: str, checkpoint_id: str) -> SessionState:
+        with self._get_session() as db:
+            row = db.get(CheckpointRow, checkpoint_id)
+            if row is None or row.session_id != session_id:
+                raise ValueError(f"Checkpoint {checkpoint_id} not found for session {session_id}")
+            state_data = json.loads(row.state_json)
+        state = SessionState.model_validate(state_data)
+        state.status = SessionStatus.PAUSED
+        self.save_session(state)
+        logger.info("Resumed session %s from checkpoint %s", session_id, checkpoint_id)
+        return state
+
+    def prune_checkpoints(
+        self, session_id: str, keep_last: int = 10
+    ) -> int:
+        with self._get_session() as db:
+            rows = (
+                db.query(CheckpointRow)
+                .filter(CheckpointRow.session_id == session_id)
+                .order_by(CheckpointRow.created_at.desc())
+                .all()
+            )
+            to_delete = rows[keep_last:]
+            count = 0
+            for row in to_delete:
+                db.delete(row)
+                count += 1
+            db.commit()
+        if count:
+            logger.info("Pruned %d old checkpoints for session %s", count, session_id)
+        return count
+
+    @staticmethod
+    def _row_to_state(row: SessionRow) -> SessionState:
+        messages = [Message.model_validate(m) for m in json.loads(row.messages_json)]
+        tool_history = [ToolCall.model_validate(t) for t in json.loads(row.tool_history_json)]
+        return SessionState(
+            session_id=row.session_id,
+            name=row.name,
+            model=row.model,
+            system_prompt=row.system_prompt,
+            tools=json.loads(row.tools),
+            status=SessionStatus(row.status),
+            memory=json.loads(row.memory_json),
+            messages=messages,
+            tool_history=tool_history,
+            created_at=row.created_at,
+            updated_at=row.updated_at,
+        )