feat: integrate API key bridging and enhance checkpointer functionality

- Introduced `bridge_provider_api_keys` function to map MISAKA_* API keys into standard environment variables for LangChain, ensuring existing keys are not overwritten.
- Enhanced checkpointer setup and teardown processes, improving error handling and logging during initialization and closure.
- Updated main application lifecycle to include API key bridging during startup.
- Added tests for the new API key bridging functionality to ensure correct behavior and environment variable management.

Author: knqiufan

agent/app/config.py

@@ -2,12 +2,16 @@
 
 from __future__ import annotations
 
+import logging
+import os
 from functools import lru_cache
 from pathlib import Path
 
 from pydantic import Field
 from pydantic_settings import BaseSettings
 
+logger = logging.getLogger(__name__)
+
 
 def _default_interrupt_config() -> dict[str, bool]:
     return {
@@ -57,4 +61,22 @@ def get_settings() -> Settings:
     return Settings()
 
 
+def bridge_provider_api_keys(settings_obj: Settings | None = None) -> None:
+    """Bridge MISAKA_* keys into standard provider env vars used by LangChain.
+
+    Does not overwrite keys that are already present in the process environment.
+    Never logs key values.
+    """
+    cfg = settings_obj or get_settings()
+    bridged: list[str] = []
+    if cfg.anthropic_api_key and not os.environ.get("ANTHROPIC_API_KEY"):
+        os.environ["ANTHROPIC_API_KEY"] = cfg.anthropic_api_key
+        bridged.append("ANTHROPIC_API_KEY")
+    if cfg.openai_api_key and not os.environ.get("OPENAI_API_KEY"):
+        os.environ["OPENAI_API_KEY"] = cfg.openai_api_key
+        bridged.append("OPENAI_API_KEY")
+    if bridged:
+        logger.info("Bridged provider API keys into process env: %s", ", ".join(bridged))
+
+
 settings = Settings()

agent/app/dependencies.py

@@ -11,6 +11,7 @@
 logger = logging.getLogger(__name__)
 
 _checkpointer: Any | None = None
+_checkpointer_cm: Any | None = None
 _store: Any | None = None
 
 
@@ -32,10 +33,15 @@ def get_store():
 
 
 def get_checkpointer():
-    """Return the AsyncSqliteSaver singleton, creating it lazily."""
-    global _checkpointer
+    """Return the process-wide AsyncSqliteSaver, or None if unavailable/uninitialized."""
+    return _checkpointer
+
+
+async def setup_checkpointer() -> None:
+    """Enter AsyncSqliteSaver context, create schema, and enable WAL via setup()."""
+    global _checkpointer, _checkpointer_cm
     if _checkpointer is not None:
-        return _checkpointer
+        return
 
     settings = get_settings()
     settings.data_dir.mkdir(parents=True, exist_ok=True)
@@ -44,40 +50,35 @@ def get_checkpointer():
         from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
     except ImportError:
         logger.warning("AsyncSqliteSaver unavailable; continuing without checkpointer")
-        return None
+        return
 
     conn_string = str(settings.checkpointer_db_path)
-    _checkpointer = AsyncSqliteSaver.from_conn_string(conn_string)
-    return _checkpointer
-
-
-async def setup_checkpointer() -> None:
-    """Ensure checkpointer schema exists."""
-    checkpointer = get_checkpointer()
-    if checkpointer is None:
-        return
-    setup = getattr(checkpointer, "setup", None)
-    if setup is None:
+    cm = AsyncSqliteSaver.from_conn_string(conn_string)
+    try:
+        saver = await cm.__aenter__()
+        await saver.setup()
+    except Exception:
+        logger.exception("Failed to initialize checkpointer")
+        try:
+            await cm.__aexit__(None, None, None)
+        except Exception:
+            logger.exception("Failed to clean up checkpointer after setup error")
         return
-    result = setup()
-    if hasattr(result, "__await__"):
-        await result
 
+    _checkpointer_cm = cm
+    _checkpointer = saver
+    logger.info("Checkpointer ready at %s (WAL enabled via setup)", conn_string)
 
-async def close_checkpointer() -> None:
-    """Close checkpointer resources if present."""
-    global _checkpointer
-    checkpointer = _checkpointer
-    if checkpointer is None:
-        return
-
-    for method_name in ("aclose", "close"):
-        method = getattr(checkpointer, method_name, None)
-        if method is None:
-            continue
-        result = method()
-        if hasattr(result, "__await__"):
-            await result
-        break
 
+async def close_checkpointer() -> None:
+    """Exit the AsyncSqliteSaver context manager and clear the singleton."""
+    global _checkpointer, _checkpointer_cm
+    cm = _checkpointer_cm
     _checkpointer = None
+    _checkpointer_cm = None
+    if cm is None:
+        return
+    try:
+        await cm.__aexit__(None, None, None)
+    except Exception:
+        logger.exception("Failed to close checkpointer")

agent/app/main.py

@@ -5,17 +5,19 @@
 
 from fastapi import FastAPI
 
-from app.config import settings
+from app.config import bridge_provider_api_keys, settings
 from app.dependencies import close_checkpointer, setup_checkpointer
 from app.routers.agent import router as agent_router
 from app.routers.health import router as health_router
 from app.routers.info import router as info_router
+from app.routers.memory import router as memory_router
 
 
 @asynccontextmanager
 async def lifespan(application: FastAPI):
     """Manage startup and shutdown lifecycle."""
     application.state.startup_time = time.time()
+    bridge_provider_api_keys()
     await setup_checkpointer()
     try:
         yield
@@ -34,3 +36,4 @@ async def lifespan(application: FastAPI):
 app.include_router(health_router)
 app.include_router(info_router)
 app.include_router(agent_router)
+app.include_router(memory_router)

agent/app/routers/memory.py

@@ -0,0 +1,149 @@
+"""PowerMem REST endpoints for search / list / delete."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, Query
+from pydantic import BaseModel, Field
+
+from app.memory import get_memory_engine
+
+router = APIRouter(prefix="/memory", tags=["memory"])
+logger = logging.getLogger(__name__)
+
+
+class MemoryItem(BaseModel):
+    id: str | None = None
+    content: str = ""
+    score: float | None = None
+    created_at: str | None = None
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+
+class MemoryListResponse(BaseModel):
+    items: list[MemoryItem]
+    offset: int = 0
+    limit: int = 20
+
+
+def _require_engine() -> Any:
+    engine = get_memory_engine()
+    if engine is None:
+        raise HTTPException(status_code=503, detail="Memory engine unavailable")
+    return engine
+
+
+def _normalize_item(item: Any) -> MemoryItem:
+    if isinstance(item, dict):
+        return MemoryItem(
+            id=_as_optional_str(item.get("id") or item.get("memory_id")),
+            content=str(item.get("content") or item.get("memory") or ""),
+            score=_as_optional_float(item.get("score")),
+            created_at=_as_optional_str(item.get("created_at")),
+            metadata=item.get("metadata") if isinstance(item.get("metadata"), dict) else {},
+        )
+
+    return MemoryItem(
+        id=_as_optional_str(getattr(item, "id", None) or getattr(item, "memory_id", None)),
+        content=str(
+            getattr(item, "content", None) or getattr(item, "memory", None) or ""
+        ),
+        score=_as_optional_float(getattr(item, "score", None)),
+        created_at=_as_optional_str(getattr(item, "created_at", None)),
+        metadata=getattr(item, "metadata", None)
+        if isinstance(getattr(item, "metadata", None), dict)
+        else {},
+    )
+
+
+def _as_optional_str(value: Any) -> str | None:
+    if value is None:
+        return None
+    return str(value)
+
+
+def _as_optional_float(value: Any) -> float | None:
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def _call_search(engine: Any, query: str, limit: int) -> list[Any]:
+    if hasattr(engine, "search"):
+        result = engine.search(query, limit=limit)
+    elif hasattr(engine, "query"):
+        result = engine.query(query, limit=limit)
+    else:
+        raise HTTPException(status_code=503, detail="Memory engine has no search API")
+    return list(result or [])
+
+
+def _call_list(engine: Any, offset: int, limit: int) -> list[Any]:
+    if hasattr(engine, "list"):
+        result = engine.list(offset=offset, limit=limit)
+    elif hasattr(engine, "get_all"):
+        result = engine.get_all()
+        result = list(result or [])[offset : offset + limit]
+    else:
+        raise HTTPException(status_code=503, detail="Memory engine has no list API")
+    return list(result or [])
+
+
+def _call_delete(engine: Any, memory_id: str) -> None:
+    if hasattr(engine, "delete"):
+        engine.delete(memory_id)
+        return
+    if hasattr(engine, "remove"):
+        engine.remove(memory_id)
+        return
+    raise HTTPException(status_code=503, detail="Memory engine has no delete API")
+
+
+@router.get("/search", response_model=MemoryListResponse)
+async def memory_search(
+    query: str = Query(..., min_length=1),
+    limit: int = Query(5, ge=1, le=100),
+) -> MemoryListResponse:
+    engine = _require_engine()
+    try:
+        items = [_normalize_item(item) for item in _call_search(engine, query, limit)]
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.exception("memory search failed")
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+    return MemoryListResponse(items=items, offset=0, limit=limit)
+
+
+@router.get("/list", response_model=MemoryListResponse)
+async def memory_list(
+    offset: int = Query(0, ge=0),
+    limit: int = Query(20, ge=1, le=100),
+) -> MemoryListResponse:
+    engine = _require_engine()
+    try:
+        items = [_normalize_item(item) for item in _call_list(engine, offset, limit)]
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.exception("memory list failed")
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+    return MemoryListResponse(items=items, offset=offset, limit=limit)
+
+
+@router.delete("/{memory_id}")
+async def memory_delete(memory_id: str) -> dict[str, str]:
+    engine = _require_engine()
+    try:
+        _call_delete(engine, memory_id)
+    except HTTPException:
+        raise
+    except Exception as exc:
+        logger.exception("memory delete failed")
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+    return {"status": "deleted", "id": memory_id}

agent/tests/test_config.py

@@ -59,3 +59,22 @@ def test_get_settings_returns_singleton():
     second = get_settings()
     assert first is second
     get_settings.cache_clear()
+
+
+def test_bridge_provider_api_keys(monkeypatch):
+    from app.config import Settings, bridge_provider_api_keys
+
+    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    cfg = Settings(
+        anthropic_api_key="sk-ant-test",
+        openai_api_key="sk-openai-test",
+    )
+    bridge_provider_api_keys(cfg)
+    assert __import__("os").environ["ANTHROPIC_API_KEY"] == "sk-ant-test"
+    assert __import__("os").environ["OPENAI_API_KEY"] == "sk-openai-test"
+
+    # Does not overwrite existing values.
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "keep-me")
+    bridge_provider_api_keys(cfg)
+    assert __import__("os").environ["ANTHROPIC_API_KEY"] == "keep-me"