feat: Cognito Agent Phase 4 - Extension System

very-simplified-stack/cognito-backend/app/core/extensions/api.py

@@ -0,0 +1,62 @@
+from typing import Callable, Awaitable, Literal, Optional, Dict, List, Any
+from pydantic import BaseModel
+
+class SessionStartPayload(BaseModel):
+    session_id: str
+    cwd: str
+    is_new: bool
+
+class BeforeToolCallPayload(BaseModel):
+    session_id: str
+    cwd: str
+    tool_name: str
+    arguments: Dict[str, Any]
+
+class AfterToolCallPayload(BaseModel):
+    session_id: str
+    cwd: str
+    tool_name: str
+    arguments: Dict[str, Any]
+    output: str
+    is_error: bool
+
+class MessageEndPayload(BaseModel):
+    session_id: str
+    cwd: str
+    role: str
+    content: str
+    tool_name: Optional[str] = None
+    tool_call_id: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+
+HookEvent = Literal["session_start", "before_tool_call", "after_tool_call", "message_end"]
+
+class ExtensionAPI:
+    def __init__(self, registry: "ExtensionRegistry", origin: Optional[str]):
+        # origin=None -> Global/Configured. origin=<cwd> -> Local al proyecto.
+        self._registry = registry
+        self._origin = origin
+
+    def register_tool(self, tool: Any) -> None:
+        self._registry.register_tool(tool, self._origin)
+
+    def register_backend(self, config: Any) -> None:
+        if self._origin is not None:
+            import logging
+            logging.getLogger("cognito.extensions").warning(
+                f"Ignored register_backend from project-local extension at {self._origin}"
+            )
+            return
+        self._registry.register_backend(config)
+
+    def register_intent(self, intent: str, backend_name: str, model: str) -> None:
+        if self._origin is not None:
+            import logging
+            logging.getLogger("cognito.extensions").warning(
+                f"Ignored register_intent from project-local extension at {self._origin}"
+            )
+            return
+        self._registry.register_intent(intent, backend_name, model)
+
+    def on(self, event: HookEvent, handler: Callable[[Any], Awaitable[Optional[str]]]) -> None:
+        self._registry.register_hook(event, handler, self._origin)

very-simplified-stack/cognito-backend/app/core/extensions/hooked_tool.py

@@ -0,0 +1,44 @@
+from app.core.tools.base import AgentTool, ToolResult
+from app.core.extensions.api import BeforeToolCallPayload, AfterToolCallPayload
+
+class HookedTool(AgentTool):
+    def __init__(self, inner, registry, session_id: str, cwd: str):
+        self.name = inner.name
+        self.description = inner.description
+        self.parameters_schema = inner.parameters_schema
+        self._inner = inner
+        self._registry = registry
+        self._session_id = session_id
+        self._cwd = cwd
+
+    async def execute(self, arguments, context):
+        veto = await self._registry.fire(
+            "before_tool_call",
+            BeforeToolCallPayload(
+                session_id=self._session_id,
+                cwd=self._cwd,
+                tool_name=self.name,
+                arguments=arguments
+            ),
+            self._cwd,
+        )
+
+        if veto:
+            return ToolResult(is_error=True, output=f"Llamada bloqueada por extensión: {veto}")
+
+        result = await self._inner.execute(arguments, context)
+
+        await self._registry.fire(
+            "after_tool_call",
+            AfterToolCallPayload(
+                session_id=self._session_id,
+                cwd=self._cwd,
+                tool_name=self.name,
+                arguments=arguments,
+                output=result.output,
+                is_error=result.is_error
+            ),
+            self._cwd,
+        )
+
+        return result

very-simplified-stack/cognito-backend/app/core/extensions/loader.py

@@ -0,0 +1,65 @@
+import os
+import json
+import logging
+import importlib.util
+from pathlib import Path
+from typing import List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from app.core.extensions.api import ExtensionAPI
+
+logger = logging.getLogger("cognito.extensions.loader")
+
+def discover_global() -> List[Path]:
+    ext_dir = Path.home() / ".cognito" / "extensions"
+    if not ext_dir.exists():
+        return []
+    return sorted([p for p in ext_dir.glob("*.py") if p.is_file()])
+
+def discover_configured() -> List[Path]:
+    config_path = Path.home() / ".cognito" / "config.json"
+    if not config_path.exists():
+        return []
+
+    try:
+        with open(config_path, "r") as f:
+            config = json.load(f)
+            paths = config.get("ExtensionPaths", [])
+
+            discovered = []
+            for p_str in paths:
+                p = Path(p_str).expanduser().resolve()
+                if p.is_file() and p.suffix == ".py":
+                    discovered.append(p)
+                elif p.is_dir():
+                    discovered.extend(sorted([f for f in p.glob("*.py") if f.is_file()]))
+            return discovered
+    except Exception as e:
+        logger.warning(f"Failed to load ExtensionPaths from config.json: {e}")
+        return []
+
+def discover_project_local(cwd: str) -> List[Path]:
+    ext_dir = Path(cwd) / ".cognito" / "extensions"
+    if not ext_dir.exists():
+        return []
+    return sorted([p for p in ext_dir.glob("*.py") if p.is_file()])
+
+def load_extension_file(path: Path, api: "ExtensionAPI") -> None:
+    try:
+        module_name = f"cognito.extensions.{path.stem}"
+        spec = importlib.util.spec_from_file_location(module_name, path)
+        if spec is None or spec.loader is None:
+            logger.warning(f"Could not load spec for extension at {path}")
+            return
+
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+
+        if hasattr(module, "register"):
+            module.register(api)
+            logger.info(f"Successfully registered extension: {path}")
+        else:
+            logger.warning(f"Extension at {path} has no register() function")
+
+    except Exception as e:
+        logger.warning(f"Failed to load extension at {path}: {e}", exc_info=True)

