runcell 0.1.17a2__py3-none-any.whl → 0.1.18a0__py3-none-any.whl

d5m_ai/_version.py

@@ -1,4 +1,4 @@
 # This file is auto-generated by Hatchling. As such, do not:
 #   - modify
 #   - track in version control e.g. be sure to add to .gitignore
-__version__ = VERSION = '0.1.17a2'
+__version__ = VERSION = '0.1.18a0'

d5m_ai/agent/proxy_handler.py

@@ -47,6 +47,7 @@ class AIJLAgentProxyHandler(WebSocketHandler):
     def initialize(self, **kwargs):
         self.connection_id: str = str(uuid.uuid4())
         self.current_request_id: str | None = None
+        self.local_only = False
 
         # Frontend WebSocket connection state
         self._is_closed = False
@@ -127,6 +128,12 @@ class AIJLAgentProxyHandler(WebSocketHandler):
         """Handle WebSocket connection opening from frontend."""
         # Get user token from local file
         user_token = self._get_user_token()
+        try:
+            local_only_param = self.get_query_argument("local_only", default="0")
+        except Exception:
+            local_only_param = "0"
+        # local_only keeps the frontend<->proxy WS for tools but avoids proxy<->remote WS.
+        self.local_only = str(local_only_param).lower() in {"1", "true", "yes"}
         
         if not user_token:
             print("[PROXY] ❌ No authentication token available")
@@ -147,11 +154,16 @@ class AIJLAgentProxyHandler(WebSocketHandler):
             "connection_id": self.connection_id
         })
         
-        # Establish connection to remote backend
-        await self._connect_to_remote_backend()
+        # Establish connection to remote backend unless running local-only tools.
+        if not self.local_only:
+            await self._connect_to_remote_backend()
+        else:
+            print("[PROXY] Local-only WebSocket enabled; skipping remote backend connection")
 
     async def _connect_to_remote_backend(self):
         """Establish WebSocket connection to remote backend server."""
+        if self.local_only:
+            return
         try:
             async with self._remote_connection_lock:
                 if self.remote_websocket is None or self._is_remote_connection_closed():
@@ -635,6 +647,9 @@ class AIJLAgentProxyHandler(WebSocketHandler):
 
     async def _forward_to_remote_backend(self, msg):
         """Forward message to remote backend."""
+        if self.local_only:
+            print("[PROXY] Local-only mode: skipping forward to remote backend")
+            return
         # Add connection ID for routing
         msg["connection_id"] = self.connection_id
         
@@ -658,19 +673,24 @@ class AIJLAgentProxyHandler(WebSocketHandler):
             request_id = msg.get("request_id")
             
             if request_id:
+                if handler.local_only:
+                    # In local-only mode there is no remote backend to reply to; resolve the waiter directly.
+                    await self._process_cell_result(handler, result)
+                    return
+
                 # This is from a remote backend tool call, forward it back
                 print(f"[PROXY] Forwarding cell_result to remote backend with request_id: {request_id}")
-                
+
                 # Clean up pending request tracking
                 if hasattr(self, '_pending_frontend_requests'):
                     self._pending_frontend_requests.discard(request_id)
-                
+
                 # Process the result to check for image data
                 if '"image/png"' in result:
                     print(f"[PROXY] Detected image/png data in cell_result")
                     modified_result = await handler.image_processor.replace_base64_with_urls(result)
                     result = modified_result
