Add resource subscriptions and live screen feed

resources/subscribe / resources/unsubscribe wire a producer
callback into a provider. ResourceProvider gains optional
subscribe/unsubscribe hooks; ChainProvider fans out to children;
LiveScreenProvider exposes autocontrol://screen/live and pushes
notifications/resources/updated every poll_seconds while at least
one client is subscribed. resources.subscribe capability flag
flipped to true so clients know to attempt subscriptions.

Author: JE-Chen

je_auto_control/utils/mcp_server/__init__.py

@@ -27,17 +27,18 @@
     MCPPrompt, MCPPromptArgument, PromptProvider, default_prompt_provider,
 )
 from je_auto_control.utils.mcp_server.resources import (
-    MCPResource, ResourceProvider, default_resource_provider,
+    LiveScreenProvider, MCPResource, ResourceProvider,
+    default_resource_provider,
 )
 from je_auto_control.utils.mcp_server.tools import (
     MCPContent, MCPTool, MCPToolAnnotations, build_default_tool_registry,
     make_plugin_tool, register_plugin_tools,
 )
 
 __all__ = [
-    "AuditLogger", "FakeState", "HttpMCPServer", "MCPContent",
-    "MCPLogBridge", "MCPPrompt", "MCPPromptArgument", "MCPResource",
-    "MCPServer", "MCPTool", "MCPToolAnnotations",
+    "AuditLogger", "FakeState", "HttpMCPServer", "LiveScreenProvider",
+    "MCPContent", "MCPLogBridge", "MCPPrompt", "MCPPromptArgument",
+    "MCPResource", "MCPServer", "MCPTool", "MCPToolAnnotations",
     "OperationCancelledError", "PluginWatcher", "PromptProvider",
     "RateLimiter", "ResourceProvider", "ToolCallContext",
     "build_default_tool_registry",

je_auto_control/utils/mcp_server/resources.py

@@ -7,10 +7,13 @@
 abstraction lets callers compose custom sources without touching the
 JSON-RPC layer.
 """
+import base64
+import io
 import json
 import os
+import threading
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 
 @dataclass(frozen=True)
@@ -45,6 +48,21 @@ def set_workspace_root(self, root: str) -> None:
         """Hook for MCP roots. Default: no-op. FS-backed providers override."""
         del root
 
+    def subscribe(self, uri: str,
+                  on_update: Callable[[], None]) -> Optional[Any]:
+        """Optional hook: start emitting ``on_update`` calls until unsubscribed.
+
+        Return a non-``None`` handle when this provider owns ``uri`` and
+        accepted the subscription. The default implementation returns
+        ``None`` (not subscribable).
+        """
+        del uri, on_update
+        return None
+
+    def unsubscribe(self, uri: str, handle: Any) -> None:
+        """Cancel a previous :meth:`subscribe` handle."""
+        del uri, handle
+
 
 class FileSystemProvider(ResourceProvider):
     """Expose ``*.json`` action files in ``root`` under ``<scheme>://files/<name>``."""
@@ -173,18 +191,107 @@ def set_workspace_root(self, root: str) -> None:
         for provider in self.providers:
             provider.set_workspace_root(root)
 
+    def subscribe(self, uri: str,
+                  on_update: Callable[[], None]) -> Optional[Any]:
+        for provider in self.providers:
+            handle = provider.subscribe(uri, on_update)
+            if handle is not None:
+                return (provider, handle)
+        return None
+
+    def unsubscribe(self, uri: str, handle: Any) -> None:
+        if not isinstance(handle, tuple) or len(handle) != 2:
+            return
+        provider, child_handle = handle
+        provider.unsubscribe(uri, child_handle)
+
+
+class LiveScreenProvider(ResourceProvider):
+    """Live screen feed at ``autocontrol://screen/live``.
+
+    ``read`` always grabs a fresh PNG (base64-encoded). Subscribers
+    receive ``on_update`` calls every ``poll_seconds`` so they can
+    re-fetch the resource and surface live state to the model.
+    """
+
+    URI = "autocontrol://screen/live"
+
+    def __init__(self, poll_seconds: float = 1.0) -> None:
+        self._poll_seconds = max(0.1, float(poll_seconds))
+        self._lock = threading.Lock()
+        self._subscribers: Dict[int, Callable[[], None]] = {}
+        self._next_handle = 1
+        self._thread: Optional[threading.Thread] = None
+        self._stop = threading.Event()
+
+    def list(self) -> List[MCPResource]:
+        return [MCPResource(
+            uri=self.URI, name="screen_live",
+            description=("Current screen as base64 PNG. Subscribe to be "
+                          "notified when it should be re-fetched."),
+            mime_type="image/png",
+        )]
+
+    def read(self, uri: str) -> Optional[Dict[str, Any]]:
+        if uri != self.URI:
+            return None
+        from je_auto_control.utils.cv2_utils.screenshot import pil_screenshot
+        image = pil_screenshot()
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+        return {"uri": uri, "mimeType": "image/png", "blob": encoded}
+
+    def subscribe(self, uri: str,
+                  on_update: Callable[[], None]) -> Optional[Any]:
+        if uri != self.URI:
+            return None
+        with self._lock:
+            handle = self._next_handle
+            self._next_handle += 1
+            self._subscribers[handle] = on_update
+            if self._thread is None or not self._thread.is_alive():
+                self._stop.clear()
+                self._thread = threading.Thread(
+                    target=self._broadcast_loop, daemon=True,
+                    name="MCPLiveScreen",
+                )
+                self._thread.start()
+        return handle
+
+    def unsubscribe(self, uri: str, handle: Any) -> None:
+        if uri != self.URI:
+            return
+        with self._lock:
+            self._subscribers.pop(int(handle), None)
+            if not self._subscribers:
+                self._stop.set()
+                self._thread = None
+
+    def _broadcast_loop(self) -> None:
+        while not self._stop.is_set():
+            with self._lock:
+                callbacks = list(self._subscribers.values())
+            for callback in callbacks:
+                try:
+                    callback()
+                except (OSError, RuntimeError, ValueError):
+                    pass
+            self._stop.wait(self._poll_seconds)
+
 
 def default_resource_provider(root: str = ".") -> ResourceProvider:
     """Return the resource provider exposed by the default MCP server."""
     return ChainProvider([
         FileSystemProvider(root=root),
         HistoryProvider(),
         CommandsProvider(),
+        LiveScreenProvider(),
     ])
 
 
 __all__ = [
     "ChainProvider", "CommandsProvider", "FileSystemProvider",
-    "HistoryProvider", "MCPResource", "ResourceProvider",
-    "default_resource_provider",
+    "HistoryProvider", "LiveScreenProvider", "MCPResource",
+    "ResourceProvider", "default_resource_provider",
 ]

je_auto_control/utils/mcp_server/server.py

@@ -86,6 +86,8 @@ def __init__(self, tools: Optional[List[MCPTool]] = None,
         self._pending_outbound: Dict[Any, Dict[str, Any]] = {}
         self._outbound_lock = threading.Lock()
         self._client_capabilities: Dict[str, Any] = {}
+        self._resource_subscriptions: Dict[str, Any] = {}
+        self._subscriptions_lock = threading.Lock()
 
     def register_tool(self, tool: MCPTool) -> None:
         """Add or replace a tool in the live registry.
@@ -379,6 +381,10 @@ def _dispatch(self, msg_id: Any, method: Optional[str],
                                    for resource in self._resources.list()]}
         if method == "resources/read":
             return self._handle_resources_read(params)
+        if method == "resources/subscribe":
+            return self._handle_resources_subscribe(params)
+        if method == "resources/unsubscribe":
+            return self._handle_resources_unsubscribe(params)
         if method == "prompts/list":
             return {"prompts": [prompt.to_descriptor()
                                  for prompt in self._prompts.list()]}
@@ -411,7 +417,7 @@ def _handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
             self._client_capabilities = client_caps
         capabilities: Dict[str, Any] = {
             "tools": {"listChanged": True},
-            "resources": {"listChanged": False, "subscribe": False},
+            "resources": {"listChanged": False, "subscribe": True},
             "prompts": {"listChanged": False},
             "sampling": {},
             "logging": {},
@@ -434,6 +440,45 @@ def _handle_resources_read(self,
             raise _MCPError(-32602, f"Unknown resource: {uri}")
         return {"contents": [content]}
 
+    def _handle_resources_subscribe(self,
+                                    params: Dict[str, Any]) -> Dict[str, Any]:
+        uri = params.get("uri")
+        if not isinstance(uri, str) or not uri:
+            raise _MCPError(-32602, "resources/subscribe requires 'uri'")
+        with self._subscriptions_lock:
+            if uri in self._resource_subscriptions:
+                return {}
+        handle = self._resources.subscribe(
+            uri, lambda u=uri: self._notify_resource_updated(u),
+        )
+        if handle is None:
+            raise _MCPError(-32602, f"Unsubscribable resource: {uri}")
+        with self._subscriptions_lock:
+            self._resource_subscriptions[uri] = handle
+        return {}
+
+    def _handle_resources_unsubscribe(self,
+                                      params: Dict[str, Any]) -> Dict[str, Any]:
+        uri = params.get("uri")
+        if not isinstance(uri, str) or not uri:
+            raise _MCPError(-32602, "resources/unsubscribe requires 'uri'")
+        with self._subscriptions_lock:
+            handle = self._resource_subscriptions.pop(uri, None)
+        if handle is not None:
+            self._resources.unsubscribe(uri, handle)
+        return {}
+
+    def _notify_resource_updated(self, uri: str) -> None:
+        notifier = self._notifier
+        if notifier is None:
+            return
+        try:
+            notifier("notifications/resources/updated", {"uri": uri})
+        except (OSError, RuntimeError, ValueError):
+            autocontrol_logger.exception(
+                "MCP failed to send resources/updated for %s", uri,
+            )
+
     def _handle_prompts_get(self, params: Dict[str, Any]) -> Dict[str, Any]:
         name = params.get("name")
         arguments = params.get("arguments") or {}

test/unit_test/headless/test_mcp_server.py

@@ -1504,6 +1504,77 @@ def test_aliases_excluded_from_read_only_registry():
     assert "mouse_pos" in names
 
 
+def test_initialize_advertises_resources_subscribe_capability():
+    server = MCPServer(tools=[])
+    response = _decode(server.handle_line(_request("initialize", params={})))
+    caps = response["result"]["capabilities"]
+    assert caps["resources"]["subscribe"] is True
+
+
+def test_resources_subscribe_and_notification_round_trip():
+    """Subscribe to a fake resource and verify the server forwards updates."""
+    from je_auto_control.utils.mcp_server.resources import (
+        ChainProvider, MCPResource, ResourceProvider,
+    )
+
+    class _FakeProvider(ResourceProvider):
+        URI = "fake://live"
+
+        def __init__(self):
+            self.callback = None
+
+        def list(self):
+            return [MCPResource(uri=self.URI, name="fake")]
+
+        def read(self, uri):
+            if uri == self.URI:
+                return {"uri": uri, "mimeType": "text/plain", "text": "hi"}
+            return None
+
+        def subscribe(self, uri, on_update):
+            if uri != self.URI:
+                return None
+            self.callback = on_update
+            return "fake-handle"
+
+        def unsubscribe(self, uri, handle):
+            self.callback = None
+
+    fake = _FakeProvider()
+    chain = ChainProvider([fake])
+    captured = []
+    server = MCPServer(tools=[], resource_provider=chain)
+    server.set_notifier(lambda method, params: captured.append((method, params)))
+
+    sub_response = _decode(server.handle_line(_request(
+        "resources/subscribe", params={"uri": "fake://live"},
+    )))
+    assert sub_response["result"] == {}
+    assert fake.callback is not None
+
+    # Simulate the provider noticing fresh content.
+    fake.callback()
+
+    methods = [event[0] for event in captured]
+    assert "notifications/resources/updated" in methods
+    update = next(e for e in captured if e[0] == "notifications/resources/updated")
+    assert update[1] == {"uri": "fake://live"}
+
+    unsub_response = _decode(server.handle_line(_request(
+        "resources/unsubscribe", params={"uri": "fake://live"},
+    )))
+    assert unsub_response["result"] == {}
+    assert fake.callback is None
+
+
+def test_resources_subscribe_rejects_unknown_uri():
+    server = MCPServer(tools=[])
+    response = _decode(server.handle_line(_request(
+        "resources/subscribe", params={"uri": "fake://nowhere"},
+    )))
+    assert response["error"]["code"] == -32602
+
+
 def test_default_registry_lists_core_automation_tools():
     names = {tool.name for tool in build_default_tool_registry()}
     expected = {