very-simplified-stack/cognito-backend/app/core/extensions/registry.py

@@ -0,0 +1,94 @@
+import logging
+from typing import Dict, List, Any, Optional, Callable, Awaitable, Literal
+from pathlib import Path
+from pydantic import BaseModel
+
+logger = logging.getLogger("cognito.extensions.registry")
+
+class ExtensionRegistry:
+    def __init__(self):
+        self._global_tools: Dict[str, Any] = {}
+        self._project_tools: Dict[str, Dict[str, Any]] = {}   # cwd -> {tool_name: tool}
+        self._hooks: Dict[str, List[tuple[Optional[str], Callable]]] = {}  # event -> [(origin, handler)]
+
+    def register_tool(self, tool: Any, origin: Optional[str]):
+        if origin is None:
+            if tool.name in self._global_tools:
+                logger.warning(f"Overwriting global tool: {tool.name}")
+            self._global_tools[tool.name] = tool
+        else:
+            if origin not in self._project_tools:
+                self._project_tools[origin] = {}
+            if tool.name in self._project_tools[origin]:
+                 logger.warning(f"Overwriting project-local tool: {tool.name} in {origin}")
+            self._project_tools[origin][tool.name] = tool
+
+    def register_hook(self, event: str, handler: Callable, origin: Optional[str]):
+        if event not in self._hooks:
+            self._hooks[event] = []
+        self._hooks[event].append((origin, handler))
+
+    def register_backend(self, config: Any):
+        from app.services.backend_router import backend_router
+        backend_router.register_backend(config)
+
+    def register_intent(self, intent: str, backend_name: str, model: str):
+        from app.services.semantic_orchestrator import semantic_orchestrator
+        semantic_orchestrator.add_intent_route(intent, backend_name, model)
+
+    def tools_for(self, cwd: str) -> List[Any]:
+        # Merge Global and project-local tools. Project-local wins.
+        tools = self._global_tools.copy()
+        if cwd in self._project_tools:
+            for name, tool in self._project_tools[cwd].items():
+                if name in tools:
+                    logger.warning(f"Project-local tool '{name}' overwriting global tool in {cwd}")
+                tools[name] = tool
+        return list(tools.values())
+
+    async def fire(self, event: str, payload: BaseModel, cwd: str) -> Optional[str]:
+        if event not in self._hooks:
+            return None
+
+        for origin, handler in self._hooks[event]:
+            if origin is None or origin == cwd:
+                try:
+                    res = await handler(payload)
+                    if event == "before_tool_call" and res:
+                        return str(res)
+                except Exception as e:
+                    logger.warning(f"Handler {handler.__name__} for event {event} failed: {e}", exc_info=True)
+        return None
+
+    def clear_project_local(self, cwd: str):
+        if cwd in self._project_tools:
+            self._project_tools[cwd].clear()
+
+        # Clear hooks for this origin
+        for event in self._hooks:
+            self._hooks[event] = [h for h in self._hooks[event] if h[0] != cwd]
+
+    def refresh(self, level: Literal["global", "configured", "project_local"],
+                cwd: Optional[str], backend_router, semantic_orchestrator) -> None:
+        from app.core.extensions.loader import (
+            discover_global, discover_configured, discover_project_local,
+            load_extension_file
+        )
+        from app.core.extensions.api import ExtensionAPI
+
+        if level == "global":
+            paths = discover_global()
+            for p in paths:
+                api = ExtensionAPI(self, None)
+                load_extension_file(p, api)
+        elif level == "configured":
+            paths = discover_configured()
+            for p in paths:
+                api = ExtensionAPI(self, None)
+                load_extension_file(p, api)
+        elif level == "project_local" and cwd:
+            self.clear_project_local(cwd)
+            paths = discover_project_local(cwd)
+            for p in paths:
+                api = ExtensionAPI(self, cwd)
+                load_extension_file(p, api)

very-simplified-stack/cognito-backend/app/services/semantic_orchestrator.py

@@ -172,12 +172,12 @@ def _parse_routing_plan(xml_text: str) -> RoutingPlan:
     )
 
 
