From ea33bb4a9c426e71a94aa77c8ce5de739fcaecbc Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 18:40:30 +0000 Subject: [PATCH] feat: Cognito Agent Phase 1 & 2 - Agent Loop + Session Management - Implement turn-based Agent Loop with tool execution (read, write, edit, bash). - Implement security layer: ProjectTrustStore, Protected Files, and sudo rejection. - Implement SessionManager with JSONL append-only persistence and fcntl locking. - Implement automatic history compaction based on token threshold. - Add SSE events for session info, text deltas, tool calls, and results. - Add endpoints for listing, getting, and forking sessions. - Comprehensive test suite for tools, sessions, compaction, and agent loop integration. Co-authored-by: Axlfc <14998495+Axlfc@users.noreply.github.com> --- .../cognito-backend/README.md | 8 + .../app/api/routes/ai_agents.py | 154 ++++++++++- .../cognito-backend/app/core/compaction.py | 74 ++++++ .../cognito-backend/app/core/events.py | 7 +- .../app/core/session_manager.py | 243 ++++++++++++++++++ .../tests/test_agent_loop_sessions.py | 100 +++++++ .../cognito-backend/tests/test_compaction.py | 39 +++ .../tests/test_session_manager.py | 66 +++++ 8 files changed, 679 insertions(+), 12 deletions(-) create mode 100644 very-simplified-stack/cognito-backend/app/core/compaction.py create mode 100644 very-simplified-stack/cognito-backend/app/core/session_manager.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_agent_loop_sessions.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_compaction.py create mode 100644 very-simplified-stack/cognito-backend/tests/test_session_manager.py diff --git a/very-simplified-stack/cognito-backend/README.md b/very-simplified-stack/cognito-backend/README.md index 276ff3c..d47ce73 100644 --- a/very-simplified-stack/cognito-backend/README.md +++ b/very-simplified-stack/cognito-backend/README.md @@ -60,6 +60,7 @@ Settings are loaded in the following order of priority: - `app/api/routes/openai_compat.py`: Core streaming and uncertainty calculation logic. - `app/services/backend_client.py`: Unified async client for Ollama and OpenAI backends. - `app/core/agent_loop.py`: Turning text generation into a tool-executing agent loop. +- `app/core/session_manager.py`: Persistence and history management for AI sessions. - `test-voice-api.ps1`: The main PowerShell profile script containing `cog` and `cogt`. - `Install-CognitoProfile.ps1`: Installer for the PowerShell environment. - `config.example.json`: Template for the user configuration file. @@ -73,6 +74,13 @@ El backend ahora incluye soporte para agentes capaces de ejecutar herramientas ( - **Body**: `{ "messages": [...], "cwd": "path/to/repo", "model_params": {} }` - **Eventos**: `text_delta`, `tool_call`, `tool_result`, `done`, `error`. +### Sesiones y Persistencia (Fase 2) +Las sesiones se guardan en `~/.cognito/sessions/` en formato JSONL (append-only) con un índice global en `index.json`. + +- **Compactado Automático**: Cuando una sesión supera el umbral de tokens (default 8000), el sistema genera automáticamente un resumen y compacta el historial para liberar ventana de contexto. +- **Continuidad**: Usa `session_id: "latest"` para continuar automáticamente la conversación más reciente en el `cwd` actual. +- **Forking**: Permite clonar una sesión existente para explorar ramas alternativas sin perder el historial original. + ### Herramientas Disponibles 1. `read`: Lee archivos del sistema (restringido al `cwd`). 2. `write`: Crea o sobrescribe archivos (requiere `trust`). diff --git a/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py b/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py index 6f27413..ffb98e4 100644 --- a/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py +++ b/very-simplified-stack/cognito-backend/app/api/routes/ai_agents.py @@ -1,5 +1,6 @@ import json from typing import List, Optional, Any, Dict +from pathlib import Path from fastapi import APIRouter, HTTPException from fastapi.responses import StreamingResponse from pydantic import BaseModel @@ -14,6 +15,9 @@ from app.core.tools.bash_tool import BashTool from app.core.project_trust import ProjectTrustStore from app.core.resource_loader import ResourceLoader +from app.core.session_manager import SessionManager +from app.core.compaction import should_compact, compact +from app.core.events import SessionInfoEvent, TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent, ErrorEvent import logging # Configure logging @@ -25,6 +29,7 @@ class AgentLoopRequest(BaseModel): messages: List[Dict[str, Any]] cwd: str + session_id: Optional[str] = None model_params: Optional[Dict[str, Any]] = None @router.post("/agent", response_model=AIResponse) @@ -49,50 +54,177 @@ async def run_ai_agent(request: AIRequest): @router.post("/agent/loop") async def run_agent_loop(request: AgentLoopRequest): """ - Agent Loop endpoint (SSE). + Agent Loop endpoint (SSE) with Session Management. """ - logger.info(f"Starting agent loop in {request.cwd}") + logger.info(f"Starting agent loop in {request.cwd} | session_id={request.session_id}") + session_manager = SessionManager() + is_new = False + session_id = request.session_id + + # 1. Resolver session_id + try: + if session_id == "latest": + session_id = session_manager.continue_recent(request.cwd) + if not session_id: + session_id = session_manager.create(request.cwd) + is_new = True + elif session_id: + metadata = session_manager.open(session_id) + if metadata.cwd != str(Path(request.cwd).resolve()): + error_msg = f"CWD mismatch: session {session_id} is for {metadata.cwd}, request is for {request.cwd}" + logger.error(error_msg) + return StreamingResponse( + iter([f"data: {json.dumps({'type': 'error', 'message': error_msg})}\n\n"]), + media_type="text/event-stream" + ) + else: + session_id = session_manager.create(request.cwd) + is_new = True + except Exception as e: + logger.error(f"Session resolution failed: {e}") + return StreamingResponse( + iter([f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"]), + media_type="text/event-stream" + ) + + # 2. Cargar historial y compactar si es necesario + effective_messages = session_manager.get_effective_messages(session_id) + if await should_compact(effective_messages): + try: + last_line = session_manager.get_last_line_index(session_id) + summary = await compact(effective_messages, backend_router=backend_router) + session_manager.append_compaction(session_id, summary, last_line) + effective_messages = session_manager.get_effective_messages(session_id) + except Exception as e: + logger.warning(f"Compaction failed for session {session_id}, continuing anyway: {e}") + + # 3. Preparar contexto y herramientas loader = ResourceLoader(request.cwd) trust_store = ProjectTrustStore() - context = ToolContext( cwd=request.cwd, trusted=trust_store.is_trusted(request.cwd), protected_files=loader.get_effective_protected_files() ) - tools = [ReadTool(), WriteTool(), EditTool(), BashTool()] - # Inyectar AGENTS.md en el system prompt si existe + # Inyectar AGENTS.md agents_md = loader.discover_agents_md() - messages = list(request.messages) + history = list(effective_messages) if agents_md: - system_msg = next((m for m in messages if m["role"] == "system"), None) + system_msg = next((m for m in history if m["role"] == "system"), None) if system_msg: system_msg["content"] += f"\n\nContext from AGENTS.md:\n{agents_md}" else: - messages.insert(0, {"role": "system", "content": f"Context from AGENTS.md:\n{agents_md}"}) + history.insert(0, {"role": "system", "content": f"Context from AGENTS.md:\n{agents_md}"}) + + # Mensajes nuevos de la request + new_messages = list(request.messages) + full_messages_for_loop = history + new_messages + + # REDESIGN: To persist assistant messages correctly, we need to capture them. + # The current agent_loop yields events but doesn't return the full objects. + async def event_generator_v2(): + # Persist incoming new messages from request + for msg in new_messages: + session_manager.append_message(session_id, msg["role"], msg["content"]) + + yield f"data: {SessionInfoEvent(session_id=session_id, is_new=is_new).model_dump_json()}\n\n" + + assistant_content = "" + current_tool_calls = [] - async def event_generator(): try: async for event in agent_loop( - messages=messages, + messages=full_messages_for_loop, tools=tools, context=context, backend_router=backend_router, model_params=request.model_params ): + if isinstance(event, TextDeltaEvent): + assistant_content += event.content + elif isinstance(event, ToolCallEvent): + current_tool_calls.append({ + "id": event.tool_call_id, + "type": "function", + "function": {"name": event.tool_name, "arguments": json.dumps(event.arguments)} + }) + elif isinstance(event, ToolResultEvent): + # Before persisting tool result, we MUST persist the assistant message that called it + if assistant_content or current_tool_calls: + session_manager.append_message( + session_id, + role="assistant", + content=assistant_content, + tool_calls=current_tool_calls if current_tool_calls else None + ) + assistant_content = "" + current_tool_calls = [] + + session_manager.append_message( + session_id, + role="tool", + content=event.output, + tool_name=event.tool_name, + tool_call_id=event.tool_call_id + ) + elif isinstance(event, DoneEvent): + if assistant_content or current_tool_calls: + session_manager.append_message( + session_id, + role="assistant", + content=assistant_content, + tool_calls=current_tool_calls if current_tool_calls else None + ) + yield f"data: {event.model_dump_json()}\n\n" except Exception as e: logger.error(f"Error in agent loop generator: {e}", exc_info=True) yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n" return StreamingResponse( - event_generator(), + event_generator_v2(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", } ) + +@router.get("/agent/sessions") +async def list_sessions(cwd: Optional[str] = None): + """ + List all sessions, optionally filtered by cwd. + """ + session_manager = SessionManager() + return session_manager.list_all(cwd=cwd) + +@router.get("/agent/sessions/{session_id}") +async def get_session(session_id: str): + """ + Get session metadata and effective messages. + """ + session_manager = SessionManager() + try: + metadata = session_manager.open(session_id) + messages = session_manager.get_effective_messages(session_id) + return { + "metadata": metadata, + "messages": messages + } + except FileNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id} not found") + +@router.post("/agent/sessions/{session_id}/fork") +async def fork_session(session_id: str): + """ + Fork an existing session into a new one. + """ + session_manager = SessionManager() + try: + new_id = session_manager.fork_from(session_id) + return {"session_id": new_id} + except FileNotFoundError: + raise HTTPException(status_code=404, detail=f"Session {session_id} not found") diff --git a/very-simplified-stack/cognito-backend/app/core/compaction.py b/very-simplified-stack/cognito-backend/app/core/compaction.py new file mode 100644 index 0000000..9709e82 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/compaction.py @@ -0,0 +1,74 @@ +import os +import json +import logging +from typing import List, Dict + +logger = logging.getLogger(__name__) + +# Configurable constants via environment variables +COMPACTION_THRESHOLD_TOKENS = int(os.getenv("COGNITO_SESSION_COMPACTION_THRESHOLD_TOKENS", "8000")) +KEEP_LAST_N_MESSAGES = int(os.getenv("COGNITO_SESSION_KEEP_LAST_N_MESSAGES", "6")) + +COMPACTION_SYSTEM_PROMPT = """Resume todos los mensajes anteriores de esta conversación de forma concisa. +ES CRÍTICO que preservas: +1. Decisiones clave tomadas por el usuario o el asistente. +2. Archivos que han sido leídos, escritos o editados. +3. Tareas pendientes que el usuario ha solicitado. +4. Cualquier restricción o advertencia mencionada (especialmente si se rechazó el acceso a un archivo protegido). + +Devuelve únicamente el texto del resumen, sin preámbulos ni explicaciones adicionales.""" + +async def should_compact(effective_messages: List[Dict], threshold_tokens: int = COMPACTION_THRESHOLD_TOKENS) -> bool: + """ + Heuristic token count: chars // 4. + Returns True if threshold is exceeded. + """ + if len(effective_messages) <= KEEP_LAST_N_MESSAGES + 2: # Don't compact very short history + return False + + total_chars = 0 + for msg in effective_messages: + content = msg.get("content", "") or "" + total_chars += len(content) + # Also count tool names and IDs if present + if "name" in msg: total_chars += len(msg["name"]) + if "tool_call_id" in msg: total_chars += len(msg["tool_call_id"]) + + estimated_tokens = total_chars // 4 + return estimated_tokens > threshold_tokens + +async def compact( + effective_messages: List[Dict], + keep_last_n: int = KEEP_LAST_N_MESSAGES, + backend_router = None, +) -> str: + """ + Summarizes messages except the last N using the backend router. + """ + if not backend_router: + from app.services.backend_router import backend_router as global_router + backend_router = global_router + + # Separate messages to compact and messages to keep + to_compact = effective_messages[:-keep_last_n] if keep_last_n > 0 else effective_messages + + # Format to_compact for the prompt + history_text = "" + for msg in to_compact: + role = msg.get("role", "unknown") + content = msg.get("content", "") + history_text += f"[{role}]: {content}\n\n" + + prompt = f"{COMPACTION_SYSTEM_PROMPT}\n\nHistorial a resumir:\n{history_text}" + + try: + logger.info(f"Triggering compaction for {len(to_compact)} messages") + result = await backend_router.generate(prompt=prompt) + summary = result.get("response", "").strip() + if not summary: + logger.warning("Compaction generated an empty summary.") + return "Resumen no disponible por error del modelo." + return summary + except Exception as e: + logger.error(f"Failed to generate compaction summary: {e}", exc_info=True) + raise e diff --git a/very-simplified-stack/cognito-backend/app/core/events.py b/very-simplified-stack/cognito-backend/app/core/events.py index 5510700..a61369b 100644 --- a/very-simplified-stack/cognito-backend/app/core/events.py +++ b/very-simplified-stack/cognito-backend/app/core/events.py @@ -28,4 +28,9 @@ class ErrorEvent(BaseModel): type: Literal["error"] = "error" message: str -AgentEvent = Union[TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent, ErrorEvent] +class SessionInfoEvent(BaseModel): + type: Literal["session_info"] = "session_info" + session_id: str + is_new: bool + +AgentEvent = Union[TextDeltaEvent, ToolCallEvent, ToolResultEvent, DoneEvent, ErrorEvent, SessionInfoEvent] diff --git a/very-simplified-stack/cognito-backend/app/core/session_manager.py b/very-simplified-stack/cognito-backend/app/core/session_manager.py new file mode 100644 index 0000000..e5282e2 --- /dev/null +++ b/very-simplified-stack/cognito-backend/app/core/session_manager.py @@ -0,0 +1,243 @@ +import json +import uuid +import logging +import fcntl +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional, List, Dict, Any +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + +class SessionMetadata(BaseModel): + session_id: str + cwd: str + created_at: str + updated_at: str + message_count: int + +class SessionManager: + def __init__(self, sessions_dir: Optional[Path] = None): + self.sessions_dir = sessions_dir or (Path.home() / ".cognito" / "sessions") + self.sessions_dir.mkdir(parents=True, exist_ok=True) + self.index_path = self.sessions_dir / "index.json" + self._ensure_index() + + def _ensure_index(self): + if not self.index_path.exists(): + with open(self.index_path, "w") as f: + json.dump({}, f) + + def _lock_file(self, f): + fcntl.flock(f, fcntl.LOCK_EX) + + def _unlock_file(self, f): + fcntl.flock(f, fcntl.LOCK_UN) + + def _get_index(self) -> Dict[str, Dict[str, Any]]: + with open(self.index_path, "r") as f: + self._lock_file(f) + try: + return json.load(f) + finally: + self._unlock_file(f) + + def _save_index(self, index: Dict[str, Dict[str, Any]]): + with open(self.index_path, "w") as f: + self._lock_file(f) + try: + json.dump(index, f, indent=2) + finally: + self._unlock_file(f) + + def create(self, cwd: str) -> str: + session_id = str(uuid.uuid4()) + now = datetime.now(timezone.utc).isoformat() + + # Create session file + session_file = self.sessions_dir / f"{session_id}.jsonl" + session_file.touch() + + # Update index + index = self._get_index() + index[session_id] = { + "cwd": str(Path(cwd).resolve()), + "created_at": now, + "updated_at": now, + "message_count": 0 + } + self._save_index(index) + + return session_id + + def open(self, session_id: str) -> SessionMetadata: + index = self._get_index() + if session_id not in index: + raise FileNotFoundError(f"Session {session_id} not found") + + data = index[session_id] + return SessionMetadata(session_id=session_id, **data) + + def fork_from(self, source_session_id: str, new_session_id: Optional[str] = None) -> str: + source_meta = self.open(source_session_id) + target_session_id = new_session_id or str(uuid.uuid4()) + + source_file = self.sessions_dir / f"{source_session_id}.jsonl" + target_file = self.sessions_dir / f"{target_session_id}.jsonl" + + # Copy file + import shutil + shutil.copy2(source_file, target_file) + + # Register in index + now = datetime.now(timezone.utc).isoformat() + index = self._get_index() + index[target_session_id] = { + "cwd": source_meta.cwd, + "created_at": now, + "updated_at": now, + "message_count": source_meta.message_count + } + self._save_index(index) + + return target_session_id + + def continue_recent(self, cwd: str) -> Optional[str]: + target_cwd = str(Path(cwd).resolve()) + index = self._get_index() + + sessions = [ + (sid, data) for sid, data in index.items() + if data["cwd"] == target_cwd + ] + + if not sessions: + return None + + sessions.sort(key=lambda x: x[1]["updated_at"], reverse=True) + return sessions[0][0] + + def list_all(self, cwd: Optional[str] = None) -> List[SessionMetadata]: + target_cwd = str(Path(cwd).resolve()) if cwd else None + index = self._get_index() + + results = [] + for sid, data in index.items(): + if target_cwd and data["cwd"] != target_cwd: + continue + results.append(SessionMetadata(session_id=sid, **data)) + + return results + + def append_message(self, session_id: str, role: str, content: str, tool_name: Optional[str] = None, tool_call_id: Optional[str] = None, tool_calls: Optional[List[Dict]] = None) -> None: + record = { + "type": "message", + "role": role, + "content": content, + "ts": datetime.now(timezone.utc).isoformat() + } + if tool_name: record["tool_name"] = tool_name + if tool_call_id: record["tool_call_id"] = tool_call_id + if tool_calls: record["tool_calls"] = tool_calls + + self._append_to_file(session_id, record) + self._update_index_metrics(session_id, message_delta=1) + + def append_compaction(self, session_id: str, summary: str, covers_through_line: int) -> None: + record = { + "type": "compaction", + "summary": summary, + "covers_through_line": covers_through_line, + "ts": datetime.now(timezone.utc).isoformat() + } + self._append_to_file(session_id, record) + self._update_index_metrics(session_id) + + def _append_to_file(self, session_id: str, record: Dict[str, Any]): + session_file = self.sessions_dir / f"{session_id}.jsonl" + with open(session_file, "a") as f: + self._lock_file(f) + try: + f.write(json.dumps(record) + "\n") + finally: + self._unlock_file(f) + + def _update_index_metrics(self, session_id: str, message_delta: int = 0): + index = self._get_index() + if session_id in index: + index[session_id]["updated_at"] = datetime.now(timezone.utc).isoformat() + index[session_id]["message_count"] += message_delta + self._save_index(index) + + def get_effective_messages(self, session_id: str) -> List[Dict[str, Any]]: + session_file = self.sessions_dir / f"{session_id}.jsonl" + if not session_file.exists(): + return [] + + messages = [] + compaction = None + compaction_line_index = -1 + + all_entries = [] + with open(session_file, "r") as f: + for i, line in enumerate(f): + if not line.strip(): continue + try: + data = json.loads(line) + all_entries.append((i, data)) + if data.get("type") == "compaction": + compaction = data + compaction_line_index = i + except json.JSONDecodeError: + logger.warning(f"Corrupt line {i} in session {session_id}") + continue + + if compaction: + summary = compaction.get("summary", "") + messages.append({"role": "system", "content": f"[Resumen de la conversación anterior]: {summary}"}) + start_line = compaction.get("covers_through_line", -1) + + for i, data in all_entries: + if i > compaction_line_index: # Only messages after the compaction record itself + messages.append(self._to_ai_message(data)) + elif i > start_line and i < compaction_line_index: + # Actually, wait. Rule says: discard all messages anterior a covers_through_line. + # Covers through line might be the line INDEX in the file. + # If I have: + # 0: msg + # 1: msg + # 2: compaction {covers_through_line: 1} + # 3: msg + # Then effective messages are [System(summary), msg(3)] + pass # Handled by the logic below + + # Revised logic for get_effective_messages with compaction: + messages = [] + messages.append({"role": "system", "content": f"[Resumen de la conversación anterior]: {compaction.get('summary', '')}"}) + for i, data in all_entries: + if i > compaction_line_index and data.get("type") == "message": + messages.append(self._to_ai_message(data)) + return messages + + else: + for i, data in all_entries: + if data.get("type") == "message": + messages.append(self._to_ai_message(data)) + return messages + + def _to_ai_message(self, data: Dict[str, Any]) -> Dict[str, Any]: + msg = {"role": data["role"], "content": data["content"]} + if "tool_name" in data: msg["name"] = data["tool_name"] + if "tool_call_id" in data: msg["tool_call_id"] = data["tool_call_id"] + if "tool_calls" in data: msg["tool_calls"] = data["tool_calls"] + return msg + + def get_last_line_index(self, session_id: str) -> int: + session_file = self.sessions_dir / f"{session_id}.jsonl" + if not session_file.exists(): + return -1 + line_count = 0 + with open(session_file, "r") as f: + for line in f: + if line.strip(): line_count += 1 + return line_count - 1 diff --git a/very-simplified-stack/cognito-backend/tests/test_agent_loop_sessions.py b/very-simplified-stack/cognito-backend/tests/test_agent_loop_sessions.py new file mode 100644 index 0000000..b9002ca --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_agent_loop_sessions.py @@ -0,0 +1,100 @@ +import pytest +import json +import tempfile +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock +from app.api.routes.ai_agents import AgentLoopRequest, run_agent_loop +from app.core.session_manager import SessionManager +from app.core.events import SessionInfoEvent, DoneEvent, TextDeltaEvent + +@pytest.fixture +def session_manager_dir(): + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) + +@pytest.mark.asyncio +async def test_agent_loop_session_integration(session_manager_dir, monkeypatch): + # Mock SessionManager's sessions_dir + monkeypatch.setattr("app.core.session_manager.Path.home", lambda: session_manager_dir) + + # Mock backend_router.generate_with_tools (used in agent_loop) + mock_router = MagicMock() + async def mock_gen(*args, **kwargs): + yield {"token": "Hello from assistant"} + mock_router.generate_with_tools = mock_gen + monkeypatch.setattr("app.api.routes.ai_agents.backend_router", mock_router) + + cwd = str(Path("/tmp/fake-repo").resolve()) + request = AgentLoopRequest( + messages=[{"role": "user", "content": "How are you?"}], + cwd=cwd + ) + + response = await run_agent_loop(request) + + # SSE stream parsing + events = [] + async for line in response.body_iterator: + if line.startswith("data: "): + events.append(json.loads(line[6:])) + + # Verify events + assert events[0]["type"] == "session_info" + session_id = events[0]["session_id"] + assert events[0]["is_new"] is True + + assert any(e["type"] == "text_delta" and "Hello" in e["content"] for e in events) + assert any(e["type"] == "done" for e in events) + + # Verify persistence + sm = SessionManager(sessions_dir=session_manager_dir / ".cognito" / "sessions") + msgs = sm.get_effective_messages(session_id) + assert len(msgs) == 2 + assert msgs[0]["role"] == "user" + assert msgs[0]["content"] == "How are you?" + assert msgs[1]["role"] == "assistant" + assert "Hello" in msgs[1]["content"] + +@pytest.mark.asyncio +async def test_agent_loop_compaction_integration(session_manager_dir, monkeypatch): + monkeypatch.setattr("app.core.session_manager.Path.home", lambda: session_manager_dir) + + sm = SessionManager(sessions_dir=session_manager_dir / ".cognito" / "sessions") + cwd = str(Path("/tmp/fake-repo").resolve()) + session_id = sm.create(cwd) + + # Add many messages to trigger compaction + for i in range(20): + sm.append_message(session_id, "user", "a" * 1000) + sm.append_message(session_id, "assistant", "b" * 1000) + + # Mock compaction + mock_router = MagicMock() + mock_router.generate = AsyncMock(return_value={"response": "Preserved summary"}) + async def mock_gen(*args, **kwargs): + yield {"token": "Response after compaction"} + mock_router.generate_with_tools = mock_gen + monkeypatch.setattr("app.api.routes.ai_agents.backend_router", mock_router) + + # Set low threshold + monkeypatch.setenv("COGNITO_SESSION_COMPACTION_THRESHOLD_TOKENS", "10") + + request = AgentLoopRequest( + messages=[{"role": "user", "content": "next msg"}], + cwd=cwd, + session_id=session_id + ) + + response = await run_agent_loop(request) + + # Consume response to trigger persistence + async for _ in response.body_iterator: + pass + + # Verify session file has compaction record + msgs = sm.get_effective_messages(session_id) + assert msgs[0]["role"] == "system" + assert "Preserved summary" in msgs[0]["content"] + + # And the last messages + the new one + assert any(m["content"] == "next msg" for m in msgs) diff --git a/very-simplified-stack/cognito-backend/tests/test_compaction.py b/very-simplified-stack/cognito-backend/tests/test_compaction.py new file mode 100644 index 0000000..ba5a3db --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_compaction.py @@ -0,0 +1,39 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock +from app.core.compaction import should_compact, compact + +@pytest.mark.asyncio +async def test_should_compact(): + msgs = [{"role": "user", "content": "a" * 100}] + assert await should_compact(msgs, threshold_tokens=1000) is False + + msgs = [{"role": "user", "content": "a" * 4000}] + # 4000 // 4 = 1000. Not greater than 1000. + # Actually, keep_last_n check might block it. + # Threshold is exceeded if estimated_tokens > threshold_tokens. + assert await should_compact(msgs, threshold_tokens=900) is False # because of KEEP_LAST_N_MESSAGES check + + msgs = [{"role": "user", "content": "a" * 100}] * 10 + assert await should_compact(msgs, threshold_tokens=10) is True + +@pytest.mark.asyncio +async def test_compact(): + backend_router = MagicMock() + backend_router.generate = AsyncMock(return_value={"response": "This is a summary"}) + + msgs = [ + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": "how are you?"}, + {"role": "assistant", "content": "fine"}, + ] + + # keep_last_n = 2 + summary = await compact(msgs, keep_last_n=2, backend_router=backend_router) + + assert summary == "This is a summary" + backend_router.generate.assert_called_once() + call_args = backend_router.generate.call_args[1] + assert "[user]: hello" in call_args["prompt"] + assert "[assistant]: hi" in call_args["prompt"] + assert "[user]: how are you?" not in call_args["prompt"] # kept diff --git a/very-simplified-stack/cognito-backend/tests/test_session_manager.py b/very-simplified-stack/cognito-backend/tests/test_session_manager.py new file mode 100644 index 0000000..9b4f2f4 --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_session_manager.py @@ -0,0 +1,66 @@ +import json +import pytest +import tempfile +from pathlib import Path +from app.core.session_manager import SessionManager + +@pytest.fixture +def session_manager(): + with tempfile.TemporaryDirectory() as tmpdir: + yield SessionManager(sessions_dir=Path(tmpdir)) + +def test_session_lifecycle(session_manager): + cwd = "/tmp/repo" + session_id = session_manager.create(cwd) + assert session_id is not None + + meta = session_manager.open(session_id) + assert meta.cwd == str(Path(cwd).resolve()) + assert meta.message_count == 0 + + session_manager.append_message(session_id, "user", "hello") + session_manager.append_message(session_id, "assistant", "hi") + + meta = session_manager.open(session_id) + assert meta.message_count == 2 + + msgs = session_manager.get_effective_messages(session_id) + assert len(msgs) == 2 + assert msgs[0]["role"] == "user" + assert msgs[1]["content"] == "hi" + +def test_continue_recent(session_manager): + cwd = "/tmp/repo" + s1 = session_manager.create(cwd) + s2 = session_manager.create(cwd) + + # Update s1 + session_manager.append_message(s1, "user", "msg") + + recent = session_manager.continue_recent(cwd) + assert recent == s1 + +def test_fork(session_manager): + cwd = "/tmp/repo" + s1 = session_manager.create(cwd) + session_manager.append_message(s1, "user", "original") + + s2 = session_manager.fork_from(s1) + assert s1 != s2 + + msgs = session_manager.get_effective_messages(s2) + assert len(msgs) == 1 + assert msgs[0]["content"] == "original" + +def test_corrupt_line(session_manager): + cwd = "/tmp/repo" + session_id = session_manager.create(cwd) + + session_file = session_manager.sessions_dir / f"{session_id}.jsonl" + with open(session_file, "a") as f: + f.write("invalid json\n") + f.write(json.dumps({"type": "message", "role": "user", "content": "valid", "ts": "..."}) + "\n") + + msgs = session_manager.get_effective_messages(session_id) + assert len(msgs) == 1 + assert msgs[0]["content"] == "valid"