diff --git a/very-simplified-stack/cognito-backend/app/core/extensions/__init__.py b/very-simplified-stack/cognito-backend/app/core/extensions/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/very-simplified-stack/cognito-backend/app/core/extensions/api.py b/very-simplified-stack/cognito-backend/app/core/extensions/api.py new file mode 100644 index 0000000..c182103 --- /dev/null +++ b/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= -> 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) diff --git a/very-simplified-stack/cognito-backend/app/core/extensions/hooked_tool.py b/very-simplified-stack/cognito-backend/app/core/extensions/hooked_tool.py new file mode 100644 index 0000000..51e83b7 --- /dev/null +++ b/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 diff --git a/very-simplified-stack/cognito-backend/app/core/extensions/loader.py b/very-simplified-stack/cognito-backend/app/core/extensions/loader.py new file mode 100644 index 0000000..0e3009e --- /dev/null +++ b/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) diff --git a/very-simplified-stack/cognito-backend/app/core/extensions/registry.py b/very-simplified-stack/cognito-backend/app/core/extensions/registry.py new file mode 100644 index 0000000..371f585 --- /dev/null +++ b/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) diff --git a/very-simplified-stack/cognito-backend/app/services/semantic_orchestrator.py b/very-simplified-stack/cognito-backend/app/services/semantic_orchestrator.py index bceb5ba..8b22459 100644 --- a/very-simplified-stack/cognito-backend/app/services/semantic_orchestrator.py +++ b/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 ] diff --git a/very-simplified-stack/cognito-backend/tests/test_extension_loader.py b/very-simplified-stack/cognito-backend/tests/test_extension_loader.py new file mode 100644 index 0000000..f4e5c5b --- /dev/null +++ b/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) diff --git a/very-simplified-stack/cognito-backend/tests/test_extension_registry.py b/very-simplified-stack/cognito-backend/tests/test_extension_registry.py new file mode 100644 index 0000000..bbd3a4c --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_extension_registry.py @@ -0,0 +1,67 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock +from pydantic import BaseModel +from app.core.extensions.registry import ExtensionRegistry + +class MockPayload(BaseModel): + data: str + +@pytest.mark.asyncio +async def test_registry_tool_isolation(): + registry = ExtensionRegistry() + + global_tool = MagicMock() + global_tool.name = "t1" + registry.register_tool(global_tool, origin=None) + + project_tool = MagicMock() + project_tool.name = "t2" + registry.register_tool(project_tool, origin="/repo/a") + + # tools_for A should have both + tools_a = registry.tools_for("/repo/a") + assert len(tools_a) == 2 + + # tools_for B should only have global + tools_b = registry.tools_for("/repo/b") + assert len(tools_b) == 1 + assert tools_b[0].name == "t1" + +@pytest.mark.asyncio +async def test_registry_hook_isolation(): + registry = ExtensionRegistry() + + h1 = AsyncMock() + registry.register_hook("session_start", h1, origin=None) + + h2 = AsyncMock() + registry.register_hook("session_start", h2, origin="/repo/a") + + payload = MockPayload(data="test") + + # Fire for A + await registry.fire("session_start", payload, "/repo/a") + assert h1.call_count == 1 + assert h2.call_count == 1 + + # Fire for B + h1.reset_mock() + h2.reset_mock() + await registry.fire("session_start", payload, "/repo/b") + assert h1.call_count == 1 + assert h2.call_count == 0 + +@pytest.mark.asyncio +async def test_registry_veto(): + registry = ExtensionRegistry() + + h1 = AsyncMock(return_value="No way") + registry.register_hook("before_tool_call", h1, origin=None) + + h2 = AsyncMock() + registry.register_hook("before_tool_call", h2, origin=None) + + veto = await registry.fire("before_tool_call", MockPayload(data="x"), "/any") + assert veto == "No way" + # Second handler should NOT be called + assert h2.call_count == 0 diff --git a/very-simplified-stack/cognito-backend/tests/test_hooked_tool.py b/very-simplified-stack/cognito-backend/tests/test_hooked_tool.py new file mode 100644 index 0000000..4777ee9 --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_hooked_tool.py @@ -0,0 +1,44 @@ +import pytest +from unittest.mock import AsyncMock, MagicMock +from app.core.extensions.hooked_tool import HookedTool +from app.core.tools.base import ToolResult + +@pytest.mark.asyncio +async def test_hooked_tool_execution(): + inner = MagicMock() + inner.name = "test_tool" + inner.description = "desc" + inner.parameters_schema = {} + inner.execute = AsyncMock(return_value=ToolResult(output="ok")) + + registry = MagicMock() + registry.fire = AsyncMock(return_value=None) + + hooked = HookedTool(inner, registry, "s1", "/cwd") + + res = await hooked.execute({"arg": 1}, MagicMock()) + + assert res.output == "ok" + assert registry.fire.call_count == 2 + assert registry.fire.await_args_list[0][0][0] == "before_tool_call" + assert registry.fire.await_args_list[1][0][0] == "after_tool_call" + assert inner.execute.called + +@pytest.mark.asyncio +async def test_hooked_tool_veto(): + inner = MagicMock() + inner.name = "test_tool" + inner.description = "desc" + inner.parameters_schema = {} + inner.execute = AsyncMock() + + registry = MagicMock() + registry.fire = AsyncMock(side_effect=[ "Blocked", None ]) + + hooked = HookedTool(inner, registry, "s1", "/cwd") + + res = await hooked.execute({}, MagicMock()) + + assert res.is_error + assert "Blocked" in res.output + assert not inner.execute.called diff --git a/very-simplified-stack/cognito-backend/tests/test_semantic_orchestrator_extra_routing.py b/very-simplified-stack/cognito-backend/tests/test_semantic_orchestrator_extra_routing.py new file mode 100644 index 0000000..9710507 --- /dev/null +++ b/very-simplified-stack/cognito-backend/tests/test_semantic_orchestrator_extra_routing.py @@ -0,0 +1,23 @@ +import pytest +from app.services.semantic_orchestrator import SemanticOrchestrator +from app.services.backend_registry import BACKENDS_BY_PRIORITY + +def test_semantic_orchestrator_default_routing(): + # extra_routing=None should use MODEL_ROUTING + orchestrator = SemanticOrchestrator(configs=BACKENDS_BY_PRIORITY) + from app.services.semantic_orchestrator import MODEL_ROUTING + assert orchestrator.routing == MODEL_ROUTING + +def test_semantic_orchestrator_extra_routing(): + extra = {"my_intent": {"backend": "b1", "model": "m1"}} + orchestrator = SemanticOrchestrator(configs=BACKENDS_BY_PRIORITY, extra_routing=extra) + + assert orchestrator.routing["my_intent"] == extra["my_intent"] + # Should still have default ones + assert "general" in orchestrator.routing + +def test_add_intent_route(): + orchestrator = SemanticOrchestrator(configs=BACKENDS_BY_PRIORITY) + orchestrator.add_intent_route("dynamic", "ext_backend", "ext_model") + + assert orchestrator.routing["dynamic"] == {"backend": "ext_backend", "model": "ext_model"}