Bump pydantic (#1593)

Co-authored-by: openhands <[email protected]>

Author: tofarr

openhands-sdk/openhands/sdk/workspace/remote/async_remote_workspace.py

@@ -1,12 +1,9 @@
-import asyncio
-import json
 from collections.abc import Generator
 from pathlib import Path
 from typing import Any
 
 import httpx
 from pydantic import PrivateAttr
-from websockets import connect
 
 from openhands.sdk.git.models import GitChange, GitDiff
 from openhands.sdk.workspace.models import CommandResult, FileOperationResult
@@ -64,8 +61,8 @@ async def execute_command(
     ) -> CommandResult:
         """Execute a bash command on the remote system.
 
-        This method conneccts a websocket client, sends a bash command, and
-        then waits for the output until the command completes.
+        This method starts a bash command via the remote agent server API,
+        then polls for the output until the command completes.
 
         Args:
             command: The bash command to execute
@@ -75,66 +72,9 @@ async def execute_command(
         Returns:
             CommandResult: Result with stdout, stderr, exit_code, and other metadata
         """
-        try:
-            result = await asyncio.wait_for(
-                self._execute_command(command, cwd, timeout), timeout=timeout
-            )
-            return result
-        except TimeoutError:
-            return CommandResult(
-                command=command,
-                exit_code=-1,
-                stdout="",
-                stderr="",
-                timeout_occurred=True,
-            )
-
-    async def _execute_command(
-        self,
-        command: str,
-        cwd: str | Path | None,
-        timeout: float,
-    ):
-        # Convert http(s) scheme to ws(s) for websocket connection
-        ws_host = self.host.replace("https://", "wss://").replace("http://", "ws://")
-        url = f"{ws_host}/sockets/bash-events?session_api_key={self.api_key}"
-        async with connect(url) as websocket:
-            payload = {
-                "command": command,
-                "timeout": int(timeout),
-            }
-            if cwd:
-                payload["cwd"] = str(cwd)
-            await websocket.send(json.dumps(payload))
-            command_id: str | None = None
-            stdout_parts: list[str] = []
-            stderr_parts: list[str] = []
-            exit_code: int | None = None
-            while exit_code is None:
-                data = await websocket.recv()
-                event = json.loads(data)
-                if event.get("kind") == "BashCommand":
-                    if command_id is None and event.get("command") == command:
-                        command_id = event.get("id")
-                if event.get("kind") == "BashOutput":
-                    if event.get("command_id") == command_id:
-                        if event.get("stdout"):
-                            stdout_parts.append(event.get("stdout"))
-                        if event.get("stderr"):
-                            stderr_parts.append(event.get("stderr"))
-                        exit_code = event.get("exit_code")
-
-            # Combine all output parts
-            stdout = "".join(stdout_parts)
-            stderr = "".join(stderr_parts)
-
-        return CommandResult(
-            command=command,
-            exit_code=exit_code,
-            stdout=stdout,
-            stderr=stderr,
-            timeout_occurred=exit_code == -1 and "timed out" in stderr,
-        )
+        generator = self._execute_command_generator(command, cwd, timeout)
+        result = await self._execute(generator)
+        return result
 
     async def file_upload(
         self,

openhands-sdk/pyproject.toml

@@ -9,7 +9,7 @@ dependencies = [
     "fastmcp>=2.11.3",
     "httpx>=0.27.0",
     "litellm>=1.80.10",
-    "pydantic>=2.11.7",
+    "pydantic>=2.12.5",
     "python-frontmatter>=1.1.0",
     "python-json-logger>=3.3.0",
     "tenacity>=9.1.2",

tests/sdk/workspace/remote/test_async_remote_workspace.py

@@ -1,7 +1,6 @@
 """Unit tests for AsyncRemoteWorkspace class."""
 
 import asyncio
-import json
 from pathlib import Path
 from unittest.mock import AsyncMock, Mock, patch
 
@@ -104,59 +103,32 @@ def test_generator():
 
 
 @pytest.mark.asyncio
-async def test_async_execute_command():
-    """Test execute_command method with websocket mocking."""
+@patch(
+    "openhands.sdk.workspace.remote.async_remote_workspace.AsyncRemoteWorkspace._execute"
+)
+async def test_async_execute_command(mock_execute):
+    """Test execute_command method calls _execute with correct generator."""
     workspace = AsyncRemoteWorkspace(
-        host="http://localhost:8000", api_key="test-key", working_dir="workspace"
+        host="http://localhost:8000", working_dir="workspace"
     )
 
-    # Create mock websocket
-    mock_websocket = AsyncMock()
-    mock_websocket.send = AsyncMock()
-    # Simulate server responses: BashCommand event followed by BashOutput event
-    mock_websocket.recv = AsyncMock(
-        side_effect=[
-            json.dumps({"kind": "BashCommand", "command": "echo hello", "id": "cmd-1"}),
-            json.dumps(
-                {
-                    "kind": "BashOutput",
-                    "command_id": "cmd-1",
-                    "stdout": "hello\n",
-                    "stderr": "",
-                    "exit_code": 0,
-                }
-            ),
-        ]
+    expected_result = CommandResult(
+        command="echo hello",
+        exit_code=0,
+        stdout="hello\n",
+        stderr="",
+        timeout_occurred=False,
     )
+    mock_execute.return_value = expected_result
 
-    mock_connect = AsyncMock()
-    mock_connect.__aenter__.return_value = mock_websocket
-    mock_connect.__aexit__.return_value = None
-
-    with patch(
-        "openhands.sdk.workspace.remote.async_remote_workspace.connect",
-        return_value=mock_connect,
-    ) as patched_connect:
-        result = await workspace.execute_command("echo hello", cwd="/tmp", timeout=30.0)
-
-    assert result.command == "echo hello"
-    assert result.exit_code == 0
-    assert result.stdout == "hello\n"
-    assert result.stderr == ""
-    assert result.timeout_occurred is False
+    result = await workspace.execute_command("echo hello", cwd="/tmp", timeout=30.0)
 
-    # Verify connect was called with ws:// scheme (converted from http://)
-    patched_connect.assert_called_once()
-    ws_url = patched_connect.call_args[0][0]
-    assert ws_url.startswith("ws://")
-    assert "localhost:8000/sockets/bash-events" in ws_url
+    assert result == expected_result
+    mock_execute.assert_called_once()
 
-    # Verify websocket was called with correct payload
-    mock_websocket.send.assert_called_once()
-    sent_payload = json.loads(mock_websocket.send.call_args[0][0])
-    assert sent_payload["command"] == "echo hello"
-    assert sent_payload["timeout"] == 30
-    assert sent_payload["cwd"] == "/tmp"
+    # Verify the generator was created correctly
+    generator_arg = mock_execute.call_args[0][0]
+    assert hasattr(generator_arg, "__next__")
 
 
 @pytest.mark.asyncio
@@ -217,56 +189,25 @@ async def test_async_file_download(mock_execute):
 
 @pytest.mark.asyncio
 async def test_async_execute_command_with_path_objects():
-    """Test execute_command works with Path objects for cwd and https:// to wss://."""
+    """Test execute_command works with Path objects for cwd."""
     workspace = AsyncRemoteWorkspace(
-        host="https://example.com", api_key="test-key", working_dir="workspace"
-    )
-
-    # Create mock websocket
-    mock_websocket = AsyncMock()
-    mock_websocket.send = AsyncMock()
-    mock_websocket.recv = AsyncMock(
-        side_effect=[
-            json.dumps({"kind": "BashCommand", "command": "ls", "id": "cmd-2"}),
-            json.dumps(
-                {
-                    "kind": "BashOutput",
-                    "command_id": "cmd-2",
-                    "stdout": "file1.txt\n",
-                    "stderr": "",
-                    "exit_code": 0,
-                }
-            ),
-        ]
+        host="http://localhost:8000", working_dir="workspace"
     )
 
-    mock_connect = AsyncMock()
-    mock_connect.__aenter__.return_value = mock_websocket
-    mock_connect.__aexit__.return_value = None
+    with patch.object(workspace, "_execute") as mock_execute:
+        expected_result = CommandResult(
+            command="ls",
+            exit_code=0,
+            stdout="file1.txt\n",
+            stderr="",
+            timeout_occurred=False,
+        )
+        mock_execute.return_value = expected_result
 
-    with patch(
-        "openhands.sdk.workspace.remote.async_remote_workspace.connect",
-        return_value=mock_connect,
-    ) as patched_connect:
         result = await workspace.execute_command("ls", cwd=Path("/tmp/test"))
 
-    assert result.command == "ls"
-    assert result.exit_code == 0
-    assert result.stdout == "file1.txt\n"
-    assert result.stderr == ""
-    assert result.timeout_occurred is False
-
-    # Verify connect was called with wss:// scheme (converted from https://)
-    patched_connect.assert_called_once()
-    ws_url = patched_connect.call_args[0][0]
-    assert ws_url.startswith("wss://")
-    assert "example.com/sockets/bash-events" in ws_url
-
-    # Verify Path object was converted to string in the payload
-    mock_websocket.send.assert_called_once()
-    sent_payload = json.loads(mock_websocket.send.call_args[0][0])
-    assert sent_payload["command"] == "ls"
-    assert sent_payload["cwd"] == "/tmp/test"  # Path is converted to string
+        assert result == expected_result
+        mock_execute.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -399,33 +340,29 @@ async def test_async_concurrent_operations():
         host="http://localhost:8000", working_dir="workspace"
     )
 
-    # Mock different results for different operations
-    command_result = CommandResult(
-        command="echo test",
-        exit_code=0,
-        stdout="test\n",
-        stderr="",
-        timeout_occurred=False,
-    )
-    upload_result = FileOperationResult(
-        success=True,
-        source_path="/local/file1.txt",
-        destination_path="/remote/file1.txt",
-        file_size=50,
-    )
-    download_result = FileOperationResult(
-        success=True,
-        source_path="/remote/file2.txt",
-        destination_path="/local/file2.txt",
-        file_size=75,
-    )
+    with patch.object(workspace, "_execute") as mock_execute:
+        # Mock different results for different operations
+        command_result = CommandResult(
+            command="echo test",
+            exit_code=0,
+            stdout="test\n",
+            stderr="",
+            timeout_occurred=False,
+        )
+        upload_result = FileOperationResult(
+            success=True,
+            source_path="/local/file1.txt",
+            destination_path="/remote/file1.txt",
+            file_size=50,
+        )
+        download_result = FileOperationResult(
+            success=True,
+            source_path="/remote/file2.txt",
+            destination_path="/local/file2.txt",
+            file_size=75,
+        )
 
-    with (
-        patch.object(workspace, "_execute_command") as mock_execute_command,
-        patch.object(workspace, "_execute") as mock_execute,
-    ):
-        mock_execute_command.return_value = command_result
-        mock_execute.side_effect = [upload_result, download_result]
+        mock_execute.side_effect = [command_result, upload_result, download_result]
 
         # Run operations concurrently
         tasks = [
@@ -439,5 +376,4 @@ async def test_async_concurrent_operations():
         assert results[0] == command_result
         assert results[1] == upload_result
         assert results[2] == download_result
-        mock_execute_command.assert_called_once()
-        assert mock_execute.call_count == 2
+        assert mock_execute.call_count == 3