Add conversation_id to automation completion callback (#3012)

Co-authored-by: openhands <openhands@all-hands.dev>

Author: malhotra5

openhands-sdk/openhands/sdk/conversation/impl/remote_conversation.py

@@ -768,6 +768,8 @@ def __init__(
                 )
             self._id = uuid.UUID(cid)
 
+            workspace.register_conversation(str(self._id))
+
         # Initialize the remote state
         self._state = RemoteState(
             self._client,

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

@@ -210,3 +210,25 @@ def default_conversation_tags(self) -> dict[str, str] | None:
             Dictionary of tag key-value pairs, or None if no default tags.
         """
         return None
+
+    def register_conversation(self, conversation_id: str) -> None:
+        """Register a conversation ID with this workspace.
+
+        Called by RemoteConversation after creation to associate the conversation
+        with the workspace. Subclasses can override to track conversation IDs
+        for callbacks or other purposes.
+
+        Args:
+            conversation_id: The conversation ID to register
+        """
+        # Default implementation is a no-op
+        pass
+
+    @property
+    def conversation_id(self) -> str | None:
+        """Get the most recently registered conversation ID.
+
+        Returns:
+            The conversation ID if one has been registered, None otherwise.
+        """
+        return None

openhands-workspace/openhands/workspace/cloud/workspace.py

@@ -151,6 +151,7 @@ class OpenHandsCloudWorkspace(RemoteWorkspace):
     _exposed_urls: list[dict[str, Any]] | None = PrivateAttr(default=None)
     _automation_callback_url: str | None = PrivateAttr(default=None)
     _automation_run_id: str | None = PrivateAttr(default=None)
+    _conversation_id: str | None = PrivateAttr(default=None)
 
     @property
     def default_conversation_tags(self) -> dict[str, str]:
@@ -796,6 +797,28 @@ def _send_settings_request(
 
         return response
 
+    def register_conversation(self, conversation_id: str) -> None:
+        """Register a conversation ID with this workspace.
+
+        Called by RemoteConversation after creation to associate the conversation
+        with the workspace. The conversation ID is included in the completion
+        callback sent to the automation service.
+
+        Args:
+            conversation_id: The conversation ID to register
+        """
+        self._conversation_id = conversation_id
+        logger.debug(f"Registered conversation: {conversation_id}")
+
+    @property
+    def conversation_id(self) -> str | None:
+        """Get the registered conversation ID.
+
+        Returns:
+            The conversation ID if one has been registered, None otherwise.
+        """
+        return self._conversation_id
+
     def __del__(self) -> None:
         self.cleanup()
 
@@ -813,6 +836,9 @@ def _send_completion_callback(
 
         Called by ``__exit__`` before ``cleanup()``.  Does nothing when
         ``AUTOMATION_CALLBACK_URL`` env var was not set.
+
+        Includes ``conversation_id`` in the payload if one was registered via
+        ``register_conversation()``.
         """
         try:
             callback_url = self._automation_callback_url
@@ -829,6 +855,10 @@ def _send_completion_callback(
         if exc_val is not None:
             payload["error"] = str(exc_val)
 
+        # Include conversation_id if one was registered
+        if self._conversation_id is not None:
+            payload["conversation_id"] = self._conversation_id
+
         try:
             headers = {"Authorization": f"Bearer {self.cloud_api_key}"}
             with httpx.Client(timeout=10.0) as cb_client:

tests/sdk/conversation/remote/test_remote_conversation.py

@@ -1321,3 +1321,25 @@ def test_remote_conversation_execute_tool_not_implemented(self, mock_ws_client):
             conversation.execute_tool("any_tool", mock_action)
 
         assert "not yet supported for RemoteConversation" in str(exc_info.value)
+
+    @patch(
+        "openhands.sdk.conversation.impl.remote_conversation.WebSocketCallbackClient"
+    )
+    def test_remote_conversation_calls_register_conversation(self, mock_ws_client):
+        """Test RemoteConversation.__init__ calls workspace.register_conversation."""
+        conversation_id = str(uuid.uuid4())
+        self.setup_mock_client(conversation_id=conversation_id)
+
+        mock_ws_instance = Mock()
+        mock_ws_client.return_value = mock_ws_instance
+
+        # Patch register_conversation at the class level to verify it gets called
+        with patch.object(RemoteWorkspace, "register_conversation") as mock_register:
+            # Create RemoteConversation - this should call register_conversation
+            _conversation = RemoteConversation(
+                agent=self.agent,
+                workspace=self.workspace,
+            )
+
+            # Verify register_conversation was called with the conversation ID
+            mock_register.assert_called_once_with(conversation_id)

tests/workspace/test_cloud_workspace.py

@@ -579,3 +579,79 @@ def test_callback_failure_does_not_raise(monkeypatch):
 
         # Should not raise
         ws.__exit__(None, None, None)
+
+
+# --- conversation_id registration tests ---
+
+
+def test_register_conversation_sets_conversation_id():
+    """register_conversation sets the _conversation_id attribute."""
+    ws = _make_local_workspace()
+
+    ws.register_conversation("conv-123")
+
+    assert ws._conversation_id == "conv-123"
+    assert ws.conversation_id == "conv-123"
+
+
+def test_conversation_id_property_returns_none_initially():
+    """conversation_id property returns None when no conversation registered."""
+    ws = _make_local_workspace()
+
+    assert ws.conversation_id is None
+
+
+def test_callback_includes_conversation_id_when_registered(monkeypatch):
+    """Callback payload includes conversation_id when registered."""
+    monkeypatch.setenv("AUTOMATION_CALLBACK_URL", "https://svc.test/complete")
+    monkeypatch.setenv("AUTOMATION_RUN_ID", "run-42")
+    ws = _make_local_workspace()
+
+    # Register a conversation
+    ws.register_conversation("conv-xyz")
+
+    mock_resp = MagicMock()
+    mock_resp.status_code = 200
+
+    with patch("httpx.Client") as MockClient:
+        mock_client = MagicMock()
+        mock_client.post.return_value = mock_resp
+        mock_client.__enter__ = MagicMock(return_value=mock_client)
+        mock_client.__exit__ = MagicMock(return_value=False)
+        MockClient.return_value = mock_client
+
+        ws.__exit__(None, None, None)
+
+        # Check the POST payload includes conversation_id
+        mock_client.post.assert_called_once()
+        payload = mock_client.post.call_args.kwargs["json"]
+        assert payload["status"] == "COMPLETED"
+        assert payload["run_id"] == "run-42"
+        assert payload["conversation_id"] == "conv-xyz"
+
+
+def test_callback_omits_conversation_id_when_not_registered(monkeypatch):
+    """Callback payload omits conversation_id when not registered."""
+    monkeypatch.setenv("AUTOMATION_CALLBACK_URL", "https://svc.test/complete")
+    monkeypatch.setenv("AUTOMATION_RUN_ID", "run-42")
+    ws = _make_local_workspace()
+
+    # Do not register a conversation
+
+    mock_resp = MagicMock()
+    mock_resp.status_code = 200
+
+    with patch("httpx.Client") as MockClient:
+        mock_client = MagicMock()
+        mock_client.post.return_value = mock_resp
+        mock_client.__enter__ = MagicMock(return_value=mock_client)
+        mock_client.__exit__ = MagicMock(return_value=False)
+        MockClient.return_value = mock_client
+
+        ws.__exit__(None, None, None)
+
+        # Check the POST payload does NOT include conversation_id
+        mock_client.post.assert_called_once()
+        payload = mock_client.post.call_args.kwargs["json"]
+        assert payload["status"] == "COMPLETED"
+        assert "conversation_id" not in payload