PC-5280: forward agent max_tokens to tool simulation, detect truncation

packages/uipath/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath"
-version = "2.14.44"
+version = "2.14.45"
 description = "Python SDK and CLI for UiPath Platform, enabling programmatic interaction with automation services, process management, and deployment tools."
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath/src/uipath/_cli/cli_debug.py

@@ -10,6 +10,7 @@
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli._utils._studio_project import StudioClient
 from uipath._cli._utils._tracing import create_trace_manager
+from uipath.eval.helpers import get_agent_max_tokens
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.eval.mocks._mock_runtime import load_simulation_config
 from uipath.platform.common import (
@@ -249,13 +250,14 @@ async def execute_debug_runtime():
                             agent_model = None
                             if schema.metadata and "settings" in schema.metadata:
                                 agent_model = schema.metadata["settings"].get("model")
+                            agent_max_tokens = get_agent_max_tokens(schema)
 
                             delegate_runtime: UiPathDebugRuntime | UiPathMockRuntime = (
                                 debug_runtime
                             )
                             if simulation_config:
                                 mocking_context = build_mocking_context(
-                                    simulation_config, agent_model
+                                    simulation_config, agent_model, agent_max_tokens
                                 )
                                 if mocking_context:
                                     delegate_runtime = UiPathMockRuntime(
@@ -264,7 +266,8 @@ async def execute_debug_runtime():
                                     )
                             else:
                                 mocking_context = load_simulation_config(
-                                    agent_model=agent_model
+                                    agent_model=agent_model,
+                                    agent_max_tokens=agent_max_tokens,
                                 )
                                 delegate_runtime = UiPathMockRuntime(
                                     delegate=debug_runtime,

packages/uipath/src/uipath/_cli/cli_run.py

@@ -9,6 +9,7 @@
 from uipath._cli._utils._common import read_resource_overwrites_from_file
 from uipath._cli._utils._debug import setup_debugging
 from uipath._cli._utils._tracing import create_trace_manager
+from uipath.eval.helpers import get_agent_max_tokens
 from uipath.eval.mocks import SimulationConfig, UiPathMockRuntime, build_mocking_context
 from uipath.platform.common import (
     ExecutionSourceContext,
@@ -316,7 +317,9 @@ async def execute() -> None:
                                         "model"
                                     )
                                 mocking_context = build_mocking_context(
-                                    simulation_config, agent_model
+                                    simulation_config,
+                                    agent_model,
+                                    get_agent_max_tokens(schema),
                                 )
                                 if mocking_context:
                                     runtime = UiPathMockRuntime(

packages/uipath/src/uipath/eval/helpers.py

@@ -302,3 +302,24 @@ def get_agent_model(schema: UiPathRuntimeSchema) -> str | None:
         return None
     except Exception:
         return None
+
+
+def get_agent_max_tokens(schema: UiPathRuntimeSchema) -> int | None:
+    """Get the agent's max output tokens from the runtime schema metadata.
+
+    Read from schema.metadata["settings"], populated by the low-code agents
+    runtime from agent.json (``max_tokens``, or ``maxTokens`` when dumped by
+    alias).
+
+    Returns:
+        The agent's max tokens, or None if not found.
+    """
+    try:
+        if schema.metadata and "settings" in schema.metadata:
+            settings = schema.metadata["settings"]
+            max_tokens = settings.get("max_tokens", settings.get("maxTokens"))
+            if isinstance(max_tokens, int) and not isinstance(max_tokens, bool):
+                return max_tokens if max_tokens > 0 else None
+        return None
+    except Exception:
+        return None

packages/uipath/src/uipath/eval/mocks/_mock_runtime.py

@@ -35,9 +35,16 @@
 
 
 def build_mocking_context(
-    config: SimulationConfig, agent_model: str | None = None
+    config: SimulationConfig,
+    agent_model: str | None = None,
+    agent_max_tokens: int | None = None,
 ) -> MockingContext | None:
-    """Build a MockingContext from a validated SimulationConfig."""
+    """Build a MockingContext from a validated SimulationConfig.
+
+    When the simulation falls back to the agent's model, the agent's max tokens
+    go with it, so the simulated tool output is not cut off at the chat
+    service's default.
+    """
     if not config.enabled:
         return None
 
@@ -68,9 +75,9 @@
     model = (
         ModelSettings(model=config.model)
         if config.model
-        else ModelSettings(model=agent_model)
+        else ModelSettings(model=agent_model, maxTokens=agent_max_tokens)
         if agent_model
         else None
     )
 
     mocking_strategy = LLMMockingStrategy(
@@ -91,7 +98,9 @@
 
 
 def build_mocking_context_from_dict(
-    simulation_data: dict[str, Any], agent_model: str | None = None
+    simulation_data: dict[str, Any],
+    agent_model: str | None = None,
+    agent_max_tokens: int | None = None,
 ) -> MockingContext | None:
     """Build a MockingContext from a simulation config dictionary.
 
@@ -100,15 +109,18 @@
     Args:
         simulation_data: Parsed simulation config (same schema as simulation.json).
         agent_model: Optional agent model name to use as fallback.
+        agent_max_tokens: Optional agent max tokens, used with agent_model.
 
     Returns:
         MockingContext if valid and enabled, None otherwise.
     """
     config = SimulationConfig.model_validate(simulation_data)
-    return build_mocking_context(config, agent_model)
+    return build_mocking_context(config, agent_model, agent_max_tokens)
 
 
-def load_simulation_config(agent_model: str | None = None) -> MockingContext | None:
+def load_simulation_config(
+    agent_model: str | None = None, agent_max_tokens: int | None = None
+) -> MockingContext | None:
     """Load simulation.json from current directory and convert to MockingContext.
 
     Returns:
@@ -124,7 +136,7 @@
         config = SimulationConfig.model_validate_json(
             simulation_path.read_text(encoding="utf-8")
         )
-        return build_mocking_context(config, agent_model)
+        return build_mocking_context(config, agent_model, agent_max_tokens)
 
     except Exception as e:
         logger.warning(f"Failed to load simulation.json: {e}")

packages/uipath/src/uipath/eval/mocks/_structured_output.py

@@ -111,7 +111,8 @@ def extract_response(response: Any) -> Any:
 
     Raises:
         ValueError: if the response carries no usable tool call or is missing the
-            wrapped ``response`` key.
+            wrapped ``response`` key (reported as a truncation when the model
+            stopped at its max_tokens limit).
     """
     choices = getattr(response, "choices", None)
     if not choices:
@@ -126,6 +127,12 @@ def extract_response(response: Any) -> Any:
 
     arguments = tool_calls[0].arguments
     if RESPONSE_KEY not in arguments:
+        if getattr(choices[0], "finish_reason", None) == "length":
+            raise ValueError(
+                "LLM response was truncated at the max_tokens limit "
+                "(finish_reason='length') before the simulated tool output was "
+                f"complete; increase the model's max tokens. Arguments: {arguments}"
+            )
         raise ValueError(
             f"Tool call arguments missing '{RESPONSE_KEY}' key: {arguments}"
         )
@@ -193,6 +200,11 @@ async def generate(
             choices = getattr(response, "choices", None)
             if choices:
                 content = choices[0].message.content
+                if getattr(choices[0], "finish_reason", None) == "length":
+                    logger.warning(
+                        "response_format output was truncated at the max_tokens "
+                        "limit; falling back to tools"
+                    )
         except Exception as e:
             logger.info("response_format path failed, falling back to tools: %s", e)
 

packages/uipath/src/uipath/eval/runtime/runtime.py

@@ -57,7 +57,7 @@
 )
 from ..evaluators.exact_match_evaluator import ExactMatchEvaluatorConfig
 from ..evaluators.output_evaluator import OutputEvaluationCriteria
-from ..helpers import get_agent_model
+from ..helpers import get_agent_max_tokens, get_agent_model
 from ..mocks._cache_manager import CacheManager
 from ..mocks._input_mocker import (
     generate_llm_input,
@@ -662,7 +662,14 @@ async def _execute_eval(
                         mocking_model = get_agent_model(self.context.runtime_schema)
                         if mocking_model:
                             mocking_strategy = mocking_strategy.model_copy(
-                                update={"model": ModelSettings(model=mocking_model)}
+                                update={
+                                    "model": ModelSettings(
+                                        model=mocking_model,
+                                        maxTokens=get_agent_max_tokens(
+                                            self.context.runtime_schema
+                                        ),
+                                    )
+                                }
                             )
 
                     agent_execution_output = await self.execute_runtime(

packages/uipath/tests/cli/eval/mocks/test_mock_runtime_model_settings.py

@@ -0,0 +1,60 @@
+"""Tests for the model settings the legacy tool-simulation mocker is built with."""
+
+from uipath.eval.mocks._mock_runtime import build_mocking_context
+from uipath.eval.mocks._simulation_model import simulation_completion_kwargs
+from uipath.eval.mocks._types import LLMMockingStrategy, SimulationConfig
+
+
+def _legacy_config(**overrides: object) -> SimulationConfig:
+    return SimulationConfig.model_validate(
+        {
+            "enabled": True,
+            "instructions": "Simulate it",
+            "toolsToSimulate": [{"name": "Get Records"}],
+            **overrides,
+        }
+    )
+
+
+def _strategy(context: object) -> LLMMockingStrategy:
+    strategy = getattr(context, "strategy", None)
+    assert isinstance(strategy, LLMMockingStrategy)
+    return strategy
+
+
+def test_agent_model_fallback_carries_agent_max_tokens():
+    context = build_mocking_context(
+        _legacy_config(), agent_model="gpt-4.1-2025-04-14", agent_max_tokens=16384
+    )
+
+    strategy = _strategy(context)
+    assert strategy.model is not None
+    assert strategy.model.model == "gpt-4.1-2025-04-14"
+    assert strategy.model.max_tokens == 16384
+    # The simulation call goes out at the agent's limit, not the 4096 default.
+    assert simulation_completion_kwargs(strategy.model) == {
+        "model": "gpt-4.1-2025-04-14",
+        "max_tokens": 16384,
+    }
+
+
+def test_agent_model_fallback_without_max_tokens_omits_it():
+    context = build_mocking_context(_legacy_config(), agent_model="gpt-4.1-2025-04-14")
+
+    strategy = _strategy(context)
+    assert strategy.model is not None
+    assert "max_tokens" not in simulation_completion_kwargs(strategy.model)
+
+
+def test_simulation_model_override_does_not_take_agent_max_tokens():
+    # A different simulation model may have a lower output cap than the agent's.
+    context = build_mocking_context(
+        _legacy_config(model="claude-haiku-4-5"),
+        agent_model="gpt-4.1-2025-04-14",
+        agent_max_tokens=16384,
+    )
+
+    strategy = _strategy(context)
+    assert strategy.model is not None
+    assert strategy.model.model == "claude-haiku-4-5"
+    assert strategy.model.max_tokens is None

packages/uipath/tests/cli/eval/mocks/test_structured_output.py

@@ -16,8 +16,14 @@
 )
 
 
-def _response(message: SimpleNamespace | None) -> SimpleNamespace:
-    choices = [] if message is None else [SimpleNamespace(message=message)]
+def _response(
+    message: SimpleNamespace | None, finish_reason: str | None = None
+) -> SimpleNamespace:
+    choices = (
+        []
+        if message is None
+        else [SimpleNamespace(message=message, finish_reason=finish_reason)]
+    )
     return SimpleNamespace(choices=choices)
 
 
@@ -124,6 +130,48 @@
         extract_response(_response(message))
 
 
+def test_extract_response_reports_truncation_at_max_tokens():
+    # A forced tool call cut off at max_tokens comes back with empty arguments.
+    message = SimpleNamespace(content=None, tool_calls=[SimpleNamespace(arguments={})])
+    with pytest.raises(ValueError, match="truncated at the max_tokens limit"):
+        extract_response(_response(message, finish_reason="length"))
+
+
+def test_extract_response_returns_value_when_finish_reason_is_length():
+    # Complete arguments are still returned; only a missing key is a truncation.
+    message = SimpleNamespace(
+        content=None,
+        tool_calls=[SimpleNamespace(arguments={RESPONSE_KEY: {"a": 1}})],
+    )
+    assert extract_response(_response(message, finish_reason="length")) == {"a": 1}
+
+
+@pytest.mark.asyncio
+async def test_generate_structured_output_reports_truncation_on_both_paths():
+    truncated_content = SimpleNamespace(content='{"records": [{"id": 1}, {"id"')
+    truncated_tool_call = SimpleNamespace(
+        content=None, tool_calls=[SimpleNamespace(arguments={})]
+    )
+    llm = _FakeLLM(
+        [
+            _response(truncated_content, finish_reason="length"),
+            _response(truncated_tool_call, finish_reason="length"),
+        ]
+    )
+
+    with pytest.raises(ValueError, match="truncated at the max_tokens limit"):
+        await generate_structured_output(
+            llm,
+            [{"role": "user", "content": "hi"}],
+            schema={"type": "object"},
+            response_format_name="OutputSchema",
+            description="desc",
+            completion_kwargs={"model": "gpt-4.1-2025-04-14", "max_tokens": 4096},
+        )
+    assert len(llm.calls) == 2
+    assert all(call["max_tokens"] == 4096 for call in llm.calls)
+
+
 @pytest.mark.asyncio
 async def test_generate_structured_output_prefers_response_format_content():
     # OpenAI returns content via response_format; no fallback call is made.

packages/uipath/tests/cli/eval/test_eval_runtime_metadata.py

@@ -12,7 +12,7 @@
 
 from uipath.core.events import EventBus
 from uipath.core.tracing import UiPathTraceManager
-from uipath.eval.helpers import get_agent_model
+from uipath.eval.helpers import get_agent_max_tokens, get_agent_model
 from uipath.eval.runtime import UiPathEvalContext, UiPathEvalRuntime
 from uipath.runtime import (
     UiPathExecuteOptions,
@@ -149,6 +149,35 @@ async def test_returns_model_consistently(self):
         assert model1 == model2 == "consistent-model"
 
 
+class TestGetAgentMaxTokens:
+    """Tests for get_agent_max_tokens function."""
+
+    def test_returns_snake_case_max_tokens(self):
+        """The agents runtime dumps settings by field name (max_tokens)."""
+        schema = MockRuntimeSchema()
+        schema.metadata = {"settings": {"model": "gpt-4.1", "max_tokens": 16384}}
+
+        assert get_agent_max_tokens(schema) == 16384
+
+    def test_returns_aliased_max_tokens(self):
+        schema = MockRuntimeSchema()
+        schema.metadata = {"settings": {"model": "gpt-4.1", "maxTokens": 8192}}
+
+        assert get_agent_max_tokens(schema) == 8192
+
+    def test_returns_none_when_missing(self):
+        schema = MockRuntimeSchema()
+
+        assert get_agent_max_tokens(schema) is None
+
+    @pytest.mark.parametrize("value", [0, -1, "16384", True, None])
+    def test_returns_none_for_unusable_values(self, value: Any):
+        schema = MockRuntimeSchema()
+        schema.metadata = {"settings": {"max_tokens": value}}
+
+        assert get_agent_max_tokens(schema) is None
+
+
 class TestGetSchema:
     """Tests for get_schema method."""