-                
+
                 await self._send_to_remote_backend({
                     "type": "tool_call_result",
                     "request_id": request_id,
@@ -689,6 +709,11 @@ class AIJLAgentProxyHandler(WebSocketHandler):
             if '"image/png"' in result:
                 modified_result = await handler.image_processor.replace_base64_with_urls(result)
                 result = modified_result
+
+            if handler.local_only:
+                # Local-only mode resolves the waiter directly instead of forwarding to remote.
+                await self._process_cell_result(handler, result)
+                return
             
             # Always forward the result back to remote backend with request_id
             if request_id:
@@ -711,6 +736,13 @@ class AIJLAgentProxyHandler(WebSocketHandler):
                 modified_result = await handler.image_processor.replace_base64_with_urls(result)
                 result = modified_result
 
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for edit_markdown_cell result")
+                return
+
             if request_id:
                 await self._send_to_remote_backend({
                     "type": "tool_call_result",
@@ -738,6 +770,13 @@ class AIJLAgentProxyHandler(WebSocketHandler):
         elif msg_type == "rerun_all_cells_result":
             result = msg.get("result", "No result provided")
             request_id = msg.get("request_id")
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for rerun_all_cells result")
+                return
             
             # If this result is from a remote backend tool call, forward it back
             if request_id:
@@ -764,6 +803,13 @@ class AIJLAgentProxyHandler(WebSocketHandler):
             else:
                 error = msg.get("error", "Unknown error occurred")
                 result = f"Error inserting markdown cell: {error}"
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for insert markdown cell result")
+                return
             
             # If this result is from a remote backend tool call, forward it back
             if request_id:
@@ -799,6 +845,14 @@ class AIJLAgentProxyHandler(WebSocketHandler):
                 result = msg.get("result", "file is opened in tab")
             else:
                 result = msg.get("error", "Unknown error occurred")
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for open tab result")
+                return
+
             if request_id:
                 print(f"[PROXY] Forwarding open_tab to remote backend")
                 await self._send_to_remote_backend({
@@ -815,6 +869,13 @@ class AIJLAgentProxyHandler(WebSocketHandler):
         elif msg_type == "insert_cell_result":
             result = msg.get("result", "No result provided")
             request_id = msg.get("request_id")
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for insert_cell result")
+                return
             
             if request_id:
                 print(f"[PROXY] Forwarding insert_cell_result to remote backend")
@@ -839,6 +900,10 @@ class AIJLAgentProxyHandler(WebSocketHandler):
                 print(f"[PROXY] Detected image/png data in run_cell_result")
                 modified_result = await handler.image_processor.replace_base64_with_urls(result)
                 result = modified_result
+
+            if handler.local_only:
+                await self._process_cell_result(handler, result)
+                return
             
             if request_id:
                 print(f"[PROXY] Forwarding run_cell_result to remote backend")
@@ -857,6 +922,13 @@ class AIJLAgentProxyHandler(WebSocketHandler):
         elif msg_type == "delete_cell_result":
             result = msg.get("result", "No result provided")
             request_id = msg.get("request_id")
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for delete_cell result")
+                return
             
             if request_id:
                 print(f"[PROXY] Forwarding delete_cell_result to remote backend")
@@ -879,6 +951,14 @@ class AIJLAgentProxyHandler(WebSocketHandler):
                 result = msg.get("result", "Terminal created")
             else:
                 result = msg.get("error", "Unknown error occurred")
+
+            if handler.local_only:
+                if handler.waiter and not handler.waiter.done():
+                    handler.waiter.set_result(result)
+                else:
+                    print(f"[PROXY] Handler waiter not available for run_long_terminal result")
+                return
+
             if request_id:
                 print(f"[PROXY] Forwarding run_long_terminal_result to remote backend")
                 await self._send_to_remote_backend({

d5m_ai/agent_v2/__init__.py

@@ -0,0 +1,11 @@
+"""
+Agent V2 proxy handlers.
+
+This module introduces a versioned HTTP/SSE agent endpoint that reuses the
+ask-mode tool loop pattern while keeping the existing WebSocket agent intact.
+"""
+
+from .proxy_handler import AIAgentV2ProxyHandler
+
+__all__ = ["AIAgentV2ProxyHandler"]
+

d5m_ai/agent_v2/proxy_handler.py

@@ -0,0 +1,263 @@
+"""
+Agent V2 HTTP/SSE proxy handler.
+
+MVP design:
+- Frontend <-> proxy may still use a local-only WebSocket for notebook tools.
+- Proxy <-> remote backend is strictly HTTP/SSE.
+- Tool execution reuses the ask-mode registry for now.
+"""
+
+from __future__ import annotations
+
+import aiohttp
+import json
+import logging
+import ssl
+from typing import Any, Dict
+
+import certifi
+
+from .._version import __version__ as client_version
+from ..chat.chat_proxy_handler import AIChatProxyHandler
+from tornado import web
+
+
+class AIAgentV2ProxyHandler(AIChatProxyHandler):
+    """
+    Versioned agent proxy that routes to `/v2/agent/chat` on the remote backend.
+
+    This intentionally subclasses the ask-mode proxy so we can validate the
+    HTTP tool loop end-to-end before migrating agent-specific tools/prompts.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        logging.info(
+            f"[AGENT-V2-PROXY] Initialized with remote backend: {self.remote_backend_url}"
+        )
+
+    async def _forward_to_remote_backend(
+        self, body_data: Dict[str, Any], stream: bool, user_token: str
+    ):
+        """Forward request to the versioned remote backend and handle response."""
+        try:
+            remote_url = f"{self.remote_backend_url}/v2/agent/chat"
+
+            if not hasattr(self, "_unified_content_index"):
+                starting_index = body_data.get("starting_content_index")
+                if starting_index is None:
+                    starting_index = body_data.get("last_content_index")
+                if starting_index is not None:
+                    try:
+                        self._unified_content_index = max(0, int(starting_index))
+                    except (TypeError, ValueError):
+                        self._unified_content_index = 0
+
+            if hasattr(self, "_unified_content_index") and self._unified_content_index > 0:
+                body_data["starting_content_index"] = self._unified_content_index
+                logging.info(
+                    "[AGENT-V2-PROXY] Continuing conversation with starting_content_index: "
+                    f"{self._unified_content_index}"
+                )
+
+            body_data["client_version"] = client_version
+
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {user_token}",
+            }
+
+            try:
+                ssl_context = ssl.create_default_context(cafile=certifi.where())
+            except Exception as exc:
+                logging.warning(
+                    f"[AGENT-V2-PROXY] Failed to create SSL context with certifi: {exc}"
+                )
+                ssl_context = ssl.create_default_context()
+
+            connector = aiohttp.TCPConnector(ssl=ssl_context)
+
+            async with aiohttp.ClientSession(connector=connector) as session:
+                async with session.post(remote_url, json=body_data, headers=headers) as response:
+                    if response.status != 200:
+                        error_text = await response.text()
+                        logging.error(
+                            f"[AGENT-V2-PROXY] Remote backend error {response.status}: {error_text}"
+                        )
+
+                        try:
+                            error_data = json.loads(error_text)
+                            error_status = error_data.get("status", "error")
+                            error_message = error_data.get("error", error_text)
+
+                            if response.status == 401:
+                                self.set_status(401)
+                                self.finish(
+                                    json.dumps(
+                                        {
+                                            "error": "Authentication failed. Please log in again.",
+                                            "status": "auth_error",
+                                        }
+                                    )
+                                )
+                            elif response.status == 402:
+                                self.set_status(402)
+                                self.finish(
+                                    json.dumps(
+                                        {
+                                            "error": error_message,
+                                            "status": "insufficient_credits",
+                                        }
+                                    )
+                                )
+                            else:
+                                self.set_status(response.status)
+                                self.finish(
+                                    json.dumps(
+                                        {"error": error_message, "status": error_status}
+                                    )
+                                )
+                        except json.JSONDecodeError:
+                            self.set_status(response.status)
+                            self.finish(
+                                json.dumps(
+                                    {
+                                        "error": f"Remote backend error: {error_text}",
+                                        "status": "error",
+                                    }
+                                )
+                            )
+                        return
+
+                    if stream:
+                        if not hasattr(self, "_unified_content_index"):
+                            self._unified_content_index = 0
+
+                        pending_tool_data = await self._handle_streaming_response(
+                            response, self._unified_content_index
+                        )
+                        if pending_tool_data:
+                            self._unified_content_index = pending_tool_data.get(
+                                "max_content_index", self._unified_content_index
+                            )
+                            await self._handle_pending_tool_call(
+                                pending_tool_data, body_data, user_token
+                            )
+                    else:
+                        response_data = await response.json()
+                        self.finish(json.dumps(response_data))
+
+        except aiohttp.ClientConnectorError as exc:
+            logging.error(
+                "[AGENT-V2-PROXY] Cannot connect to remote backend at "
+                f"{self.remote_backend_url}: {exc}"
+            )
+            self.set_status(503)
+            self.finish(
+                json.dumps(
+                    {
+                        "error": (
+                            "Cannot connect to agent v2 remote backend. "
+                            f"Please ensure the remote backend server is running at {self.remote_backend_url}"
+                        ),
+                        "status": "error",
+                    }
+                )
+            )
+        except Exception as exc:
+            logging.error(f"[AGENT-V2-PROXY] Error forwarding to remote backend: {exc}")
+            self.set_status(500)
+            self.finish(json.dumps({"error": str(exc), "status": "error"}))
+
+
+class AIAgentV2SessionProxyHandler(AIChatProxyHandler):
+    """Proxy handler for agent v2 session endpoints."""
+
+    @web.authenticated
+    async def get(self):
+        user_token = await self.get_user_token() or self.get_test_token_from_env()
+        if not user_token:
+            self.set_status(401)
+            self.finish(
+                json.dumps(
+                    {
+                        "error": "Authentication required. Please log in again.",
+                        "status": "auth_required",
+                    }
+                )
+            )
+            return
+
+        remote_url = f"{self.remote_backend_url}/v2/agent/session"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {user_token}",
+        }
+
+        try:
+            ssl_context = ssl.create_default_context(cafile=certifi.where())
+        except Exception as exc:
+            logging.warning(
+                f"[AGENT-V2-PROXY] Failed to create SSL context with certifi: {exc}"
+            )
+            ssl_context = ssl.create_default_context()
+
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.get(remote_url, params=self.request.query, headers=headers) as response:
+                payload_text = await response.text()
+                if response.status != 200:
+                    self.set_status(response.status)
+                    self.finish(payload_text)
+                    return
+                self.finish(payload_text)
+
+    @web.authenticated
+    async def post(self):
+        user_token = await self.get_user_token() or self.get_test_token_from_env()
+        if not user_token:
+            self.set_status(401)
+            self.finish(
+                json.dumps(
+                    {
+                        "error": "Authentication required. Please log in again.",
+                        "status": "auth_required",
+                    }
+                )
+            )
+            return
+
+        try:
+            body_data = json.loads(self.request.body.decode("utf-8")) if self.request.body else {}
+        except Exception:
+            body_data = {}
+
+        remote_url = f"{self.remote_backend_url}/v2/agent/session"
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {user_token}",
+        }
+
+        try:
+            ssl_context = ssl.create_default_context(cafile=certifi.where())
+        except Exception as exc:
+            logging.warning(
+                f"[AGENT-V2-PROXY] Failed to create SSL context with certifi: {exc}"
+            )
+            ssl_context = ssl.create_default_context()
+
+        connector = aiohttp.TCPConnector(ssl=ssl_context)
+        async with aiohttp.ClientSession(connector=connector) as session:
+            async with session.post(remote_url, json=body_data, headers=headers) as response:
+                payload_text = await response.text()
+                if response.status != 200:
+                    self.set_status(response.status)
+                    self.finish(payload_text)
+                    return
+                self.finish(payload_text)
+
+    async def options(self):
+        self.set_header("Access-Control-Allow-Origin", "*")
+        self.set_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+        self.set_header("Access-Control-Allow-Headers", "Content-Type, Authorization, X-User-Token")
+        self.set_status(200)

d5m_ai/chat/chat_proxy_handler.py

@@ -25,13 +25,26 @@ from typing import Dict, Any, Optional
 from jupyter_server.base.handlers import APIHandler
 from tornado import web
 
-from ..utils import build_remote_backend_url
+from ..utils import build_remote_backend_url, get_server_url_from_handler
 from ..auth.token_handler import get_current_user_token_string
 from ..tools.registry import run_ask_tool
 from .._version import __version__ as client_version
+from ..agent.image_processor import ImageProcessor
 
 logging.basicConfig(level=logging.INFO)
 
+FRONTEND_NOTEBOOK_TOOLS = {
+    "cell_execute",
+    "edit_cell",
+    "edit_markdown_cell",
+    "insert_markdown_cell",
+    "insert_cell",
+    "run_cell",
+    "delete_cell",
+    "open_tab",
+    "rerun_all_cells",
+}
+
 
 class AIChatProxyHandler(APIHandler):
     """Proxy handler for AI chat requests that forwards to remote backend."""
@@ -43,6 +56,10 @@ class AIChatProxyHandler(APIHandler):
         
         logging.info(f"[CHAT-PROXY] Initialized with remote backend: {self.remote_backend_url}")
         logging.info("[CHAT-PROXY] Token authentication enabled - will forward tokens to remote backend")
+
+        # Lazily initialized when we first need to process cell images.
+        self._image_processor: ImageProcessor | None = None
+        self._image_processor_connection_id: str | None = None
     
     async def get_user_token(self) -> Optional[str]:
         """
@@ -429,10 +446,79 @@ class AIChatProxyHandler(APIHandler):
         self.write(f"data: {event_text}\n\n")
         await self.flush()
 
+    def _normalize_frontend_tool_args(
+        self, tool_name: str, tool_args: Dict[str, Any]
+    ) -> tuple[Dict[str, Any], str | None]:
+        """Normalize frontend tool args so Jupyter contents paths are used."""
+        if not tool_args or tool_name not in FRONTEND_NOTEBOOK_TOOLS:
+            return tool_args, None
+
+        normalized = dict(tool_args)
+        file_path = normalized.get("file_path")
+        if file_path:
+            resolved_path, error_message = self._resolve_contents_path(file_path)
+            if error_message:
+                return {}, error_message
+            normalized["file_path"] = resolved_path
+        return normalized, None
+
+    def _resolve_contents_path(self, file_path: str) -> tuple[str, str | None]:
+        """Convert absolute filesystem paths to Jupyter contents paths when possible."""
+        from ..agent.path_utils import resolve_contents_path
+
+        root_dir = getattr(getattr(self, "contents_manager", None), "root_dir", None)
+        if not root_dir:
+            contents_manager = self.settings.get("contents_manager")
+            root_dir = getattr(contents_manager, "root_dir", None) if contents_manager else None
+        if not root_dir:
+            root_dir = (
+                self.settings.get("server_root_dir")
+                or self.settings.get("root_dir")
+                or os.getcwd()
+            )
+        return resolve_contents_path(file_path, root_dir)
+
+    def _get_image_processor(self, connection_id: str | None) -> ImageProcessor:
+        """Create/reuse an ImageProcessor for replacing base64 image outputs."""
+        if not self._image_processor or self._image_processor_connection_id != connection_id:
+            server_url = get_server_url_from_handler(
+                self,
+                fallback=os.environ.get("SERVER_URL", "http://localhost:8888"),
+            )
+            self._image_processor = ImageProcessor(connection_id or "unknown", server_url)
+            self._image_processor_connection_id = connection_id
+        return self._image_processor
+
+    async def _replace_base64_images(
+        self,
+        tool_name: str,
+        tool_result: Any,
+        sse_output: Any,
+        connection_id: str | None,
+    ) -> tuple[Any, Any]:
+        """Replace base64 image payloads with URLs for notebook tool outputs."""
+        if tool_name not in FRONTEND_NOTEBOOK_TOOLS:
+            return tool_result, sse_output
+
+        processor = None
+
+        if isinstance(tool_result, str) and '"image/png"' in tool_result:
+            processor = self._get_image_processor(connection_id)
+            tool_result = await processor.replace_base64_with_urls(tool_result)
+
+        if isinstance(sse_output, str) and '"image/png"' in sse_output:
+            if processor is None:
+                processor = self._get_image_processor(connection_id)
+            sse_output = await processor.replace_base64_with_urls(sse_output)
+
+        return tool_result, sse_output
+
     async def _handle_pending_tool_call(self, pending_data: Dict[str, Any], body_data: Dict[str, Any], user_token: str):
         """Execute local tools and continue the conversation with the remote backend."""
-        tool_calls = pending_data.get('tool_calls', []) or []
-        if not tool_calls:
+        client_tool_calls = pending_data.get('tool_calls', []) or []
+        server_tool_calls = pending_data.get('server_tool_calls', []) or []
+        server_tool_results = pending_data.get('server_tool_results', []) or []
+        if not client_tool_calls and not server_tool_calls:
             logging.warning("[CHAT-PROXY] Pending tool payload without tool_calls")
             return
 
@@ -448,7 +534,7 @@ class AIChatProxyHandler(APIHandler):
             thinking_blocks = deepcopy(thinking_blocks)
 
         assistant_tool_calls = []
-        for tool_call in tool_calls:
+        for tool_call in server_tool_calls + client_tool_calls:
             function_info = tool_call.get('function', {}) or {}
             assistant_tool_calls.append({
                 'id': tool_call.get('id', ''),
@@ -477,7 +563,16 @@ class AIChatProxyHandler(APIHandler):
 
         messages.append(assistant_message)
 
-        for idx, tool_call in enumerate(tool_calls):
+        if server_tool_results:
+            for idx, tool_result in enumerate(server_tool_results):
+                call_id = tool_result.get('tool_call_id') or tool_result.get('call_id') or f'server_tool_call_{idx}'
+                messages.append({
+                    'role': 'tool',
+                    'tool_call_id': call_id,
+                    'content': tool_result.get('content', '')
+                })
+
+        for idx, tool_call in enumerate(client_tool_calls):
             function_info = tool_call.get('function', {}) or {}
             tool_name = function_info.get('name', '')
             raw_arguments = function_info.get('arguments', '') or ''
@@ -487,14 +582,35 @@ class AIChatProxyHandler(APIHandler):
             except json.JSONDecodeError:
                 logging.warning(f"[CHAT-PROXY] Failed to parse tool arguments for {tool_name}: {raw_arguments}")
                 parsed_args = {}
-
-            tool_run = await run_ask_tool(tool_name, parsed_args)
-            if tool_run is None:
-                tool_result = f"Error: Unsupported tool '{tool_name}' in ask mode"
+            # Inject connection_id for tools that require frontend WebSocket (e.g., cell_execute).
+            if 'connection_id' not in parsed_args:
+                # This comes from the ask-mode local-only WS established by the frontend.
+                parsed_args['connection_id'] = body_data.get('connection_id')
+
+            normalized_args, error_message = self._normalize_frontend_tool_args(
+                tool_name, parsed_args
+            )
+            if error_message:
+                logging.warning(
+                    f"[CHAT-PROXY] Path resolution error for {tool_name}: {error_message}"
+                )
+                tool_result = error_message
                 sse_output = tool_result
             else:
-                tool_result = tool_run["llm_output"]
-                sse_output = tool_run["sse_output"]
+                tool_run = await run_ask_tool(tool_name, normalized_args)
+                if tool_run is None:
+                    tool_result = f"Error: Unsupported tool '{tool_name}' in ask mode"
+                    sse_output = tool_result
+                else:
+                    tool_result = tool_run["llm_output"]
+                    sse_output = tool_run["sse_output"]
+
+            tool_result, sse_output = await self._replace_base64_images(
+                tool_name,
+                tool_result,
+                sse_output,
+                body_data.get("connection_id"),
+            )
 
             call_id = tool_call.get('id', f'tool_call_{idx}')