pixelweaver 0.1.0

package/server/src/pixelweaver/mcp_registry.py

@@ -0,0 +1,247 @@
+"""MCP command registry for PixelWeaver.
+
+Maps MCP tool names to internal state mutations and read operations.
+Each tool handler receives validated arguments and returns a result dict.
+
+Tools are split into three categories:
+  - Raw command tools: one per drawing/layer/frame/project/canvas/history/export command
+  - Curated high-level tools: combine multiple raw tools behind a friendlier API
+  - Read/introspection tools: return data without mutating state
+"""
+
+from __future__ import annotations
+
+import base64
+import time
+import uuid
+from typing import Any
+
+from pixelweaver.connections import ConnectionManager
+from pixelweaver.state import ServerState
+from pixelweaver.storage import export_frame_png
+
+# ---------------------------------------------------------------------------
+# Tool definition helpers
+# ---------------------------------------------------------------------------
+
+
+class ToolDef:
+    """Lightweight descriptor for an MCP tool."""
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        parameters: dict[str, Any],
+        *,
+        mutates: bool = True,
+    ) -> None:
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+        # Whether the tool mutates state (controls thumbnail inclusion)
+        self.mutates = mutates
+
+
+# ---------------------------------------------------------------------------
+# Registry
+# ---------------------------------------------------------------------------
+
+
+class MCPCommandRegistry:
+    """Registry that maps MCP tool names to command dispatch."""
+
+    def __init__(self, state: ServerState, connections: ConnectionManager) -> None:
+        self.state = state
+        self.connections = connections
+        self._handlers: dict[str, Any] = {}
+        self._tool_defs: dict[str, ToolDef] = {}
+        self._register_all()
+
+    # -- public API ----------------------------------------------------------
+
+    async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        """Execute an MCP tool and return the result."""
+        handler = self._handlers.get(tool_name)
+        if handler is None:
+            return {"success": False, "error": f"Unknown tool: {tool_name}"}
+
+        tool_def = self._tool_defs.get(tool_name)
+
+        # Mutating tools take the shared lock so MCP and WebSocket mutation
+        # paths can't interleave a partial state change. Read-only tools skip
+        # the lock to keep introspection cheap and non-blocking.
+        try:
+            if tool_def and tool_def.mutates:
+                async with self.state.mutation_lock:
+                    result = await handler(arguments)
+            else:
+                result = await handler(arguments)
+        except Exception as exc:
+            return {"success": False, "error": str(exc)}
+
+        # Attach thumbnail for mutation operations
+        if tool_def and tool_def.mutates:
+            thumb = self._get_active_canvas_thumbnail()
+            if thumb:
+                result["thumbnail"] = thumb
+
+        return result
+
+    def get_tool_definitions(self) -> list[dict[str, Any]]:
+        """Return MCP-compatible tool definitions for all registered tools."""
+        defs = []
+        for td in self._tool_defs.values():
+            defs.append({
+                "name": td.name,
+                "description": td.description,
+                "inputSchema": td.parameters,
+            })
+        return defs
+
+    def get_tool_def(self, name: str) -> ToolDef | None:
+        """Get a single tool definition by name."""
+        return self._tool_defs.get(name)
+
+    # -- registration --------------------------------------------------------
+
+    def _register(self, tool_def: ToolDef, handler: Any) -> None:
+        """Register a tool definition and its async handler."""
+        self._tool_defs[tool_def.name] = tool_def
+        self._handlers[tool_def.name] = handler
+
+    def _register_all(self) -> None:
+        """Register every tool via domain modules."""
+        from .mcp_drawing_tools import register_drawing_tools
+        from .mcp_export_tools import register_export_tools
+        from .mcp_frame_tools import register_frame_tools
+        from .mcp_history_tools import register_history_tools
+        from .mcp_layer_tools import register_layer_tools
+        from .mcp_project_tools import register_project_tools
+        from .mcp_read_tools import register_read_tools
+
+        register_drawing_tools(self)
+        register_layer_tools(self)
+        register_frame_tools(self)
+        register_project_tools(self)
+        register_history_tools(self)
+        register_export_tools(self)
+        register_read_tools(self)
+
+    # -- helpers -------------------------------------------------------------
+
+    def _get_active_canvas_thumbnail(self) -> str | None:
+        """Get a thumbnail of the first canvas in the active project."""
+        project = self.state.get_active_project()
+        if project is None:
+            return None
+        canvas = next(iter(project.canvases.values()), None)
+        if canvas is None:
+            return None
+        # Composite all visible layers
+        try:
+            png_bytes = export_frame_png(project, canvas.name)
+            # Re-encode as base64 thumbnail (the export already is a PNG)
+            return base64.b64encode(png_bytes).decode("ascii")
+        except Exception:
+            return None
+
+    def _make_command(
+        self,
+        cmd_type: str,
+        params: dict[str, Any],
+        *,
+        plugin: str = "mcp",
+    ) -> dict[str, Any]:
+        """Build a command dict matching the protocol format."""
+        return {
+            "type": cmd_type,
+            "plugin": plugin,
+            "version": "1.0.0",
+            "params": params,
+            "id": str(uuid.uuid4()),
+            "timestamp": int(time.time() * 1000),
+        }
+
+    async def _dispatch_command(
+        self,
+        cmd_type: str,
+        params: dict[str, Any],
+        *,
+        plugin: str = "mcp",
+    ) -> dict[str, Any]:
+        """Create a command, store in history, and broadcast to WebSocket clients."""
+        project = self.state.get_active_project()
+        if project is None:
+            return {"success": False, "error": "No active project"}
+
+        cmd = self._make_command(cmd_type, params, plugin=plugin)
+        project.command_history.append(cmd)
+        project.redo_stack.clear()
+
+        # Broadcast to all connected WebSocket clients
+        await self.connections.broadcast({
+            "type": "command_broadcast",
+            "command": cmd,
+            "source": "mcp",
+        })
+
+        return {"success": True, "command_id": cmd["id"]}
+
+    def _require_active_project(self) -> dict[str, Any] | None:
+        """Return an error dict if no project is active, else None."""
+        if self.state.get_active_project() is None:
+            return {"success": False, "error": "No active project"}
+        return None
+
+    def _register_drawing_tool(
+        self, name: str, description: str, schema: dict[str, Any],
+    ) -> None:
+        """Register a single drawing tool that dispatches a command.
+
+        Automatically injects an optional ``frame_index`` parameter so
+        MCP clients can target a specific animation frame.  When provided,
+        the server sets ``canvas.current_frame_index`` before dispatching
+        the command (so the frontend draws on the correct frame) and
+        includes ``frame_index`` in the dispatched params for future
+        frontend-side routing.
+        """
+        # Inject frame_index into every drawing tool schema
+        schema["properties"]["frame_index"] = {
+            "type": "integer",
+            "description": "Target frame index (default: current frame)",
+        }
+
+        tool_def = ToolDef(name, description, schema, mutates=True)
+
+        async def handler(args: dict[str, Any], _name: str = name) -> dict[str, Any]:
+            err = self._require_active_project()
+            if err:
+                return err
+
+            project = self.state.get_active_project()
+            assert project is not None  # guaranteed by _require_active_project
+            canvas = next(iter(project.canvases.values()))
+
+            # Resolve frame_index: use provided value or fall back to current
+            frame_index = args.pop("frame_index", canvas.current_frame_index)
+
+            # Validate bounds
+            if frame_index < 0 or frame_index >= len(canvas.frames):
+                return {
+                    "success": False,
+                    "error": (
+                        f"frame_index {frame_index} out of range "
+                        f"[0, {len(canvas.frames)})"
+                    ),
+                }
+
+            # Point the canvas at the target frame before dispatch
+            canvas.current_frame_index = frame_index
+
+            # Include frame_index in dispatched params for frontend use
+            args["frame_index"] = frame_index
+
+            return await self._dispatch_command(_name, args)
+
+        self._register(tool_def, handler)

