Implement compliant SessionStore API endpoints, completions concurrency limit, mirror_error alert, and nightly DB retention policy

backend.py

@@ -23,7 +23,8 @@ import time
 import re
 import collections
 from pathlib import Path
-from typing import AsyncIterator, Optional
+from typing import AsyncIterator, Optional, List, Dict, Any
+from pydantic import BaseModel
 
 from fastapi import FastAPI, Request, Header, HTTPException
 from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
@@ -364,16 +365,16 @@ async def init_db():
         )
         async with db_pool.acquire() as conn:
             await conn.execute("""
-                CREATE TABLE IF NOT EXISTS agent_sessions (
-                    id BIGSERIAL PRIMARY KEY,
-                    session_id TEXT NOT NULL,
-                    role TEXT NOT NULL,
-                    content TEXT,
-                    tool_calls JSONB,
-                    tool_call_id TEXT,
-                    created_at TIMESTAMPTZ DEFAULT NOW()
+                CREATE TABLE IF NOT EXISTS agent_session_entries (
+                    id          BIGSERIAL PRIMARY KEY,
+                    project_key TEXT        NOT NULL,
+                    session_id  TEXT        NOT NULL,
+                    subpath     TEXT,
+                    entry       JSONB       NOT NULL,
+                    created_at  TIMESTAMPTZ NOT NULL DEFAULT now()
                 );
-                CREATE INDEX IF NOT EXISTS idx_agent_sessions_sid ON agent_sessions(session_id);
+                CREATE INDEX IF NOT EXISTS idx_session_key ON agent_session_entries (project_key, session_id, subpath, id);
+                CREATE INDEX IF NOT EXISTS idx_project_session ON agent_session_entries (project_key, session_id);
             """)
     except Exception as e:
         print(f"[DB] Warning: Could not initialize database: {e}")
@@ -382,42 +383,43 @@ async def init_db():
 
 async def save_message(session_id: str, role: str, content: str = None, 
                        tool_calls: list = None, tool_call_id: str = None):
-    """Save a message to the session store."""
+    """Save a message to the session store using the unified schema."""
     if not db_pool:
         return
+    msg = {"role": role}
+    if content is not None:
+        msg["content"] = content
+    if tool_calls:
+        msg["tool_calls"] = tool_calls
+    if tool_call_id:
+        msg["tool_call_id"] = tool_call_id
+        
     try:
         async with db_pool.acquire() as conn:
             await conn.execute(
-                "INSERT INTO agent_sessions (session_id, role, content, tool_calls, tool_call_id) VALUES ($1, $2, $3, $4, $5)",
-                session_id, role, content,
-                json.dumps(tool_calls) if tool_calls else None,
-                tool_call_id,
+                "INSERT INTO agent_session_entries (project_key, session_id, subpath, entry) VALUES ($1, $2, $3, $4)",
+                "fastapi-completions",
+                session_id,
+                None,
+                json.dumps(msg)
             )
     except Exception as e:
         print(f"[DB] Warning: Could not save message: {e}")
 
 
 async def load_session(session_id: str) -> list:
-    """Load conversation history from the session store."""
+    """Load conversation history from the session store using the unified schema."""
     if not db_pool:
         return []
     try:
         async with db_pool.acquire() as conn:
             rows = await conn.fetch(
-                "SELECT role, content, tool_calls, tool_call_id FROM agent_sessions WHERE session_id = $1 ORDER BY id",
+                "SELECT entry FROM agent_session_entries WHERE project_key = $1 AND session_id = $2 AND subpath IS NOT DISTINCT FROM $3 ORDER BY id",
+                "fastapi-completions",
                 session_id,
+                None
             )
-            messages = []
-            for row in rows:
-                msg = {"role": row["role"]}
-                if row["content"]:
-                    msg["content"] = row["content"]
-                if row["tool_calls"]:
-                    msg["tool_calls"] = json.loads(row["tool_calls"])
-                if row["tool_call_id"]:
-                    msg["tool_call_id"] = row["tool_call_id"]
-                messages.append(msg)
-            return messages
+            return [json.loads(row["entry"]) for row in rows]
     except Exception as e:
         print(f"[DB] Warning: Could not load session: {e}")
         return []
@@ -621,7 +623,8 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
             if is_agentic:
                 kwargs["tools"] = TOOLS
                 kwargs["tool_choice"] = "auto"
