feat(chat): add keys to bridge

Made-with: Cursor

Author: scottcmg

src/uipath/_cli/_chat/_bridge.py

@@ -6,7 +6,7 @@
 import os
 import uuid
 from typing import Any
-from urllib.parse import urlparse
+from urllib.parse import urlencode, urlparse
 
 from uipath.core.chat import (
     UiPathConversationEvent,
@@ -18,6 +18,10 @@
     UiPathConversationToolCallConfirmationInterruptStartEvent,
     UiPathConversationToolCallConfirmationValue,
 )
+from uipath._utils.constants import (
+    ENV_FOLDER_KEY,
+    ENV_JOB_KEY,
+)
 from uipath.core.triggers import UiPathResumeTrigger
 from uipath.runtime.chat import UiPathChatProtocol
 from uipath.runtime.context import UiPathRuntimeContext
@@ -389,7 +393,6 @@ def get_chat_bridge(
     assert context.conversation_id is not None, "conversation_id must be set in context"
     assert context.exchange_id is not None, "exchange_id must be set in context"
 
-    # Extract host from UIPATH_URL
     base_url = os.environ.get("UIPATH_URL")
     if not base_url:
         raise RuntimeError(
@@ -401,9 +404,20 @@ def get_chat_bridge(
         raise RuntimeError(f"Invalid UIPATH_URL format: {base_url}")
 
     host = parsed.netloc
+    conversation_id = context.conversation_id
+    folder_key = os.environ.get(ENV_FOLDER_KEY)
+    job_key = os.environ.get(ENV_JOB_KEY)
+
+    # Build query params for CAS for use verifying the requesting agent should have access when runAsMe=False
+    query_params: dict[str, str] = {
+        "conversationId": conversation_id,
+        "folderKey": folder_key or "",
+        "jobKey": job_key or "",
+    }
+    query_params = {k: v for k, v in query_params.items() if v}
+    query_string = urlencode(query_params)
 
-    # Construct WebSocket URL for CAS
-    websocket_url = f"wss://{host}?conversationId={context.conversation_id}"
+    websocket_url = f"wss://{host}?{query_string}"
     websocket_path = "autopilotforeveryone_/websocket_/socket.io"
 
     if os.environ.get("CAS_WEBSOCKET_HOST"):
@@ -420,13 +434,13 @@ def get_chat_bridge(
         or os.environ.get("UIPATH_TENANT_ID", ""),
         "X-UiPath-Internal-AccountId": f"{context.org_id}"
         or os.environ.get("UIPATH_ORGANIZATION_ID", ""),
-        "X-UiPath-ConversationId": context.conversation_id,
+        "X-UiPath-ConversationId": conversation_id,
     }
 
     return SocketIOChatBridge(
         websocket_url=websocket_url,
         websocket_path=websocket_path,
-        conversation_id=context.conversation_id,
+        conversation_id=conversation_id,
         exchange_id=context.exchange_id,
         headers=headers,
     )

tests/cli/chat/test_bridge.py

@@ -165,6 +165,25 @@ def test_get_chat_bridge_includes_conversation_id_in_url(
 
         assert "conversationId=my-conversation-id" in bridge.websocket_url
 
+    def test_get_chat_bridge_includes_folder_key_job_key_when_set(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """folderKey and jobKey are included in URL when env vars are set (CAS contract)."""
+        monkeypatch.setenv("UIPATH_URL", "https://cloud.uipath.com/org/tenant")
+        monkeypatch.setenv("UIPATH_ACCESS_TOKEN", "test-token")
+        monkeypatch.setenv("UIPATH_FOLDER_KEY", "folder-guid-456")
+        monkeypatch.setenv("UIPATH_JOB_KEY", "job-guid-789")
+        monkeypatch.delenv("CAS_WEBSOCKET_HOST", raising=False)
+
+        context = MockRuntimeContext(conversation_id="conv-id")
+
+        bridge = cast(SocketIOChatBridge, get_chat_bridge(cast(Any, context)))
+
+        assert "conversationId=conv-id" in bridge.websocket_url
+        assert "folderKey=folder-guid-456" in bridge.websocket_url
+        assert "jobKey=job-guid-789" in bridge.websocket_url
+        assert "robotKey" not in bridge.websocket_url
+
 
 class TestGetChatBridge:
     """Tests for get_chat_bridge factory function."""