-def _resolve_route(intent: str, client_map: Dict[str, BackendClient]) -> BackendClient:
+def _resolve_route(intent: str, client_map: Dict[str, BackendClient], routing_table: Dict[str, Dict[str, str]]) -> BackendClient:
     """
     Given an intent string, return the appropriate BackendClient with the
     correct model injected. Falls back to 'general' then to first available.
     """
-    route = MODEL_ROUTING.get(intent) or MODEL_ROUTING.get("general")
+    route = routing_table.get(intent) or routing_table.get("general")
     backend_name = route["backend"]
     model_name   = route["model"]
 
@@ -205,10 +205,16 @@ class SemanticOrchestrator:
                            respecting dependency order.
     """
 
-    def __init__(self, configs: List[BackendConfig]):
+    def __init__(self, configs: List[BackendConfig], extra_routing: Optional[Dict[str, Dict[str, str]]] = None):
         self._client_map = _build_client_map(configs)
+        self.routing = {**MODEL_ROUTING, **(extra_routing or {})}
         self._orchestrator_client = self._build_orchestrator_client()
 
+    def add_intent_route(self, intent: str, backend_name: str, model: str) -> None:
+        """Adds a new intent route at runtime."""
+        self.routing[intent] = {"backend": backend_name, "model": model}
+        logger.info(f"Intent route '{intent}' -> {backend_name}/{model} added via extension.")
+
     def _build_orchestrator_client(self) -> BackendClient:
         import copy
         base = self._client_map.get(ORCHESTRATOR_BACKEND)
@@ -261,7 +267,7 @@ async def _execute_plan(self, plan: RoutingPlan, original_input: str) -> Dict[st
 
             # Execute ready tasks concurrently
             async def _run(task: SubTask) -> tuple[str, str]:
-                client = _resolve_route(task.intent, self._client_map)
+                client = _resolve_route(task.intent, self._client_map, self.routing)
                 # Enrich the input slice with context from upstream tasks
                 context = ""
                 for dep_id in task.depends_on:
@@ -327,8 +333,8 @@ async def process(self, request: AIRequest) -> AIResponse:
                 "task_id":     t.id,
                 "description": t.description,
                 "intent":      t.intent,
-                "backend":     MODEL_ROUTING.get(t.intent, MODEL_ROUTING["general"])["backend"],
-                "model":       MODEL_ROUTING.get(t.intent, MODEL_ROUTING["general"])["model"],
+                "backend":     self.routing.get(t.intent, self.routing["general"])["backend"],
+                "model":       self.routing.get(t.intent, self.routing["general"])["model"],
             }
             for t in plan.subtasks
         ]

very-simplified-stack/cognito-backend/tests/test_extension_loader.py

@@ -0,0 +1,63 @@
+import pytest
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock
+from app.core.extensions.loader import discover_global, discover_configured, discover_project_local, load_extension_file
+from app.core.extensions.api import ExtensionAPI
+
+def test_discover_global(monkeypatch):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_home = Path(tmpdir)
+        monkeypatch.setattr(Path, "home", lambda: tmp_home)
+
+        ext_dir = tmp_home / ".cognito" / "extensions"
+        ext_dir.mkdir(parents=True)
+        (ext_dir / "ext1.py").touch()
+        (ext_dir / "ext2.py").touch()
+        (ext_dir / "README.md").touch() # skip non-py
+
+        paths = discover_global()
+        assert len(paths) == 2
+        assert paths[0].name == "ext1.py"
+        assert paths[1].name == "ext2.py"
+
+def test_discover_configured(monkeypatch):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_home = Path(tmpdir)
+        monkeypatch.setattr(Path, "home", lambda: tmp_home)
+
+        ext_file = tmp_home / "my_ext.py"
+        ext_file.touch()
+
+        ext_dir = tmp_home / "extra_exts"
+        ext_dir.mkdir()
+        (ext_dir / "a.py").touch()
+
+        config_dir = tmp_home / ".cognito"
+        config_dir.mkdir()
+        config_file = config_dir / "config.json"
+        import json
+        config_file.write_text(json.dumps({
+            "ExtensionPaths": [str(ext_file), str(ext_dir)]
+        }))
+
+        paths = discover_configured()
+        assert len(paths) == 2
+        assert ext_file in paths
+        assert (ext_dir / "a.py") in paths
+
+def test_load_broken_extension():
+    api = MagicMock()
+    with tempfile.NamedTemporaryFile(suffix=".py", mode="w") as tmp:
+        tmp.write("def register(api): raise RuntimeError('Boom')")
+        tmp.flush()
+
+        # Should not raise
+        load_extension_file(Path(tmp.name), api)
+
+    with tempfile.NamedTemporaryFile(suffix=".py", mode="w") as tmp:
+        tmp.write("invalid python code")
+        tmp.flush()
+
+        # Should not raise
+        load_extension_file(Path(tmp.name), api)