package/server/src/pixelweaver/mcp_resources.py

@@ -0,0 +1,312 @@
+"""MCP resources for PixelWeaver.
+
+Exposes project state as MCP resources so that LLM clients can read and
+subscribe to state changes.  Resources are registered on the FastMCP
+``mcp`` instance; subscriptions are wired through the low-level server
+so that mutating tool calls can push ``resource/updated`` notifications.
+
+Resource URIs
+-------------
+- ``project://state``                        -- full project overview
+- ``project://canvas/{name}``                -- single canvas metadata
+- ``project://canvas/{name}/frame/{index}``  -- composited frame as PNG
+- ``project://palette``                      -- colours in use on canvas
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import TYPE_CHECKING, Any
+
+from pydantic import AnyUrl
+
+from pixelweaver.mcp_bridge import CollabServerUnreachableError
+
+if TYPE_CHECKING:
+    from mcp.server.fastmcp import FastMCP
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Subscription tracking
+# ---------------------------------------------------------------------------
+# Tracks which resource URIs the client has subscribed to.  Populated via
+# the low-level subscribe/unsubscribe handlers registered in
+# ``init_mcp_subscriptions``.
+_subscribed_uris: set[str] = set()
+
+
+def get_subscribed_uris() -> frozenset[str]:
+    """Return a snapshot of currently subscribed URIs (for testing)."""
+    return frozenset(_subscribed_uris)
+
+
+# ---------------------------------------------------------------------------
+# Helpers shared by resource handlers and the notification path
+# ---------------------------------------------------------------------------
+
+def _unreachable_error(exc: Exception) -> str:
+    return json.dumps({
+        "error": (
+            f"Collab server unreachable: {exc}. "
+            "Make sure the PixelWeaver server is running "
+            "(pixelweaver serve --port 7779)."
+        ),
+    })
+
+
+async def _synced_state():
+    """Pull latest state and return it, or raise on failure."""
+    # Import here to avoid circular imports at module level.
+    from pixelweaver.mcp_server import _ensure_initialized, _state, _sync_from_collab
+
+    _ensure_initialized()
+    await _sync_from_collab()
+    return _state
+
+
+# ---------------------------------------------------------------------------
+# Resource registration
+# ---------------------------------------------------------------------------
+
+def init_mcp_resources(mcp: FastMCP) -> None:
+    """Register all MCP resources on the FastMCP server instance.
+
+    Called from ``mcp_server.init_mcp_tools`` so resources and tools are
+    initialised together.
+    """
+
+    # -- project://state ------------------------------------------------
+    @mcp.resource(
+        "project://state",
+        name="project_state",
+        description=(
+            "Full project overview: project name, dimensions, canvas list "
+            "with layer trees, frame counts, and playback settings."
+        ),
+        mime_type="application/json",
+    )
+    async def project_state() -> str:
+        try:
+            state = await _synced_state()
+        except CollabServerUnreachableError as exc:
+            return _unreachable_error(exc)
+
+        project = state.get_active_project()
+        if project is None:
+            return json.dumps({"error": "No active project"})
+
+        canvases_info: dict[str, Any] = {}
+        for cname, canvas in project.canvases.items():
+            canvases_info[cname] = {
+                "name": canvas.name,
+                "width": canvas.width,
+                "height": canvas.height,
+                "layer_count": len(canvas.layers),
+                "layers": [
+                    {
+                        "id": layer["id"],
+                        "name": layer.get("name", ""),
+                        "visible": layer.get("visible", True),
+                        "opacity": layer.get("opacity", 1.0),
+                    }
+                    for layer in canvas.layers
+                ],
+                "frame_count": canvas.frame_count,
+                "current_frame_index": canvas.current_frame_index,
+                "global_fps": canvas.global_fps,
+            }
+
+        return json.dumps(
+            {
+                "name": project.name,
+                "width": project.width,
+                "height": project.height,
+                "canvases": canvases_info,
+            },
+            indent=2,
+        )
+
+    # -- project://canvas/{name} ----------------------------------------
+    @mcp.resource(
+        "project://canvas/{name}",
+        name="canvas_info",
+        description=(
+            "Metadata for a single canvas: dimensions, layer tree, "
+            "frame count, playback state."
+        ),
+        mime_type="application/json",
+    )
+    async def canvas_info(name: str) -> str:
+        try:
+            state = await _synced_state()
+        except CollabServerUnreachableError as exc:
+            return _unreachable_error(exc)
+
+        project = state.get_active_project()
+        if project is None:
+            return json.dumps({"error": "No active project"})
+
+        canvas = project.canvases.get(name)
+        if canvas is None:
+            return json.dumps({"error": f"Canvas {name!r} not found"})
+
+        return json.dumps(
+            {
+                "name": canvas.name,
+                "width": canvas.width,
+                "height": canvas.height,
+                "layers": [
+                    {
+                        "id": layer["id"],
+                        "name": layer.get("name", ""),
+                        "visible": layer.get("visible", True),
+                        "opacity": layer.get("opacity", 1.0),
+                        "locked": layer.get("locked", False),
+                        "blend_mode": layer.get("blendMode", "normal"),
+                    }
+                    for layer in canvas.layers
+                ],
+                "frame_count": canvas.frame_count,
+                "current_frame_index": canvas.current_frame_index,
+                "global_fps": canvas.global_fps,
+            },
+            indent=2,
+        )
+
+    # -- project://canvas/{name}/frame/{index} --------------------------
+    @mcp.resource(
+        "project://canvas/{name}/frame/{index}",
+        name="frame_png",
+        description="Composited frame pixel data as a base64-encoded PNG.",
+        mime_type="image/png",
+    )
+    async def frame_png(name: str, index: str) -> bytes:
+        """Return composited PNG bytes for the given canvas frame."""
+        from pixelweaver.storage import export_frame_png
+
+        try:
+            state = await _synced_state()
+        except CollabServerUnreachableError as exc:
+            # Resources returning bytes cannot easily embed error JSON;
+            # raise so the MCP framework surfaces the error to the client.
+            raise RuntimeError(_unreachable_error(exc)) from exc
+
+        project = state.get_active_project()
+        if project is None:
+            raise RuntimeError("No active project")
+
+        try:
+            frame_index = int(index)
+        except ValueError:
+            raise RuntimeError(f"Invalid frame index: {index!r}")
+
+        return export_frame_png(project, name, frame_index)
+
+    # -- project://palette -----------------------------------------------
+    @mcp.resource(
+        "project://palette",
+        name="palette",
+        description="Unique colours currently present on the active canvas.",
+        mime_type="application/json",
+    )
+    async def palette_resource() -> str:
+        try:
+            state = await _synced_state()
+        except CollabServerUnreachableError as exc:
+            return _unreachable_error(exc)
+
+        project = state.get_active_project()
+        if project is None:
+            return json.dumps({"error": "No active project"})
+
+        canvas = next(iter(project.canvases.values()), None)
+        if canvas is None:
+            return json.dumps({"error": "No canvas in project"})
+
+        max_colors = 256
+        colors: set[str] = set()
+
+        for layer in canvas.layers:
+            data = canvas.current_frame().pixel_data.get(layer["id"])
+            if data is None:
+                continue
+            for i in range(0, len(data), 4):
+                r, g, b, a = data[i], data[i + 1], data[i + 2], data[i + 3]
+                if a > 0:
+                    colors.add(f"#{r:02x}{g:02x}{b:02x}{a:02x}")
+                if len(colors) >= max_colors:
+                    break
+            if len(colors) >= max_colors:
+                break
+
+        return json.dumps({"colors": sorted(colors)}, indent=2)
+
+
+# ---------------------------------------------------------------------------
+# Subscription handlers and notification support
+# ---------------------------------------------------------------------------
+
+def init_mcp_subscriptions(mcp: FastMCP) -> None:
+    """Wire subscribe/unsubscribe handlers and enable the capability.
+
+    Must be called after ``init_mcp_resources`` but before the server
+    starts accepting connections.
+    """
+    low_level = mcp._mcp_server
+
+    @low_level.subscribe_resource()
+    async def handle_subscribe(uri: AnyUrl) -> None:
+        uri_str = str(uri)
+        _subscribed_uris.add(uri_str)
+        logger.info("Client subscribed to resource: %s", uri_str)
+
+    @low_level.unsubscribe_resource()
+    async def handle_unsubscribe(uri: AnyUrl) -> None:
+        uri_str = str(uri)
+        _subscribed_uris.discard(uri_str)
+        logger.info("Client unsubscribed from resource: %s", uri_str)
+
+    # Patch the capability so the server advertises subscribe=True.
+    # The low-level server hardcodes subscribe=False in get_capabilities;
+    # we wrap it to flip the flag after the fact.
+    _original_get_capabilities = low_level.get_capabilities
+
+    def _patched_get_capabilities(notification_options, experimental_capabilities):
+        caps = _original_get_capabilities(notification_options, experimental_capabilities)
+        if caps.resources is not None:
+            caps.resources.subscribe = True
+        return caps
+
+    low_level.get_capabilities = _patched_get_capabilities
+
+
+async def notify_resource_subscribers(mcp: FastMCP) -> None:
+    """Send ``resource/updated`` for every currently-subscribed URI.
+
+    Called after a mutating tool pushes state back to the collab server,
+    so the MCP client knows to re-read any resources it is watching.
+
+    This must run inside a request context (i.e. while a tool handler is
+    executing) so we can access the session.
+    """
+    if not _subscribed_uris:
+        return
+
+    try:
+        session = mcp._mcp_server.request_context.session
+    except LookupError:
+        # No active request context -- nothing we can do.
+        logger.debug("No request context; skipping resource notifications")
+        return
+
+    for uri_str in _subscribed_uris:
+        try:
+            uri = AnyUrl(uri_str)
+            await session.send_resource_updated(uri=uri)
+            logger.debug("Sent resource/updated for %s", uri_str)
+        except Exception:
+            logger.warning(
+                "Failed to send resource/updated for %s", uri_str, exc_info=True
+            )