-            response = await nim_client.chat.completions.create(**kwargs)
+            async with completions_semaphore:
+                response = await nim_client.chat.completions.create(**kwargs)
             content = response.choices[0].message.content or ""
             await save_message(session_id, "assistant", content)
             ACTIVE_SESSIONS.discard(session_id)
@@ -640,50 +643,51 @@ async def chat_completions(request: Request, authorization: str = Header(None)):
     # Streaming + agentic loop
     async def generate() -> AsyncIterator[str]:
         nonlocal final_messages
-        try:
-            for round_num in range(MAX_TOOL_ROUNDS + 1):
-                kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
-                if is_agentic:
-                    kwargs["tools"] = TOOLS
-                    kwargs["tool_choice"] = "auto"
-
-                # Collect streamed response
-                full_content = ""
-                tool_calls_raw = {}  # index -> {id, name, arguments_str}
-                
-                async for chunk in await nim_client.chat.completions.create(**kwargs):
-                    choice = chunk.choices[0] if chunk.choices else None
-                    if not choice:
-                        continue
-                    delta = choice.delta
-                    
-                    # Stream text content to client
-                    if delta and delta.content:
-                        full_content += delta.content
-                        yield make_chunk(request_id, requested_model, delta.content)
-                    
-                    # Collect tool calls
-                    if delta and delta.tool_calls:
-                        for tc in delta.tool_calls:
-                            idx = tc.index
-                            if idx not in tool_calls_raw:
-                                tool_calls_raw[idx] = {
-                                    "id": tc.id or f"call_{uuid.uuid4().hex[:8]}",
-                                    "name": tc.function.name if tc.function and tc.function.name else "",
-                                    "arguments": ""
-                                }
-                            if tc.function and tc.function.name:
-                                tool_calls_raw[idx]["name"] = tc.function.name
-                            if tc.id:
-                                tool_calls_raw[idx]["id"] = tc.id
-                            if tc.function and tc.function.arguments:
-                                tool_calls_raw[idx]["arguments"] += tc.function.arguments
+        async with completions_semaphore:
+            try:
+                for round_num in range(MAX_TOOL_ROUNDS + 1):
+                    kwargs = {"model": requested_model, "messages": final_messages, "stream": True}
+                    if is_agentic:
+                        kwargs["tools"] = TOOLS
+                        kwargs["tool_choice"] = "auto"
+
+                    # Collect streamed response
+                    full_content = ""
+                    tool_calls_raw = {}  # index -> {id, name, arguments_str}
                     
-                    # Check for finish
-                    if choice.finish_reason == "stop":
-                        break
-                    if choice.finish_reason == "tool_calls":
-                        break
+                    async for chunk in await nim_client.chat.completions.create(**kwargs):
+                        choice = chunk.choices[0] if chunk.choices else None
+                        if not choice:
+                            continue
+                        delta = choice.delta
+                        
+                        # Stream text content to client
+                        if delta and delta.content:
+                            full_content += delta.content
+                            yield make_chunk(request_id, requested_model, delta.content)
+                        
+                        # Collect tool calls
+                        if delta and delta.tool_calls:
+                            for tc in delta.tool_calls:
+                                idx = tc.index
+                                if idx not in tool_calls_raw:
+                                    tool_calls_raw[idx] = {
+                                        "id": tc.id or f"call_{uuid.uuid4().hex[:8]}",
+                                        "name": tc.function.name if tc.function and tc.function.name else "",
+                                        "arguments": ""
+                                    }
+                                if tc.function and tc.function.name:
+                                    tool_calls_raw[idx]["name"] = tc.function.name
+                                if tc.id:
+                                    tool_calls_raw[idx]["id"] = tc.id
+                                if tc.function and tc.function.arguments:
+                                    tool_calls_raw[idx]["arguments"] += tc.function.arguments
+                        
+                        # Check for finish
+                        if choice.finish_reason == "stop":
+                            break
+                        if choice.finish_reason == "tool_calls":
+                            break
 
                 # If no tool calls, we're done
                 if not tool_calls_raw:
