diff --git a/src/uipath_langchain/_utils/_request_mixin.py b/src/uipath_langchain/_utils/_request_mixin.py index b787ee715..79025644c 100644 --- a/src/uipath_langchain/_utils/_request_mixin.py +++ b/src/uipath_langchain/_utils/_request_mixin.py @@ -4,6 +4,7 @@ import os import time from typing import Any, AsyncIterator, Dict, Iterator, Mapping +from urllib.parse import quote import httpx import openai @@ -79,7 +80,7 @@ class UiPathRequestMixin(BaseModel): default_headers: Mapping[str, str] | None = { "X-UiPath-Streaming-Enabled": "false", "X-UiPath-JobKey": os.getenv("UIPATH_JOB_KEY", ""), - "X-UiPath-ProcessKey": os.getenv("UIPATH_PROCESS_KEY", ""), + "X-UiPath-ProcessKey": quote(os.getenv("UIPATH_PROCESS_KEY", ""), safe=""), } model_name: str | None = Field( default_factory=lambda: os.getenv( diff --git a/src/uipath_langchain/chat/bedrock.py b/src/uipath_langchain/chat/bedrock.py index 6ab68e20d..eac216506 100644 --- a/src/uipath_langchain/chat/bedrock.py +++ b/src/uipath_langchain/chat/bedrock.py @@ -2,6 +2,7 @@ import os from collections.abc import Iterator from typing import Any, Optional +from urllib.parse import quote from langchain_core.callbacks import CallbackManagerForLLMRun from langchain_core.messages import BaseMessage @@ -143,7 +144,7 @@ def _modify_request(self, request, **kwargs): if job_key: headers["X-UiPath-JobKey"] = job_key if process_key: - headers["X-UiPath-ProcessKey"] = process_key + headers["X-UiPath-ProcessKey"] = quote(process_key, safe="") request.headers.update(headers) diff --git a/src/uipath_langchain/chat/openai.py b/src/uipath_langchain/chat/openai.py index 3025162a7..7167a9b78 100644 --- a/src/uipath_langchain/chat/openai.py +++ b/src/uipath_langchain/chat/openai.py @@ -1,6 +1,7 @@ import logging import os from typing import Optional +from urllib.parse import quote import httpx from langchain_openai import AzureChatOpenAI @@ -162,7 +163,7 @@ def _build_headers(self, token: str) -> dict[str, str]: if job_key := os.getenv("UIPATH_JOB_KEY"): headers["X-UiPath-JobKey"] = job_key if process_key := os.getenv("UIPATH_PROCESS_KEY"): - headers["X-UiPath-ProcessKey"] = process_key + headers["X-UiPath-ProcessKey"] = quote(process_key, safe="") # Allow extra_headers to override defaults headers.update(self._extra_headers) diff --git a/src/uipath_langchain/chat/vertex.py b/src/uipath_langchain/chat/vertex.py index c93133ccb..770c352ca 100644 --- a/src/uipath_langchain/chat/vertex.py +++ b/src/uipath_langchain/chat/vertex.py @@ -2,6 +2,7 @@ import os from collections.abc import AsyncIterator, Iterator from typing import Any, Optional +from urllib.parse import quote import httpx from langchain_core.callbacks import ( @@ -266,7 +267,7 @@ def _build_headers( if job_key := os.getenv("UIPATH_JOB_KEY"): headers["X-UiPath-JobKey"] = job_key if process_key := os.getenv("UIPATH_PROCESS_KEY"): - headers["X-UiPath-ProcessKey"] = process_key + headers["X-UiPath-ProcessKey"] = quote(process_key, safe="") return headers @staticmethod diff --git a/tests/chat/test_header_encoding.py b/tests/chat/test_header_encoding.py new file mode 100644 index 000000000..a0752e90e --- /dev/null +++ b/tests/chat/test_header_encoding.py @@ -0,0 +1,133 @@ +import os +from unittest.mock import MagicMock, patch +from urllib.parse import quote + +import pytest + +from uipath_langchain.chat.openai import UiPathChatOpenAI + +NON_ASCII_PROCESS_KEY = "Solution.17.agent.GetCompanyIdAgent-請-test" +ASCII_PROCESS_KEY = "Solution.17.agent.MyAgent-test" + +BASE_ENV = { + "UIPATH_URL": "https://cloud.uipath.com/org/tenant", + "UIPATH_ORGANIZATION_ID": "org-id", + "UIPATH_TENANT_ID": "tenant-id", + "UIPATH_ACCESS_TOKEN": "test-token", +} + + +class TestOpenAIHeaderEncoding: + """Verify UiPathChatOpenAI percent-encodes non-ASCII header values.""" + + def _build_headers_with_process_key(self, process_key: str) -> dict[str, str]: + env = {**BASE_ENV, "UIPATH_PROCESS_KEY": process_key} + with patch.dict(os.environ, env, clear=False): + obj = object.__new__(UiPathChatOpenAI) + obj._agenthub_config = None + obj._byo_connection_id = None + obj._extra_headers = {} + return obj._build_headers("fake-token") + + def test_ascii_process_key_unchanged(self) -> None: + headers = self._build_headers_with_process_key(ASCII_PROCESS_KEY) + assert headers["X-UiPath-ProcessKey"] == quote(ASCII_PROCESS_KEY, safe="") + + def test_non_ascii_process_key_encoded(self) -> None: + headers = self._build_headers_with_process_key(NON_ASCII_PROCESS_KEY) + value = headers["X-UiPath-ProcessKey"] + assert "請" not in value + assert value == quote(NON_ASCII_PROCESS_KEY, safe="") + assert "%E8%AB%8B" in value + + def test_header_value_is_ascii_safe(self) -> None: + headers = self._build_headers_with_process_key(NON_ASCII_PROCESS_KEY) + value = headers["X-UiPath-ProcessKey"] + value.encode("ascii") + + def test_missing_process_key_omitted(self) -> None: + env = {**BASE_ENV} + env.pop("UIPATH_PROCESS_KEY", None) + with patch.dict(os.environ, env, clear=True): + obj = object.__new__(UiPathChatOpenAI) + obj._agenthub_config = None + obj._byo_connection_id = None + obj._extra_headers = {} + headers = obj._build_headers("fake-token") + assert "X-UiPath-ProcessKey" not in headers + + +class TestVertexHeaderEncoding: + """Verify UiPathChatVertex._build_headers percent-encodes non-ASCII header values.""" + + @pytest.fixture(autouse=True) + def _skip_if_no_google(self) -> None: + pytest.importorskip("google.genai", reason="google-genai not installed") + + def test_non_ascii_process_key_encoded(self) -> None: + from uipath_langchain.chat.vertex import UiPathChatVertex + + env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY} + with patch.dict(os.environ, env, clear=False): + headers = UiPathChatVertex._build_headers("fake-token") + value = headers["X-UiPath-ProcessKey"] + assert "請" not in value + assert "%E8%AB%8B" in value + value.encode("ascii") + + def test_ascii_process_key_unchanged(self) -> None: + from uipath_langchain.chat.vertex import UiPathChatVertex + + env = {**BASE_ENV, "UIPATH_PROCESS_KEY": ASCII_PROCESS_KEY} + with patch.dict(os.environ, env, clear=False): + headers = UiPathChatVertex._build_headers("fake-token") + assert headers["X-UiPath-ProcessKey"] == quote(ASCII_PROCESS_KEY, safe="") + + +class TestBedrockHeaderEncoding: + """Verify AwsBedrockCompletionsPassthroughClient percent-encodes non-ASCII header values.""" + + def test_non_ascii_process_key_encoded(self) -> None: + pytest.importorskip("botocore", reason="botocore not installed") + from uipath_langchain.chat.bedrock import AwsBedrockCompletionsPassthroughClient + + env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY} + with ( + patch.dict(os.environ, env, clear=False), + patch( + "uipath_langchain.chat.bedrock.boto3.client", return_value=MagicMock() + ), + ): + client = AwsBedrockCompletionsPassthroughClient( + model="test-model", + token="fake-token", + api_flavor="converse", + ) + request = MagicMock() + request.url = "https://example.com/converse" + request.headers = {} + client._modify_request(request) + + value = request.headers["X-UiPath-ProcessKey"] + assert "請" not in value + assert "%E8%AB%8B" in value + value.encode("ascii") + + +class TestRequestMixinHeaderEncoding: + """Verify UiPathRequestMixin default_headers percent-encodes non-ASCII values.""" + + def test_non_ascii_process_key_encoded_in_defaults(self) -> None: + env = {**BASE_ENV, "UIPATH_PROCESS_KEY": NON_ASCII_PROCESS_KEY} + with patch.dict(os.environ, env, clear=False): + import importlib + + import uipath_langchain._utils._request_mixin as mod + + importlib.reload(mod) + value = mod.UiPathRequestMixin.model_fields["default_headers"].default[ + "X-UiPath-ProcessKey" + ] + assert "請" not in value + assert "%E8%AB%8B" in value + value.encode("ascii")