@@ -1176,6 +1180,123 @@ Select a file from the sidebar explorer on the left to read its code contents in
 </html>
 """
 
+# ---------------------------------------------------------------------------
+# Concurrency Settings
+# ---------------------------------------------------------------------------
+completions_semaphore = asyncio.Semaphore(2)
+
+
+# ---------------------------------------------------------------------------
+# SessionStore API (Claude Agent SDK / Claude Code Compatible)
+# ---------------------------------------------------------------------------
+
+class SessionAppendRequest(BaseModel):
+    project_key: str
+    session_id: str
+    subpath: Optional[str] = None
+    entries: List[Dict[str, Any]]
+
+@app.post("/api/sessions/append")
+async def append_session_entries(req: SessionAppendRequest, authorization: str = Header(None)):
+    auth(authorization)
+    if not db_pool:
+        raise HTTPException(status_code=500, detail="Database not connected")
+    try:
+        # Check for mirror_error and log/alert if present
+        for e in req.entries:
+            if e.get("type") == "system" and e.get("subtype") == "mirror_error":
+                log_activity(f"[ALERT] Claude Agent SDK reported mirror_error: {e.get('message')}")
+                
+        rows = [(req.project_key, req.session_id, req.subpath, json.dumps(e)) for e in req.entries]
+        async with db_pool.acquire() as conn:
+            await conn.executemany(
+                "INSERT INTO agent_session_entries (project_key, session_id, subpath, entry) VALUES ($1, $2, $3, $4)",
+                rows
+            )
+        return {"status": "success"}
+    except Exception as e:
+        log_activity(f"[SessionStore Error] Failed append: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+class SessionLoadRequest(BaseModel):
+    project_key: str
+    session_id: str
+    subpath: Optional[str] = None
+
+@app.post("/api/sessions/load")
+async def load_session_entries(req: SessionLoadRequest, authorization: str = Header(None)):
+    auth(authorization)
+    if not db_pool:
+        raise HTTPException(status_code=500, detail="Database not connected")
+    try:
+        async with db_pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT entry FROM agent_session_entries WHERE project_key=$1 AND session_id=$2 AND subpath IS NOT DISTINCT FROM $3 ORDER BY id",
+                req.project_key, req.session_id, req.subpath
+            )
+        return {"entries": [json.loads(r["entry"]) for r in rows]}
+    except Exception as e:
+        log_activity(f"[SessionStore Error] Failed load: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+class SessionListRequest(BaseModel):
+    project_key: str
+
+@app.post("/api/sessions/list")
+async def list_sessions(req: SessionListRequest, authorization: str = Header(None)):
+    auth(authorization)
+    if not db_pool:
+        raise HTTPException(status_code=500, detail="Database not connected")
+    try:
+        async with db_pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT DISTINCT session_id FROM agent_session_entries WHERE project_key=$1",
+                req.project_key
+            )
+        return {"sessions": [r["session_id"] for r in rows]}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+class SessionDeleteRequest(BaseModel):
+    project_key: str
+    session_id: str
+
+@app.post("/api/sessions/delete")
+async def delete_session(req: SessionDeleteRequest, authorization: str = Header(None)):
+    auth(authorization)
+    if not db_pool:
+        raise HTTPException(status_code=500, detail="Database not connected")
+    try:
+        async with db_pool.acquire() as conn:
+            await conn.execute(
+                "DELETE FROM agent_session_entries WHERE project_key=$1 AND session_id=$2",
+                req.project_key, req.session_id
+            )
+        return {"status": "success"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# ---------------------------------------------------------------------------
+# Database Retention & Heartbeat loops
+# ---------------------------------------------------------------------------
+
+async def db_cleanup_loop():
+    log_activity("Database Retention Cleanup task started")
+    while True:
+        try:
+            if db_pool:
+                async with db_pool.acquire() as conn:
+                    # Purge session entries older than 30 days
+                    result = await conn.execute(
+                        "DELETE FROM agent_session_entries WHERE created_at < NOW() - INTERVAL '30 days'"
+                    )
+                    log_activity(f"[Cleanup] Nightly retention sweep complete. Status: {result}")
+        except Exception as e:
+            log_activity(f"[Cleanup Warning] Failed to run retention cleanup: {e}")
+        await asyncio.sleep(86400) # Every 24 hours
+
+
 @app.get("/", response_class=HTMLResponse)
 async def dashboard():
     return HTMLResponse(content=DASHBOARD_HTML)
@@ -1346,6 +1467,8 @@ async def startup_event():
     threading.Thread(target=run_watchdog, daemon=True).start()
     # Start the db keep-alive loop on FastAPI event loop
     asyncio.create_task(db_heartbeat_loop())
+    # Start the db nightly retention cleanup loop
+    asyncio.create_task(db_cleanup_loop())
 
 
 # ---------------------------------------------------------------------------