From c8cad0dcc41eb7550df8f2541b74592b3222af13 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:15:13 +0300 Subject: [PATCH 01/18] feat(cuga-lite): native function-calling execution mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `advanced_features.cuga_lite_execution_mode = "function_calling"` (default "codeact"; the default path is byte-identical). In this mode the model emits native tool_calls instead of a Python block: - `tool_exec` node: runs each tool_call and replies with a ToolMessage that carries the call's id, under the same envelope as the sandbox (tracker session, run/thread/block budgets, tool_call_timeout, budget fields on every exit). Every failure — unknown tool, bad arguments, a tool exception, a timeout, a spent budget, provider-side invalid_tool_calls — is a ToolMessage, never an exception, and one failing call does not abort its siblings. - one seam in the shared call_model, `execute_call_model_fc`: None on the base adapter (Supervisor and every CodeAct run unchanged); CugaLite builds real message objects (system prompt, few-shot demos, sanitized history) so tool_calls survive the round trip, then routes on the response: tool_calls -> tool_exec, text -> END. A fenced code block with no tool_calls is a mode violation: never executed, one corrective turn. - the node is always wired on the CugaLite graph and dormant in codeact, so the mode can be switched per invoke without recompiling. - a short dedicated prompt (`prompts/fc_prompt.jinja2`) replaces the CodeAct prompt in this mode; tools travel natively via bind_tools, whose mode is upgraded from "none" to "all" only when it resolved to none. Opt-in fragments via `cuga_lite_fc_prompt_fragments` (`evidence_first`). - configurable per invoke (`cuga_lite_execution_mode`), per model profile, in settings.toml, and from the SDK (`execution_mode=` on the constructor, invoke() and stream()). Not in this change: human-in-the-loop tool approval is not wired on the function-calling path; tools run in-process and sequentially. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../cuga_agent_core/graph/graph_nodes.py | 28 ++ .../cuga_agent_core/graph/shared_graph.py | 17 +- .../cuga_agent_core/graph/shared_nodes.py | 21 ++ .../tests/graph/test_call_model_fc_seam.py | 138 ++++++++ .../nodes/cuga_lite/adapter/graph_adapter.py | 250 +++++++++++++- .../nodes/cuga_lite/adapter/prepare_node.py | 28 +- .../nodes/cuga_lite/adapter/tool_exec_node.py | 263 ++++++++++++++ .../nodes/cuga_lite/cuga_lite_graph.py | 5 + .../nodes/cuga_lite/model_runtime_profile.py | 122 +++++++ .../nodes/cuga_lite/prompt_utils.py | 47 +++ .../nodes/cuga_lite/prompts/fc_prompt.jinja2 | 21 ++ .../tests/test_execution_mode_config.py | 239 +++++++++++++ .../tests/test_fc_call_model_hook.py | 204 +++++++++++ .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 236 +++++++++++++ .../nodes/cuga_lite/tests/test_fc_prompt.py | 53 +++ .../cuga_lite/tests/test_tool_exec_node.py | 323 ++++++++++++++++++ src/cuga/config.py | 3 + src/cuga/sdk.py | 34 ++ src/cuga/settings.toml | 7 + .../test_variables_addendum_not_persisted.py | 1 + 20 files changed, 2034 insertions(+), 6 deletions(-) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py index f92620abf..10a000b72 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py @@ -133,6 +133,34 @@ async def ainvoke_model(self, bound: Any, messages: list, invoke_config: dict) - """ return await bound.ainvoke(messages, config=invoke_config) + async def execute_call_model_fc( + self, + *, + state: Any, + config: Any, + configurable: dict, + active_model: Any, + bound: Any, + invoke_config: dict, + system_content: str, + modified_messages: list, + budget_exhausted: bool, + playbook_fired: bool, + ) -> Optional[Command]: + """Native function-calling seam. + + Called by the shared ``call_model`` after the bound model is resolved and + before it is invoked. Returning ``None`` means "not my turn": the CodeAct + path below runs exactly as before. Returning a ``Command`` short-circuits + ``call_model`` entirely. + + The base returns ``None``, so the Supervisor graph and every CodeAct run + never enter function-calling code. CugaLite overrides this to run the + native ``tool_calls`` -> ``ToolMessage`` loop when + ``cuga_lite_execution_mode`` resolves to ``function_calling``. + """ + return None + # ── Post-invocation hooks ────────────────────────────────────────────── def normalize_response(self, response: Any) -> Tuple[str, Optional[str]]: diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py index c367e5556..acf1c2f4f 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py @@ -4,6 +4,9 @@ START → prepare --Command--> call_model ↔ execute (loop) → END +CugaLite may add a fourth node, ``tool_exec`` (native function-calling +mode), which loops back into ``call_model`` the same way. + Both CugaLite and CugaSupervisor share this structure. The nodes themselves are provided by the caller (produced by adapter factories), so the graph builder stays graph-agnostic. @@ -15,7 +18,7 @@ from __future__ import annotations -from typing import Callable, Type +from typing import Callable, Optional, Type from langgraph.graph import START, StateGraph @@ -29,6 +32,7 @@ def build_agent_graph( prepare_node: Callable, call_model_node: Callable, execute_node: Callable, + tool_exec_node: Optional[Callable] = None, ) -> StateGraph: """Wire and return an UNCOMPILED 3-node agent StateGraph. @@ -42,6 +46,11 @@ def build_agent_graph( call_model_node: Async node function for the call_model step (use ``create_call_model_node`` from ``shared_nodes.py``). execute_node: Async node function for the execute/sandbox step. + tool_exec_node: Optional node that executes native ``tool_calls`` and + replies with ``ToolMessage``s (CugaLite function-calling mode). When + supplied it is added as ``"tool_exec"`` with a static edge back to + ``call_model``; ``call_model`` only routes to it in that mode, so it + is dormant on every CodeAct run. Omitted by the Supervisor graph. Returns: An uncompiled ``StateGraph``. Call ``.compile(checkpointer=...)`` @@ -58,4 +67,10 @@ def build_agent_graph( # Execute node returns a state update (not Command); loop back for the NL answer. graph.add_edge(adapter.execute_node_name, "call_model") + if tool_exec_node is not None: + graph.add_node("tool_exec", tool_exec_node) + # Same shape as the sandbox edge: the node returns a state update, and the + # loop closes back into call_model for the next model turn. + graph.add_edge("tool_exec", "call_model") + return graph diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py index 395fddfff..f7abdcaf9 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py @@ -226,6 +226,27 @@ async def call_model(state: Any, config: RunnableConfig = None) -> Command: # Pass the full node config so LangChain keeps parent_run_id linkage for # Langfuse. Passing only {"callbacks": [...]} starts orphan root traces. invoke_config = config if config is not None else {} + + # ── Native function-calling seam ─────────────────────────────────── + # The base adapter returns None, so this is a no-op for the Supervisor + # graph and for every CodeAct run. CugaLite returns a Command here when + # cuga_lite_execution_mode resolves to function_calling, and the CodeAct + # path below is skipped entirely — never blended. + fc_command = await adapter.execute_call_model_fc( + state=state, + config=config, + configurable=configurable, + active_model=active_model, + bound=bound, + invoke_config=invoke_config, + system_content=system_content, + modified_messages=modified_messages, + budget_exhausted=budget_exhausted, + playbook_fired=playbook_fired, + ) + if fc_command is not None: + return fc_command + response = await adapter.ainvoke_model(bound, messages_for_model, invoke_config) # ── Normalise response ───────────────────────────────────────────── diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py new file mode 100644 index 000000000..316c2bc96 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py @@ -0,0 +1,138 @@ +"""The ``execute_call_model_fc`` seam in the shared ``call_model`` node. + +The base adapter returns ``None`` and the CodeAct path runs exactly as before — +that is the Supervisor graph and every default CugaLite run. An adapter that +returns a ``Command`` short-circuits the node: the CodeAct invocation, response +normalisation and routing never run, so the two paths are never blended. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_core.messages import BaseMessage, HumanMessage +from langgraph.graph import END +from langgraph.types import Command + +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import CoreGraphAdapter +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import create_call_model_node + +pytestmark = pytest.mark.unit + +_SUMMARIZE = "cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes.apply_context_summarization" +_BOUND = object() # what resolve_bind_tools hands back; the hook must receive exactly this + + +class _BaseAdapter(CoreGraphAdapter): + messages_key = "chat_messages" + execute_node_name = "sandbox" + metadata_key = "cuga_lite_metadata" + + def get_messages(self, state: Any) -> List[BaseMessage]: + return list(state.chat_messages or []) + + def resolve_max_steps(self, state: Any, override: Optional[int]) -> int: + return override if override is not None else 50 + + async def resolve_bind_tools(self, state, model, configurable, config): + return model + + +class _ShortCircuitAdapter(_BaseAdapter): + def __init__(self): + self.seen_kwargs = None + + async def resolve_bind_tools(self, state, model, configurable, config): + return _BOUND + + async def execute_call_model_fc(self, **kwargs) -> Optional[Command]: + self.seen_kwargs = kwargs + return Command(goto="tool_exec", update={"chat_messages": [], "script": None}) + + +def _make_state(): + vm = MagicMock() + vm.get_variable_names.return_value = [] + return SimpleNamespace( + chat_messages=[HumanMessage(content="do task")], + step_count=0, + prepared_prompt="You are a helpful agent.", + cuga_lite_metadata={}, + variables_storage=None, + variable_counter_state=None, + variable_creation_order=None, + variables_manager=vm, + tool_budget_exhausted=False, + ) + + +def _mock_model(content: str): + model = MagicMock() + model.ainvoke = AsyncMock(return_value=SimpleNamespace(content=content, additional_kwargs={})) + return model + + +def _mock_settings(): + return SimpleNamespace( + advanced_features=SimpleNamespace(cuga_lite_max_steps=50), + policy=SimpleNamespace(enabled=False), + ) + + +@pytest.mark.asyncio +async def test_base_hook_returns_none(): + assert ( + await _BaseAdapter().execute_call_model_fc( + state=None, + config=None, + configurable={}, + active_model=None, + bound=None, + invoke_config={}, + system_content="", + modified_messages=[], + budget_exhausted=False, + playbook_fired=False, + ) + is None + ) + + +@pytest.mark.asyncio +@patch(_SUMMARIZE, new_callable=AsyncMock) +async def test_none_from_the_hook_leaves_the_codeact_path_untouched(mock_summarize): + mock_summarize.side_effect = lambda messages, *a, **kw: messages + model = _mock_model("Plain answer, no code.") + node = create_call_model_node(_BaseAdapter(), model, _mock_settings()) + + result = await node(_make_state(), config=None) + + assert model.ainvoke.await_count == 1, "the CodeAct invocation must still happen" + assert isinstance(result, Command) and result.goto == END + assert result.update["final_answer"] == "Plain answer, no code." + + +@pytest.mark.asyncio +@patch(_SUMMARIZE, new_callable=AsyncMock) +async def test_a_command_from_the_hook_short_circuits_the_node(mock_summarize): + mock_summarize.side_effect = lambda messages, *a, **kw: messages + adapter = _ShortCircuitAdapter() + model = _mock_model("this must never be requested through the CodeAct path") + node = create_call_model_node(adapter, model, _mock_settings()) + + result = await node(_make_state(), config={"configurable": {"cuga_lite_max_steps": 7}}) + + assert isinstance(result, Command) and result.goto == "tool_exec" + assert model.ainvoke.await_count == 0, "the CodeAct invocation ran alongside the FC turn" + + kw = adapter.seen_kwargs + assert kw["bound"] is _BOUND, "the hook must get the bind_tools-resolved model" + assert isinstance(kw["system_content"], str) and "helpful agent" in kw["system_content"] + assert all(isinstance(m, BaseMessage) for m in kw["modified_messages"]), ( + "history as message objects, not dicts" + ) + assert kw["configurable"] == {"cuga_lite_max_steps": 7} + assert kw["budget_exhausted"] is False and kw["playbook_fired"] is False diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 5ba3fecd0..943ff0df8 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -6,9 +6,12 @@ from __future__ import annotations +import json from typing import Any, Dict, List, Optional, Tuple -from langchain_core.messages import BaseMessage, HumanMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage +from langgraph.graph import END +from langgraph.types import Command from loguru import logger from cuga.backend.activity_tracker.tracker import Step @@ -19,15 +22,26 @@ from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( EXECUTION_OUTPUT_PREFIX, CoreGraphAdapter, + enforce_step_limit, ) -from cuga.backend.cuga_graph.utils.harmony import strip_harmony_tokens +from cuga.backend.cuga_graph.utils.harmony import contains_harmony_tokens, strip_harmony_tokens from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.prepare_node import create_prepare_tools_and_apps_node from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.response_utils import ( clean_empty_response_retry_meta, extract_code_from_response_tool_calls, ) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import create_sandbox_node -from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import resolve_model_with_bind_tools +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import create_tool_exec_node +from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import ( + _bind_tools_mode_from_settings, + resolve_model_with_bind_tools, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( + EXECUTION_MODE_FUNCTION_CALLING, + resolve_execution_mode, + resolved_runtime_model_name, + runtime_defaults_for_model, +) from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.find_tools import _first_user_message_text from cuga.backend.cuga_graph.nodes.cuga_lite.nl_auto_continue_classifier import ( BLOCKED_CLAIM_CORRECTION, @@ -40,6 +54,64 @@ from cuga.config import settings +_REASONING_KEYS = ("reasoning_content", "reasoning") + +FC_MODE_VIOLATION_CORRECTION = ( + "You wrote a code block, but this session executes tools through native function-calling only. " + "Code is never run here. Issue the tool call natively instead, or give the final answer as plain text." +) + + +def _sanitize_for_replay(messages: List[BaseMessage]) -> List[BaseMessage]: + """Copy history for replay to the provider. + + Assistant turns are replayed verbatim (``tool_calls`` intact) — but their + ``reasoning_content`` is dropped and harmony framing stripped from string + content: gpt-oss emits both, and strict OpenAI-compatible proxies 400 when + they come back in a later round. Irrelevant in CodeAct, whose history is + flattened to text before it leaves; in function-calling mode the raw turns + are replayed from round 2 onward, so it matters here. + """ + out: List[BaseMessage] = [] + for m in messages: + ak = getattr(m, "additional_kwargs", None) or {} + drop_reasoning = any(k in ak for k in _REASONING_KEYS) + content = getattr(m, "content", None) + new_content = ( + strip_harmony_tokens(content) if isinstance(content, str) and "<|" in content else content + ) + if not drop_reasoning and new_content is content: + out.append(m) + continue + update: Dict[str, Any] = {} + if drop_reasoning: + update["additional_kwargs"] = {k: v for k, v in ak.items() if k not in _REASONING_KEYS} + if new_content is not content: + update["content"] = new_content + try: + out.append(m.model_copy(update=update)) + except Exception: + out.append(m) + return out + + +def _few_shot_to_messages(few_shot: List[Any]) -> List[BaseMessage]: + """The prepare node's normalized ``{role, content}`` demos, as real chat messages.""" + msgs: List[BaseMessage] = [] + for ex in few_shot or []: + if not isinstance(ex, dict): + continue + role = (ex.get("role") or "").strip().lower() + content = ex.get("content") or "" + if not content: + continue + if role in ("user", "human"): + msgs.append(HumanMessage(content=content)) + elif role in ("assistant", "ai"): + msgs.append(AIMessage(content=content)) + return msgs + + def _format_observed_tool_shapes_block(shapes: Dict[str, str]) -> str: lines = ["", "---", "", "## Observed tool output shapes (this session)", ""] for name, description in shapes.items(): @@ -153,6 +225,14 @@ async def resolve_bind_tools( configurable: dict, config: Any = None, ) -> Any: + # Function-calling mode is inert without advertised tools, and the shipped + # default is bind mode "none". Upgrade only when the *resolved* mode is none + # — an explicit non-none choice from configurable, a model profile or + # settings is respected. Additive: no effect in codeact. + if self._execution_mode(configurable) == EXECUTION_MODE_FUNCTION_CALLING: + if self._resolved_bind_mode(configurable) == "none": + configurable = {**(configurable or {}), "cuga_lite_bind_tools_mode": "all"} + logger.info("[fc] bind_tools mode was 'none'; upgraded to 'all' for function-calling mode") try: return await resolve_model_with_bind_tools( active_model, @@ -240,6 +320,170 @@ def _any_execution_ran(self, state: Any) -> bool: return True return False + # ── Native function-calling mode ────────────────────────────────────── + + def _runtime_model_name(self, configurable: dict) -> str: + return resolved_runtime_model_name( + configurable_llm=(configurable or {}).get("llm"), graph_default_model=self._model + ) + + def _execution_mode(self, configurable: dict) -> str: + return resolve_execution_mode(configurable, self._runtime_model_name(configurable)) + + def _resolved_bind_mode(self, configurable: dict) -> str: + cfg = configurable or {} + prof = runtime_defaults_for_model(self._runtime_model_name(configurable)) + for candidate in (cfg.get("cuga_lite_bind_tools_mode"), prof.get("cuga_lite_bind_tools_mode")): + if candidate is not None and str(candidate).strip(): + return str(candidate).strip().lower() + return _bind_tools_mode_from_settings() + + async def execute_call_model_fc( + self, + *, + state: Any, + config: Any, + configurable: dict, + active_model: Any, + bound: Any, + invoke_config: dict, + system_content: str, + modified_messages: list, + budget_exhausted: bool, + playbook_fired: bool, + ) -> Optional[Command]: + """Function-calling turn: invoke with real message objects, route on ``tool_calls``. + + Returns ``None`` in codeact mode so the shared CodeAct path runs untouched. + Otherwise builds the outbound list from ``system_content`` (the FC prompt + that ``prepare`` selected), the few-shot demos as chat messages, and the + sanitized history — hands them straight to ``bound.ainvoke`` so the shared + dict serializer (which flattens ``tool_calls``) is never involved — then: + + - ``tool_calls`` present -> ``Command(goto="tool_exec")``, assistant turn kept verbatim + - otherwise -> final answer, ``END`` + - a fenced code block with no ``tool_calls`` is a mode violation: it is + never executed; the model gets one corrective turn instead. + """ + if self._execution_mode(configurable) != EXECUTION_MODE_FUNCTION_CALLING: + return None + + from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import ( + TOOL_BUDGET_EXHAUSTED_INSTRUCTION, + ) + + msgs: List[BaseMessage] = [SystemMessage(content=system_content)] + msgs.extend(_few_shot_to_messages(self.get_few_shot_messages(state))) + msgs.extend(_sanitize_for_replay(list(modified_messages))) + if budget_exhausted: + # Outbound only, like call_model's CodeAct path: never persisted. + msgs.append(HumanMessage(content=TOOL_BUDGET_EXHAUSTED_INSTRUCTION)) + + clamp_watsonx_completion_for_messages(bound, msgs) + response = await bound.ainvoke(msgs, config=invoke_config) + + tool_calls = list(getattr(response, "tool_calls", None) or []) + invalid_tool_calls = list(getattr(response, "invalid_tool_calls", None) or []) + content = strip_harmony_tokens(normalize_assistant_text(getattr(response, "content", "")) or "") + reasoning = normalize_assistant_text( + (getattr(response, "additional_kwargs", None) or {}).get("reasoning_content") + ) + if not isinstance(response, AIMessage): + response = AIMessage(content=content, tool_calls=tool_calls) + + if budget_exhausted: + # No tools were bound for the grace turn, so any tool_calls are noise. + tool_calls, invalid_tool_calls = [], [] + + if tool_calls or invalid_tool_calls: + try: + self._tracker.collect_step( + step=Step( + name="Assistant_tool_calls", + data=json.dumps(tool_calls, ensure_ascii=False, default=str), + ) + ) + except Exception as exc: + logger.debug(f"AgentGraphAdapter fc tracker error: {exc}") + else: + self.on_response_processed(state, None, content, reasoning) + + final_messages: list = list(modified_messages) + [response] + new_step_count: int = state.step_count + 1 + limit_cmd = ( + None + if budget_exhausted + else enforce_step_limit( + self, + state=state, + messages=final_messages, + new_step_count=new_step_count, + limit=self.resolve_max_steps(state, (configurable or {}).get("cuga_lite_max_steps")), + ) + ) + if limit_cmd is not None: + return limit_cmd + + meta_update = {self.metadata_key: self.build_metadata_update(state, playbook_fired=playbook_fired)} + + if tool_calls or invalid_tool_calls: + logger.info("[fc] {} native tool_call(s) -> tool_exec", len(tool_calls) + len(invalid_tool_calls)) + return Command( + goto="tool_exec", + update={ + self.messages_key: final_messages, + "script": None, + "step_count": new_step_count, + **meta_update, + }, + ) + + if "```" in content and not budget_exhausted: + # Mode violation: never execute code here. One corrective turn, charged + # as a step so it cannot loop past cuga_lite_max_steps. + meta = dict(self.get_metadata(state)) + violations = int(meta.get("fc_mode_violations", 0) or 0) + 1 + self.set_metadata(state, {**meta, "fc_mode_violations": violations}) + logger.warning("[fc] mode violation #{}: code block emitted in function-calling mode", violations) + return Command( + goto="call_model", + update={ + self.messages_key: final_messages + [HumanMessage(content=FC_MODE_VIOLATION_CORRECTION)], + "script": None, + "final_answer": "", + "execution_complete": False, + "step_count": new_step_count, + self.metadata_key: self.build_metadata_update(state, playbook_fired=playbook_fired), + }, + ) + + final_answer = content + if not final_answer.strip() and reasoning and not contains_harmony_tokens(reasoning): + final_answer = reasoning.strip() + if not final_answer.strip(): + for m in reversed(modified_messages): + if isinstance(m, ToolMessage) and isinstance(m.content, str) and m.content.strip(): + final_answer = m.content + break + if not content.strip() and final_answer: + final_messages = list(modified_messages) + [AIMessage(content=final_answer)] + + logger.info("[fc] no tool_calls -> final answer (END)") + return Command( + goto=END, + update={ + self.messages_key: final_messages, + "script": None, + "final_answer": final_answer, + "execution_complete": True, + "step_count": new_step_count, + **meta_update, + }, + ) + + def build_tool_exec_node(self): + return create_tool_exec_node(self) + def build_prepare_node(self, lc_bind_tools_meta: dict): return create_prepare_tools_and_apps_node(self, lc_bind_tools_meta) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index b05784d47..dda9efc69 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -33,7 +33,12 @@ _get_knowledge_tool_scope_context, _knowledge_scope_instruction, ) -from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import resolved_runtime_model_name +from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( + EXECUTION_MODE_FUNCTION_CALLING, + resolve_execution_mode, + resolve_fc_prompt_fragments, + resolved_runtime_model_name, +) from cuga.backend.cuga_graph.nodes.cuga_lite.providers.langchain import DirectLangChainToolsProvider from cuga.backend.cuga_graph.nodes.cuga_lite.providers.toolguard import ToolGuardingToolProvider from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( @@ -45,6 +50,7 @@ create_mcp_prompt, format_apps_for_prompt, normalize_mcp_few_shot_examples, + render_fc_prompt, resolve_cuga_lite_few_shots_enabled, ) from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.app_auth import call_authenticate_apps @@ -730,7 +736,25 @@ async def _wrapped(*args, **kwargs): # Create prompt dynamically dynamic_prompt = adapter._static_prompt - if not dynamic_prompt: + _execution_mode = resolve_execution_mode(configurable, _runtime_model_name) + if not dynamic_prompt and _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: + # Function-calling mode: the CodeAct prompt is never rendered. Tools + # travel natively via bind_tools, so this prompt is only the behavioural + # contract plus the same instructions / special_instructions. + _fragments = resolve_fc_prompt_fragments(configurable, _runtime_model_name) + dynamic_prompt = render_fc_prompt( + instructions=effective_instructions, + special_instructions=special_instructions_final, + is_autonomous_subtask=settings.advanced_features.force_autonomous_mode + or is_autonomous_subtask, + fragments=_fragments, + ) + logger.info( + "Prepared CugaLite function-calling prompt: fragments={} prompt_chars={}", + _fragments, + len(dynamic_prompt), + ) + elif not dynamic_prompt: dynamic_prompt = create_mcp_prompt( tools_for_prompt, allow_user_clarification=True, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py new file mode 100644 index 000000000..7c83e6dc9 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py @@ -0,0 +1,263 @@ +"""``tool_exec`` node — executes native ``tool_calls`` in function-calling mode. + +The counterpart of ``sandbox_node`` for ``cuga_lite_execution_mode = +"function_calling"``: where the sandbox runs a generated Python block, this +node runs the ``tool_calls`` the model attached to its last ``AIMessage`` and +replies to each with a ``ToolMessage`` carrying the call's exact +``tool_call_id``, so the conversation stays replayable to the provider. + +Guard envelope +-------------- +The budgets and the tracker are *not* on the callables in ``_tools_context`` — +they are applied by ``eval_with_tools_async`` and the sandbox node. A node that +merely called the callable would bypass the per-run / per-thread / per-block +tool-call caps and the run receipt, exactly the bug #560 shipped with on the +supervisor graph. So this node mirrors the sandbox node's envelope in the same +order: ``start_tracking`` -> ``seed_call_budget`` -> ``seed_block_budget`` -> +each call under ``counted_tool_call`` and a timeout -> ``stop_tracking`` -> +``_budget_updates()`` on every exit. + +Every outcome is a ``ToolMessage``, never a raised exception: unknown tool, +malformed arguments, a tool exception, a timeout, or a spent budget. One +failing call does not abort its siblings, and every id the model issued — +including provider-side ``invalid_tool_calls`` — receives a reply. +""" + +from __future__ import annotations + +import asyncio +import inspect +import json +from typing import Any, Callable, Dict, List, Optional + +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.runnables import RunnableConfig +from loguru import logger + +from cuga.backend.activity_tracker.tracker import Step +from cuga.backend.cuga_graph.nodes.cuga_agent_core.execution.code_extraction import make_tool_awaitable +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( + append_chat_messages_with_step_limit as core_append_with_step_limit, + create_error_command as core_create_error_command, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import _budget_updates +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( + ToolCallBudgetExceeded, + ToolCallTracker, + counted_tool_call, +) +from cuga.config import settings + +TRUNCATION_MARKER = "\n... [result truncated to {limit} characters]" + + +def _tool_call_parts(call: Any) -> tuple[str, str, Any]: + """``(id, name, args)`` from a LangChain tool_call dict or a raw OpenAI-shaped one.""" + if not isinstance(call, dict): + return "", "", None + fn = call.get("function") if isinstance(call.get("function"), dict) else {} + call_id = str(call.get("id") or "") + name = str(call.get("name") or fn.get("name") or "") + args = call.get("args") if "args" in call else fn.get("arguments") + return call_id, name, args + + +def _coerce_args(args: Any) -> tuple[Optional[Dict[str, Any]], Optional[str]]: + """Return ``(kwargs, error)``. Strings are JSON-decoded; anything else is rejected.""" + if args is None: + return {}, None + if isinstance(args, dict): + return args, None + if isinstance(args, str): + text = args.strip() + if not text: + return {}, None + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + return ( + None, + f"Could not parse tool arguments as JSON ({exc}). Re-issue the call with valid JSON arguments.", + ) + if isinstance(parsed, dict): + return parsed, None + return None, f"Tool arguments must be a JSON object, got {type(parsed).__name__}." + return None, f"Tool arguments must be an object, got {type(args).__name__}." + + +def _stringify(result: Any, limit: int) -> str: + if isinstance(result, str): + text = result + else: + try: + text = json.dumps(result, ensure_ascii=False, default=str) + except Exception: + text = str(result) + if limit and len(text) > limit: + return text[:limit] + TRUNCATION_MARKER.format(limit=limit) + return text + + +def _error_message(text: str, *, call_id: str, name: str) -> ToolMessage: + return ToolMessage(content=text, tool_call_id=call_id, name=name or "unknown", status="error") + + +def create_tool_exec_node(adapter: Any) -> Callable: + async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): + configurable = config.get("configurable", {}) if config else {} + track_tool_calls = configurable.get("track_tool_calls", False) + max_steps = configurable.get("cuga_lite_max_steps") if "cuga_lite_max_steps" in configurable else None + timeout = getattr(settings.advanced_features, "tool_call_timeout", 30) or None + output_limit = int(getattr(settings.advanced_features, "execution_output_max_length", 0) or 0) + + messages = adapter.get_messages(state) + last = messages[-1] if messages else None + calls: List[Any] = ( + list(getattr(last, "tool_calls", None) or []) if isinstance(last, AIMessage) else [] + ) + invalid: List[Any] = ( + list(getattr(last, "invalid_tool_calls", None) or []) if isinstance(last, AIMessage) else [] + ) + + # ── Guard envelope (mirrors sandbox_node, same order) ────────────── + ToolCallTracker.start_tracking( + enabled=bool(track_tool_calls), + timings_only=track_tool_calls == "timings_only", + ) + ToolCallTracker.seed_call_budget( + getattr(state, "tool_calls_used_run", 0), + getattr(state, "tool_calls_used_thread", 0), + ) + # One turn's batch of calls is the function-calling analogue of one code block. + ToolCallTracker.seed_block_budget() + + results: List[ToolMessage] = [] + executed = 0 + budget_stop: Optional[str] = None + try: + for index, call in enumerate(calls): + call_id, name, raw_args = _tool_call_parts(call) + call_id = call_id or f"call_{index}" + if not name: + results.append( + _error_message("Tool call carried no tool name.", call_id=call_id, name="") + ) + continue + if budget_stop: + results.append(_error_message(budget_stop, call_id=call_id, name=name)) + continue + + fn = adapter._tools_context.get(name) + if fn is None or not callable(fn): + known = ", ".join(sorted(k for k in adapter._tools_context if not k.startswith("_"))) + results.append( + _error_message( + f"Unknown tool '{name}'. Choose one of the provided tools: {known or '(none)'}.", + call_id=call_id, + name=name, + ) + ) + continue + + kwargs, arg_error = _coerce_args(raw_args) + if arg_error: + results.append(_error_message(arg_error, call_id=call_id, name=name)) + continue + + awaitable = fn if inspect.iscoroutinefunction(fn) else make_tool_awaitable(fn) + counted = counted_tool_call(awaitable) + try: + result = await asyncio.wait_for(counted(**kwargs), timeout=timeout) + executed += 1 + results.append( + ToolMessage(content=_stringify(result, output_limit), tool_call_id=call_id, name=name) + ) + except asyncio.TimeoutError: + results.append( + _error_message( + f"Tool '{name}' timed out after {timeout}s. Try a narrower call or a different tool.", + call_id=call_id, + name=name, + ) + ) + except ToolCallBudgetExceeded as exc: + # Every remaining id still gets a reply — without invoking anything. + if exc.scope != "block": + budget_stop = str(exc) + results.append(_error_message(str(exc), call_id=call_id, name=name)) + except TypeError as exc: + results.append( + _error_message( + f"Tool '{name}' rejected these arguments: {exc}. Check the parameter names and types.", + call_id=call_id, + name=name, + ) + ) + except Exception as exc: # one failing call must not abort its siblings + results.append(_error_message(f"Tool '{name}' failed: {exc}", call_id=call_id, name=name)) + + for index, bad in enumerate(invalid): + call_id, name, _ = _tool_call_parts(bad) + reason = (bad.get("error") if isinstance(bad, dict) else None) or "malformed tool call" + results.append( + _error_message( + f"The provider could not parse this tool call ({reason}). Re-issue it with valid arguments.", + call_id=call_id or f"invalid_{index}", + name=name, + ) + ) + finally: + execution_tool_calls = ToolCallTracker.stop_tracking() + + if not calls and not invalid: + # Nothing executable: advance the step and hand back to the model so a + # mis-route can never become an infinite no-op loop. + logger.warning("tool_exec entered with no tool_calls on the last message") + return {"step_count": state.step_count + 1, "script": None, **_budget_updates()} + + try: + adapter._tracker.collect_step( + step=Step( + name="Tool_results", + data=json.dumps( + [ + { + "tool_call_id": m.tool_call_id, + "name": m.name, + "status": m.status, + "content": m.content, + } + for m in results + ], + ensure_ascii=False, + default=str, + ), + ) + ) + except Exception as exc: + logger.debug(f"tool_exec tracker error: {exc}") + + updated_messages, error_message = core_append_with_step_limit(adapter, state, results, max_steps) + accumulated_tool_calls = (state.tool_calls or []) + (execution_tool_calls if track_tool_calls else []) + + if error_message: + return core_create_error_command( + adapter, + updated_messages, + error_message, + state.step_count, + additional_updates={"tool_calls": accumulated_tool_calls, **_budget_updates()}, + ) + + logger.info( + "[fc] tool_exec: {} call(s) -> {} ToolMessage(s), {} executed", len(calls), len(results), executed + ) + return { + adapter.messages_key: updated_messages, + "script": None, + "step_count": state.step_count + 1, + "tool_calls": accumulated_tool_calls, + **_budget_updates(), + } + + return tool_exec diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py index b855d6362..b3adcbe56 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py @@ -229,6 +229,10 @@ def create_cuga_lite_graph( prepare_node = adapter.build_prepare_node(lc_bind_tools_meta) sandbox_node = adapter.build_sandbox_node(thread_id, apps_list) call_model_node = _create_shared_call_model_node(adapter, model, settings) + # Native function-calling node. Always wired so the mode can be switched per + # invoke without recompiling; call_model only routes to it in that mode, so + # it is dormant on every CodeAct run. + tool_exec_node = adapter.build_tool_exec_node() return build_agent_graph( adapter=adapter, @@ -236,4 +240,5 @@ def create_cuga_lite_graph( prepare_node=prepare_node, call_model_node=call_model_node, execute_node=sandbox_node, + tool_exec_node=tool_exec_node, ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py index 0b161104d..ef6c98f51 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py @@ -114,3 +114,125 @@ def resolve_bind_tools_fields( inc = settings_include_fn() return mode_s, apps_list, tool_names_list, _bool_coerce(inc) + + +# ── Execution mode ─────────────────────────────────────────────────────────── +# +# Same three-layer resolution as the bind_tools keys above: ``configurable`` +# (per invoke) overrides the per-model runtime profile overrides ``settings``. +# Unknown values never raise — they fall back to the default with a warning, +# so a typo in an env var cannot take an agent down. + +EXECUTION_MODE_CODEACT = "codeact" +EXECUTION_MODE_FUNCTION_CALLING = "function_calling" + +_EXECUTION_MODE_ALIASES: Dict[str, str] = { + "codeact": EXECUTION_MODE_CODEACT, + "code_act": EXECUTION_MODE_CODEACT, + "code": EXECUTION_MODE_CODEACT, + "sandbox": EXECUTION_MODE_CODEACT, + "function_calling": EXECUTION_MODE_FUNCTION_CALLING, + "functioncalling": EXECUTION_MODE_FUNCTION_CALLING, + "function-calling": EXECUTION_MODE_FUNCTION_CALLING, + "fc": EXECUTION_MODE_FUNCTION_CALLING, + "native": EXECUTION_MODE_FUNCTION_CALLING, + "native_tool_calling": EXECUTION_MODE_FUNCTION_CALLING, + "tool_calling": EXECUTION_MODE_FUNCTION_CALLING, + "toolcalling": EXECUTION_MODE_FUNCTION_CALLING, +} + +FC_PROMPT_FRAGMENT_EVIDENCE_FIRST = "evidence_first" +KNOWN_FC_PROMPT_FRAGMENTS = frozenset({FC_PROMPT_FRAGMENT_EVIDENCE_FIRST}) + + +def _settings_value(key: str, default: Any) -> Any: + """Read ``settings.advanced_features.`` lazily; never raises.""" + try: + from cuga.config import settings + + return getattr(settings.advanced_features, key, default) + except Exception: + return default + + +def _layered( + key: str, + configurable: Optional[Dict[str, Any]], + model_name: Optional[str], + settings_fn: Optional[Callable[[], Any]], + default: Any, +) -> Any: + """configurable > per-model profile > settings > default, skipping empties.""" + cfg = configurable or {} + prof = runtime_defaults_for_model((model_name or "").strip()) + for candidate in ( + cfg.get(key), + prof.get(key), + settings_fn() if settings_fn else _settings_value(key, None), + ): + if candidate is None: + continue + if isinstance(candidate, str) and not candidate.strip(): + continue + if isinstance(candidate, (list, tuple)) and len(candidate) == 0: + continue + return candidate + return default + + +def normalize_execution_mode(raw: Any) -> str: + """Map any accepted spelling to a canonical mode; unknown -> codeact + warning.""" + if raw is None: + return EXECUTION_MODE_CODEACT + key = str(raw).strip().lower() + if key in _EXECUTION_MODE_ALIASES: + return _EXECUTION_MODE_ALIASES[key] + try: + from loguru import logger + + logger.warning( + "cuga_lite_execution_mode={!r} is not recognised; falling back to {!r}", + raw, + EXECUTION_MODE_CODEACT, + ) + except Exception: + pass + return EXECUTION_MODE_CODEACT + + +def resolve_execution_mode( + configurable: Optional[Dict[str, Any]], + model_name: Optional[str] = None, + *, + settings_mode_fn: Optional[Callable[[], Any]] = None, +) -> str: + """Resolved execution mode: ``codeact`` or ``function_calling``.""" + raw = _layered( + "cuga_lite_execution_mode", configurable, model_name, settings_mode_fn, EXECUTION_MODE_CODEACT + ) + return normalize_execution_mode(raw) + + +def resolve_fc_prompt_fragments( + configurable: Optional[Dict[str, Any]], + model_name: Optional[str] = None, + *, + settings_fn: Optional[Callable[[], Any]] = None, +) -> List[str]: + """Opt-in function-calling prompt fragments, unknown names dropped with a warning.""" + raw = _layered("cuga_lite_fc_prompt_fragments", configurable, model_name, settings_fn, []) + names = _normalize_bind_tools_string_list(raw) + out: List[str] = [] + for n in names: + key = n.strip().lower() + if key in KNOWN_FC_PROMPT_FRAGMENTS: + if key not in out: + out.append(key) + else: + try: + from loguru import logger + + logger.warning("Unknown cuga_lite_fc_prompt_fragments entry {!r} ignored", n) + except Exception: + pass + return out diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py index 3dfd13880..bf2d7e01b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py @@ -5,6 +5,7 @@ import functools import json import os +from functools import lru_cache from typing import Any, Dict, List, Optional from cuga.config import settings @@ -897,3 +898,49 @@ def create_mcp_prompt( } ).to_string() return prompt + + +# ── Function-calling prompt ─────────────────────────────────────────────────── + + +@lru_cache(maxsize=1) +def _fc_prompt_template(): + from pathlib import Path + + from cuga.backend.llm.utils.helpers import load_one_prompt + + path = Path(__file__).parent / "prompts" / "fc_prompt.jinja2" + return load_one_prompt(str(path), relative_to_caller=False) + + +def render_fc_prompt( + *, + instructions: Optional[str] = None, + special_instructions: Optional[str] = None, + base_prompt: Optional[str] = None, + is_autonomous_subtask: bool = False, + fragments: Optional[List[str]] = None, +) -> str: + """System prompt for ``cuga_lite_execution_mode = "function_calling"``. + + Deliberately short. Tool names, parameters and types travel in the API's + native ``tools`` parameter (``bind_tools``), which the model was trained to + read — so, unlike the CodeAct prompt, nothing about the tool catalog is + rendered as text. The prompt only sets the behavioural contract: + + - the generic preamble (always), + - opt-in fragments from ``cuga_lite_fc_prompt_fragments`` (``evidence_first``), + - then the agent's ``instructions`` and the caller's ``special_instructions``. + + Answer-format contracts belong in ``special_instructions``; no + benchmark- or deployment-specific wording ships in the template. + """ + frags = {f.strip().lower() for f in (fragments or [])} + rendered = _fc_prompt_template().format( + base_prompt=base_prompt or "", + is_autonomous_subtask=bool(is_autonomous_subtask), + evidence_first="evidence_first" in frags, + instructions=(instructions or "").strip(), + special_instructions=(special_instructions or "").strip(), + ) + return rendered.strip() + "\n" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 new file mode 100644 index 000000000..9e7bc41c8 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 @@ -0,0 +1,21 @@ +{% if base_prompt %}{{ base_prompt }} + +{% endif %}# ROLE + +You are Cuga Agent, a helpful assistant that answers the user's request by calling the provided tools. +{% if is_autonomous_subtask %} +**You are executing this task autonomously - work independently without asking the user any questions or requesting decisions. Complete the entire task before returning to the user.** +{% endif %} +# HOW TO ACT +- Use the available tools via native function-calling. Do NOT write Python code, code fences, pseudo-code, or JSON tool-call blocks in your text — issue real tool calls only. +{% if evidence_first %}- You MUST call at least one tool before giving a final answer, and base the answer STRICTLY on values returned by your tool calls — never answer from prior knowledge, memory, or assumption. If you have not yet observed the needed value in a tool result, make another tool call. +- Never guess an intermediate value; only use values returned by a previous tool call. +{% endif %}- If a tool returns an error, reason about why and try a corrected call. + +# FINAL ANSWER +- When your tool results contain everything needed, STOP calling tools and give the final answer with no further tool call. +{% if instructions %} +{{ instructions }} +{% endif %}{% if special_instructions %} +{{ special_instructions }} +{% endif %} \ No newline at end of file diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py new file mode 100644 index 000000000..fb4e6df73 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py @@ -0,0 +1,239 @@ +"""Configuration surfaces for ``cuga_lite_execution_mode``. + +One assertion per control plane — settings.toml, validators, ``configurable``, +the per-model runtime profile, the SDK constructor / ``invoke`` / ``stream`` — +so a knob wired into only some surfaces cannot ship (the usual "my config does +nothing" bug). Plus the graph wiring: the ``tool_exec`` node is present on the +CugaLite graph regardless of mode, and absent when a graph does not opt in. +""" + +from __future__ import annotations + +import inspect +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from cuga.backend.cuga_graph.nodes.cuga_lite import model_runtime_profile as mrp +from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( + EXECUTION_MODE_CODEACT, + EXECUTION_MODE_FUNCTION_CALLING, + normalize_execution_mode, + resolve_execution_mode, + resolve_fc_prompt_fragments, +) + +pytestmark = pytest.mark.unit + + +# ── 1. settings.toml + validators: the shipped default is the old behaviour ── + + +def test_settings_toml_defaults_to_codeact(): + from cuga.config import settings + + assert settings.advanced_features.cuga_lite_execution_mode == "codeact" + assert list(settings.advanced_features.cuga_lite_fc_prompt_fragments) == [] + + +def test_validators_supply_the_same_defaults_without_the_keys(): + """A settings.toml without the keys must resolve identically, not AttributeError.""" + from cuga.config import validators + + declared = {v.names[0]: v.default for v in validators if v.names} + assert declared["advanced_features.cuga_lite_execution_mode"] == "codeact" + assert declared["advanced_features.cuga_lite_fc_prompt_fragments"] == [] + + +def test_real_settings_resolve_to_the_defaults(): + assert resolve_execution_mode({}) == EXECUTION_MODE_CODEACT + assert resolve_fc_prompt_fragments({}) == [] + + +# ── 2. spellings: forgiving in, canonical out, never raising ───────────────── + + +@pytest.mark.parametrize( + "raw, expected", + [ + ("codeact", EXECUTION_MODE_CODEACT), + ("CodeAct", EXECUTION_MODE_CODEACT), + ("code_act", EXECUTION_MODE_CODEACT), + ("function_calling", EXECUTION_MODE_FUNCTION_CALLING), + (" Function-Calling ", EXECUTION_MODE_FUNCTION_CALLING), + ("fc", EXECUTION_MODE_FUNCTION_CALLING), + ("native", EXECUTION_MODE_FUNCTION_CALLING), + ("tool_calling", EXECUTION_MODE_FUNCTION_CALLING), + ], +) +def test_execution_mode_aliases(raw, expected): + assert normalize_execution_mode(raw) == expected + + +def test_unknown_execution_mode_falls_back_to_codeact_without_raising(): + assert normalize_execution_mode("banana") == EXECUTION_MODE_CODEACT + assert normalize_execution_mode(None) == EXECUTION_MODE_CODEACT + assert normalize_execution_mode(42) == EXECUTION_MODE_CODEACT + + +def test_fragments_are_normalised_deduped_and_unknown_dropped(): + cfg = {"cuga_lite_fc_prompt_fragments": ["evidence_first", "Evidence_First", "not_a_fragment"]} + assert resolve_fc_prompt_fragments(cfg) == ["evidence_first"] + assert resolve_fc_prompt_fragments({"cuga_lite_fc_prompt_fragments": "evidence_first"}) == [ + "evidence_first" + ] + + +# ── 3. precedence: configurable > per-model profile > settings ─────────────── + + +def test_configurable_beats_settings(): + assert ( + resolve_execution_mode({"cuga_lite_execution_mode": "fc"}, settings_mode_fn=lambda: "codeact") + == EXECUTION_MODE_FUNCTION_CALLING + ) + + +def test_settings_apply_when_configurable_is_silent_or_blank(): + assert ( + resolve_execution_mode({}, settings_mode_fn=lambda: "function_calling") + == EXECUTION_MODE_FUNCTION_CALLING + ) + # An empty string is "unset", not "codeact" — the next layer decides. + assert ( + resolve_execution_mode({"cuga_lite_execution_mode": ""}, settings_mode_fn=lambda: "function_calling") + == EXECUTION_MODE_FUNCTION_CALLING + ) + + +def test_model_profile_sits_between_configurable_and_settings(monkeypatch): + monkeypatch.setattr( + mrp, + "runtime_defaults_for_model", + lambda name: {"cuga_lite_execution_mode": "fc"} if name == "profiled-model" else {}, + ) + # profile beats settings + assert resolve_execution_mode({}, "profiled-model", settings_mode_fn=lambda: "codeact") == ( + EXECUTION_MODE_FUNCTION_CALLING + ) + # configurable beats profile + assert ( + resolve_execution_mode( + {"cuga_lite_execution_mode": "codeact"}, "profiled-model", settings_mode_fn=lambda: "fc" + ) + == EXECUTION_MODE_CODEACT + ) + # an unprofiled model falls through to settings + assert ( + resolve_execution_mode({}, "other-model", settings_mode_fn=lambda: "codeact") + == EXECUTION_MODE_CODEACT + ) + + +# ── 4. SDK: constructor default, per-call override, raw keys win ───────────── + + +def test_cuga_agent_accepts_execution_mode_everywhere(): + from cuga.sdk import CugaAgent + + for method in (CugaAgent.__init__, CugaAgent.invoke, CugaAgent.stream): + params = inspect.signature(method).parameters + assert "execution_mode" in params, method.__name__ + assert params["execution_mode"].default is None + + +def test_apply_execution_mode_writes_the_configurable_key(): + from cuga.sdk import CugaAgent + + agent = SimpleNamespace(_execution_mode="function_calling") + run_config = {"configurable": {}} + CugaAgent._apply_execution_mode(agent, run_config) + assert run_config["configurable"] == {"cuga_lite_execution_mode": "function_calling"} + + +def test_per_invoke_value_overrides_the_constructor_default(): + from cuga.sdk import CugaAgent + + agent = SimpleNamespace(_execution_mode="function_calling") + run_config = {"configurable": {}} + CugaAgent._apply_execution_mode(agent, run_config, execution_mode="codeact") + assert run_config["configurable"]["cuga_lite_execution_mode"] == "codeact" + + +def test_explicit_raw_keys_are_never_clobbered(): + from cuga.sdk import CugaAgent + + agent = SimpleNamespace(_execution_mode="function_calling") + run_config = {"configurable": {"cuga_lite_execution_mode": "codeact"}} + CugaAgent._apply_execution_mode(agent, run_config, execution_mode="fc") + assert run_config["configurable"]["cuga_lite_execution_mode"] == "codeact" + + +def test_apply_execution_mode_is_a_noop_when_unconfigured(): + from cuga.sdk import CugaAgent + + agent = SimpleNamespace(_execution_mode=None) + run_config = {"configurable": {}} + CugaAgent._apply_execution_mode(agent, run_config) + assert run_config["configurable"] == {}, "settings.toml must decide when nothing is passed" + + +# ── 5. graph wiring ────────────────────────────────────────────────────────── + + +def _dummy_nodes(): + async def prepare(state, config=None): + return {} + + async def call_model(state, config=None): + return {} + + async def execute(state, config=None): + return {} + + async def tool_exec(state, config=None): + return {} + + return prepare, call_model, execute, tool_exec + + +def test_build_agent_graph_adds_tool_exec_only_when_given(): + from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_graph import build_agent_graph + from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState + + prepare, call_model, execute, tool_exec = _dummy_nodes() + adapter = SimpleNamespace(execute_node_name="sandbox") + + without = build_agent_graph( + adapter=adapter, + state_class=CugaLiteState, + prepare_node=prepare, + call_model_node=call_model, + execute_node=execute, + ) + assert "tool_exec" not in without.nodes, "a graph that does not opt in must be unchanged" + + with_node = build_agent_graph( + adapter=adapter, + state_class=CugaLiteState, + prepare_node=prepare, + call_model_node=call_model, + execute_node=execute, + tool_exec_node=tool_exec, + ) + assert "tool_exec" in with_node.nodes + assert ("tool_exec", "call_model") in with_node.edges, "tool results always go back to the model" + + +def test_cuga_lite_graph_always_wires_tool_exec_so_mode_can_switch_per_invoke(): + from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import create_cuga_lite_graph + + provider = MagicMock() + provider.get_all_tools = AsyncMock(return_value=[]) + provider.get_apps = AsyncMock(return_value=[]) + provider.get_tools = AsyncMock(return_value=[]) + graph = create_cuga_lite_graph(model=MagicMock(), tool_provider=provider, apps_list=[]) + + assert {"prepare", "call_model", "sandbox", "tool_exec"} <= set(graph.nodes) + assert ("tool_exec", "call_model") in graph.edges diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py new file mode 100644 index 000000000..1fe916fff --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -0,0 +1,204 @@ +"""CugaLite's ``execute_call_model_fc``: the function-calling turn. + +Codeact: ``None``, nothing invoked. Function-calling: the bound model gets real +message objects (system prompt, few-shot demos, sanitized history), and the +response routes on ``tool_calls`` — to ``tool_exec`` with the assistant turn +kept verbatim, or to END as the final answer. A fenced code block with no +``tool_calls`` is a mode violation: never executed, one corrective turn. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, List +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langgraph.graph import END + +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import TOOL_BUDGET_EXHAUSTED_INSTRUCTION +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import graph_adapter as ga +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( + FC_MODE_VIOLATION_CORRECTION, + AgentGraphAdapter, + _sanitize_for_replay, +) + +pytestmark = pytest.mark.unit + +FC = {"cuga_lite_execution_mode": "function_calling"} + + +class _Model: + """Bare scripted model: no MagicMock attribute magic, so helper probes see a plain object.""" + + def __init__(self, response: Any): + self.response = response + self.seen: List[list] = [] + + async def ainvoke(self, messages, config=None, **kwargs): + self.seen.append(list(messages)) + return self.response + + +def _adapter() -> AgentGraphAdapter: + return AgentGraphAdapter( + tracker=MagicMock(), + base_callbacks=[], + task_todos_ref=[], + tools_context_ref={}, + base_tool_provider=None, + ) + + +def _state(messages=None, few_shots=None): + return SimpleNamespace( + chat_messages=messages or [HumanMessage(content="what is 1+2?")], + step_count=0, + cuga_lite_metadata={}, + mcp_few_shot_messages=few_shots or [], + cuga_lite_max_steps=None, + ) + + +async def _turn(adapter, model, state, configurable, *, budget_exhausted=False, system="SYS"): + return await adapter.execute_call_model_fc( + state=state, + config=None, + configurable=configurable, + active_model=model, + bound=model, + invoke_config={}, + system_content=system, + modified_messages=list(state.chat_messages), + budget_exhausted=budget_exhausted, + playbook_fired=False, + ) + + +_CALL = {"name": "add", "args": {"a": 1, "b": 2}, "id": "c1", "type": "tool_call"} + + +@pytest.mark.asyncio +async def test_codeact_returns_none_and_invokes_nothing(): + model = _Model(AIMessage(content="never")) + assert await _turn(_adapter(), model, _state(), {}) is None + assert await _turn(_adapter(), model, _state(), {"cuga_lite_execution_mode": "codeact"}) is None + assert model.seen == [] + + +@pytest.mark.asyncio +async def test_tool_calls_route_to_tool_exec_with_the_assistant_turn_verbatim(): + response = AIMessage(content="", tool_calls=[_CALL]) + model = _Model(response) + state = _state() + + cmd = await _turn(_adapter(), model, state, FC) + + assert cmd.goto == "tool_exec" + persisted = cmd.update["chat_messages"] + assert persisted[:-1] == state.chat_messages + assert persisted[-1] is response, "tool_calls must reach tool_exec untouched, ids included" + assert cmd.update["script"] is None and cmd.update["step_count"] == 1 + + (outbound,) = model.seen + assert isinstance(outbound[0], SystemMessage) and outbound[0].content == "SYS" + assert outbound[1:] == state.chat_messages, "history goes out as message objects, not the CodeAct dicts" + + +@pytest.mark.asyncio +async def test_plain_text_is_the_final_answer(): + cmd = await _turn(_adapter(), _Model(AIMessage(content="It is 3.")), _state(), FC) + + assert cmd.goto == END + assert cmd.update["final_answer"] == "It is 3." and cmd.update["execution_complete"] is True + assert isinstance(cmd.update["chat_messages"][-1], AIMessage) + + +@pytest.mark.asyncio +async def test_empty_answer_falls_back_to_the_last_tool_result(): + history = [ + HumanMessage(content="q"), + AIMessage(content="", tool_calls=[_CALL]), + ToolMessage(content="3", tool_call_id="c1", name="add"), + ] + cmd = await _turn(_adapter(), _Model(AIMessage(content="")), _state(history), FC) + + assert cmd.goto == END and cmd.update["final_answer"] == "3" + + +@pytest.mark.asyncio +async def test_code_block_without_tool_calls_is_a_mode_violation_never_executed(): + state = _state() + cmd = await _turn(_adapter(), _Model(AIMessage(content="```python\nawait add(a=1, b=2)\n```")), state, FC) + + assert cmd.goto == "call_model", "one corrective turn, not the sandbox" + assert cmd.update["script"] is None + assert cmd.update["chat_messages"][-1] == HumanMessage(content=FC_MODE_VIOLATION_CORRECTION) + assert cmd.update["cuga_lite_metadata"]["fc_mode_violations"] == 1 + assert cmd.update["step_count"] == 1, "charged as a step so it cannot loop forever" + + +@pytest.mark.asyncio +async def test_budget_exhausted_turn_ignores_tool_calls_and_sends_the_instruction_outbound_only(): + model = _Model(AIMessage(content="From what I have: 3.", tool_calls=[_CALL])) + state = _state() + + cmd = await _turn(_adapter(), model, state, FC, budget_exhausted=True) + + assert cmd.goto == END and cmd.update["final_answer"] == "From what I have: 3." + (outbound,) = model.seen + assert outbound[-1] == HumanMessage(content=TOOL_BUDGET_EXHAUSTED_INSTRUCTION) + assert all(m.content != TOOL_BUDGET_EXHAUSTED_INSTRUCTION for m in cmd.update["chat_messages"]), ( + "never persisted" + ) + + +@pytest.mark.asyncio +async def test_few_shot_demos_go_out_as_chat_messages_before_history(): + model = _Model(AIMessage(content="ok")) + few = [{"role": "user", "content": "demo question"}, {"role": "assistant", "content": "demo answer"}] + + await _turn(_adapter(), model, _state(few_shots=few), FC) + + (outbound,) = model.seen + assert outbound[1] == HumanMessage(content="demo question") + assert outbound[2] == AIMessage(content="demo answer") + assert outbound[3].content == "what is 1+2?" + + +def test_replay_sanitizer_drops_reasoning_and_strips_harmony_without_touching_history(monkeypatch): + monkeypatch.setattr(ga, "strip_harmony_tokens", lambda s: s.replace("<|channel|>final", "")) + kept = HumanMessage(content="plain") + noisy = AIMessage( + content="<|channel|>finalanswer", additional_kwargs={"reasoning_content": "thinking...", "x": 1} + ) + + out = _sanitize_for_replay([kept, noisy]) + + assert out[0] is kept, "untouched messages are passed through by identity" + assert out[1].content == "answer" and out[1].additional_kwargs == {"x": 1} + assert noisy.additional_kwargs["reasoning_content"] == "thinking...", "persisted history is never mutated" + + +@pytest.mark.asyncio +async def test_bind_mode_is_upgraded_from_none_only_in_function_calling_mode(): + captured = [] + + async def fake_resolve(active_model, **kwargs): + captured.append(kwargs["configurable"]) + return "BOUND" + + with patch.object(ga, "resolve_model_with_bind_tools", new=AsyncMock(side_effect=fake_resolve)): + adapter = _adapter() + assert await adapter.resolve_bind_tools(_state(), object(), dict(FC), None) == "BOUND" + await adapter.resolve_bind_tools( + _state(), object(), {**FC, "cuga_lite_bind_tools_mode": "find_tools"}, None + ) + await adapter.resolve_bind_tools(_state(), object(), {}, None) + + fc_default, fc_explicit, codeact = captured + assert fc_default["cuga_lite_bind_tools_mode"] == "all", "FC is inert without advertised tools" + assert fc_explicit["cuga_lite_bind_tools_mode"] == "find_tools", "an explicit choice is respected" + assert "cuga_lite_bind_tools_mode" not in codeact, "codeact is byte-identical" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py new file mode 100644 index 000000000..6b03d05df --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -0,0 +1,236 @@ +"""Function-calling mode end to end: the real CugaLite graph, a scripted model. + +Node tests pin what each node returns; only a real run proves the loop — +prepare selects the FC prompt, call_model routes native ``tool_calls`` to +``tool_exec``, the ``ToolMessage`` replies are replayed to the model with their +ids, and a text reply ends the run. The scripted model asserts on what it is +sent, so a broken replay fails loudly instead of the test passing by accident. + +The last test is the feature-off golden: with no mode configured the graph +behaves exactly as before — dict-serialised CodeAct turns, the sandbox, and a +``tool_exec`` node that is present but never entered. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import StructuredTool +from langgraph.checkpoint.memory import MemorySaver + +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( + FC_MODE_VIOLATION_CORRECTION, + AgentGraphAdapter, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState, create_cuga_lite_graph +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module + +pytestmark = pytest.mark.unit + +CALLS: list = [] + + +class _ScriptedModel: + """Queued responses; an entry may be a callable that receives the outbound messages.""" + + def __init__(self, responses): + self._responses = list(responses) + self.invocations = 0 + self.seen: list = [] + self.bound_tool_names: list = [] + + def bind_tools(self, tools, **kwargs): + self.bound_tool_names = [getattr(t, "name", str(t)) for t in tools] + return self + + async def ainvoke(self, messages, config=None, **kwargs): + self.invocations += 1 + self.seen.append(list(messages)) + if not self._responses: + raise AssertionError( + f"model asked for response #{self.invocations} — the run did not end when it should" + ) + nxt = self._responses.pop(0) + return nxt(messages) if callable(nxt) else nxt + + +def _provider(*tools): + provider = MagicMock() + provider.get_all_tools = AsyncMock(return_value=list(tools)) + provider.get_apps = AsyncMock(return_value=[]) + provider.get_tools = AsyncMock(return_value=[]) + provider.app_name = "test_app" + return provider + + +def _echo_tool(): + async def echo(value: int) -> int: + """Echo a value.""" + CALLS.append(("echo", value)) + return value + + return StructuredTool.from_function(coroutine=echo, name="echo", description="Echo a value.") + + +def _tc(name, args, call_id): + return {"name": name, "args": args, "id": call_id, "type": "tool_call"} + + +def _config(thread_id, **extra): + return {"configurable": {"thread_id": thread_id, "enable_todos": False, **extra}} + + +def _last_tool_message(messages): + assert isinstance(messages[-1], ToolMessage), ( + f"expected a ToolMessage last, got {type(messages[-1]).__name__}" + ) + return messages[-1] + + +@pytest.fixture(autouse=True) +def _quiet(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + CALLS.clear() + yield + CALLS.clear() + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + + +def _run(model, config, *tools): + graph = create_cuga_lite_graph( + model=model, tool_provider=_provider(*tools), apps_list=[], thread_id="t" + ).compile(checkpointer=MemorySaver()) + return graph.ainvoke(CugaLiteState(chat_messages=[HumanMessage(content="use the tools")]), config=config) + + +@pytest.mark.asyncio +async def test_single_hop_tool_call_round_trip(): + model = _ScriptedModel( + [ + AIMessage(content="", tool_calls=[_tc("echo", {"value": 7}, "call_1")]), + AIMessage(content="The value is 7."), + ] + ) + + result = await _run(model, _config("fc-1", cuga_lite_execution_mode="function_calling"), _echo_tool()) + + assert CALLS == [("echo", 7)] + assert result["final_answer"] == "The value is 7." + assert result["script"] is None and model.invocations == 2 + assert "echo" in model.bound_tool_names, ( + "FC mode must advertise the tools natively (bind mode none -> all)" + ) + + kinds = [type(m).__name__ for m in result["chat_messages"]] + assert kinds == ["HumanMessage", "AIMessage", "ToolMessage", "AIMessage"], kinds + assert result["chat_messages"][2].tool_call_id == "call_1" and result["chat_messages"][2].content == "7" + + first, second = model.seen + assert isinstance(first[0], SystemMessage) and "native function-calling" in first[0].content + assert _last_tool_message(second).tool_call_id == "call_1", "the ToolMessage must be replayed with its id" + assert "native function-calling" in result["prepared_prompt"] + assert "```python" not in result["prepared_prompt"], "the CodeAct prompt must not be rendered in FC mode" + + +@pytest.mark.asyncio +async def test_three_hop_chain_each_hop_depends_on_the_previous_result(): + async def lookup_id(name: str) -> int: + """Employee id by name.""" + CALLS.append(("lookup_id", name)) + return 42 + + async def lookup_manager(employee_id: int) -> str: + """Manager id for an employee.""" + CALLS.append(("lookup_manager", employee_id)) + return "M-42" + + async def lookup_email(manager: str) -> str: + """Email for a manager id.""" + CALLS.append(("lookup_email", manager)) + return "m42@example.com" + + tools = [ + StructuredTool.from_function(coroutine=lookup_id, name="lookup_id", description="id by name"), + StructuredTool.from_function(coroutine=lookup_manager, name="lookup_manager", description="manager"), + StructuredTool.from_function(coroutine=lookup_email, name="lookup_email", description="email"), + ] + + def hop2(messages): + assert _last_tool_message(messages).content == "42" + return AIMessage(content="", tool_calls=[_tc("lookup_manager", {"employee_id": 42}, "c2")]) + + def hop3(messages): + assert _last_tool_message(messages).content == "M-42" + return AIMessage(content="", tool_calls=[_tc("lookup_email", {"manager": "M-42"}, "c3")]) + + def final(messages): + assert _last_tool_message(messages).content == "m42@example.com" + return AIMessage(content="Alice's manager can be reached at m42@example.com.") + + model = _ScriptedModel( + [AIMessage(content="", tool_calls=[_tc("lookup_id", {"name": "alice"}, "c1")]), hop2, hop3, final] + ) + + result = await _run(model, _config("fc-3", cuga_lite_execution_mode="function_calling"), *tools) + + assert CALLS == [("lookup_id", "alice"), ("lookup_manager", 42), ("lookup_email", "M-42")] + assert result["final_answer"].endswith("m42@example.com.") + assert [m.tool_call_id for m in result["chat_messages"] if isinstance(m, ToolMessage)] == [ + "c1", + "c2", + "c3", + ] + assert model.invocations == 4 + + +@pytest.mark.asyncio +async def test_code_block_in_fc_mode_is_corrected_not_executed(): + model = _ScriptedModel( + [AIMessage(content="```python\nawait echo(value=9)\n```"), AIMessage(content="Understood: done.")] + ) + + result = await _run(model, _config("fc-v", cuga_lite_execution_mode="function_calling"), _echo_tool()) + + assert CALLS == [], "code must never run in function-calling mode" + assert any( + isinstance(m, HumanMessage) and m.content == FC_MODE_VIOLATION_CORRECTION + for m in result["chat_messages"] + ) + assert result["final_answer"] == "Understood: done." and model.invocations == 2 + + +@pytest.mark.asyncio +async def test_feature_off_is_the_old_codeact_run_and_tool_exec_is_never_entered(monkeypatch): + async def forbidden(state, config=None): + raise AssertionError("tool_exec entered on a CodeAct run") + + monkeypatch.setattr(AgentGraphAdapter, "build_tool_exec_node", lambda self: forbidden) + model = _ScriptedModel( + [ + AIMessage(content="```python\nr = await echo(value=3)\nprint(r)\n```"), + AIMessage(content="It printed 3."), + ] + ) + + graph = create_cuga_lite_graph( + model=model, tool_provider=_provider(_echo_tool()), apps_list=[], thread_id="t" + ) + assert "tool_exec" in graph.nodes, "the node is wired (dormant) so the mode can switch per invoke" + result = await graph.compile(checkpointer=MemorySaver()).ainvoke( + CugaLiteState(chat_messages=[HumanMessage(content="use the tools")]), config=_config("codeact-golden") + ) + + assert CALLS == [("echo", 3)] + assert result["final_answer"] == "It printed 3." + kinds = [type(m).__name__ for m in result["chat_messages"]] + assert kinds == ["HumanMessage", "AIMessage", "HumanMessage", "AIMessage"], kinds + assert all(isinstance(m, dict) for m in model.seen[0]), ( + "CodeAct still serialises turns through the dict path" + ) + assert "```python" in result["prepared_prompt"], "the CodeAct prompt is unchanged" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py new file mode 100644 index 000000000..34103cd9e --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py @@ -0,0 +1,53 @@ +"""The function-calling system prompt: short, generic, assembled from configuration.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from cuga.backend.cuga_graph.nodes.cuga_lite.prompt_utils import render_fc_prompt + +pytestmark = pytest.mark.unit + +_TEMPLATE = Path(__file__).resolve().parents[1] / "prompts" / "fc_prompt.jinja2" +_EVIDENCE_LINE = "You MUST call at least one tool before giving a final answer" + + +def test_default_prompt_is_the_bare_contract(): + text = render_fc_prompt() + assert text.startswith("# ROLE") + assert "native function-calling" in text and "# FINAL ANSWER" in text + assert "Do NOT write Python code" in text + assert _EVIDENCE_LINE not in text + assert "autonomously" not in text + + +def test_fragments_are_opt_in_and_unknown_names_are_ignored(): + assert _EVIDENCE_LINE in render_fc_prompt(fragments=["evidence_first"]) + assert _EVIDENCE_LINE in render_fc_prompt(fragments=["Evidence_First"]) + assert _EVIDENCE_LINE not in render_fc_prompt(fragments=["something_else"]) + + +def test_autonomous_subtask_line(): + assert "autonomously" in render_fc_prompt(is_autonomous_subtask=True) + + +def test_instructions_then_special_instructions_close_the_prompt(): + text = render_fc_prompt( + fragments=["evidence_first"], + instructions=" Agent instructions. ", + special_instructions="Answer format: one line.", + ) + body = text.rstrip("\n") + assert body.endswith("Agent instructions.\n\nAnswer format: one line.") + assert text.index(_EVIDENCE_LINE) < text.index("# FINAL ANSWER") + assert "\n\n\n" not in text, "no blank-line runs from unset optional sections" + + +def test_template_ships_no_benchmark_or_deployment_wording(): + text = _TEMPLATE.read_text() + assert "[[" not in text + assert "vakra" not in text.lower() + for module in ("adapter/tool_exec_node.py", "adapter/graph_adapter.py", "prompt_utils.py"): + assert "vakra" not in (_TEMPLATE.parents[1] / module).read_text().lower(), module diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py new file mode 100644 index 000000000..90649ce3d --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py @@ -0,0 +1,323 @@ +"""``tool_exec`` node: every ``tool_call`` gets a ``ToolMessage``, inside the guard envelope. + +The node is the function-calling counterpart of the sandbox, so the assertions +mirror what #560 pinned for the sandbox: budgets and the tracker apply to every +call even though the callables in ``_tools_context`` are bare, and every exit +carries the budget fields. On top of that, the FC-specific contract: ids are +preserved, one failing call never aborts its siblings, provider-side invalid +calls are answered. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any, List, Optional +from unittest.mock import MagicMock + +import pytest +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage +from langgraph.types import Command + +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import ( + create_tool_exec_node, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( + ToolCallTracker, + make_recording_awaitable, +) + +pytestmark = pytest.mark.unit + + +class _Adapter: + messages_key = "chat_messages" + metadata_key = "cuga_lite_metadata" + execute_node_name = "sandbox" + sender_name = "CugaLite" + + def __init__(self, tools: dict, max_steps: int = 50): + self._tools_context = tools + self._tracker = MagicMock() + self._model = None + self._max_steps = max_steps + + def get_messages(self, state: Any) -> List[BaseMessage]: + return list(state.chat_messages or []) + + def resolve_max_steps(self, state: Any, override: Optional[int]) -> int: + return override if override is not None else self._max_steps + + +def _state(last: AIMessage, step_count: int = 0): + return SimpleNamespace( + chat_messages=[HumanMessage(content="add 1 and 2"), last], + step_count=step_count, + tool_calls=[], + tool_calls_used_run=0, + tool_calls_used_thread=0, + ) + + +def _call(name: str, args: dict, call_id: str) -> dict: + return {"name": name, "args": args, "id": call_id, "type": "tool_call"} + + +def _tool_messages(result) -> List[ToolMessage]: + return [m for m in result["chat_messages"] if isinstance(m, ToolMessage)] + + +@pytest.fixture(autouse=True) +def _reset_budget_contexts(): + yield + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + + +def _caps(monkeypatch, *, block=0, run=0, thread=0): + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_block", block, raising=False) + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_run", run, raising=False) + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_thread", thread, raising=False) + + +# ── happy path ─────────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_every_call_runs_and_ids_are_preserved(monkeypatch): + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + def greet(name: str) -> str: # sync tools are made awaitable, like the sandbox does + return f"hi {name}" + + node = create_tool_exec_node(_Adapter({"add": add, "greet": greet})) + last = AIMessage( + content="", tool_calls=[_call("add", {"a": 1, "b": 2}, "c1"), _call("greet", {"name": "bob"}, "c2")] + ) + + result = await node(_state(last), config={"configurable": {}}) + + msgs = _tool_messages(result) + assert [(m.tool_call_id, m.name, m.content, m.status) for m in msgs] == [ + ("c1", "add", "3", "success"), + ("c2", "greet", "hi bob", "success"), + ] + assert result["chat_messages"][:2] == _state(last).chat_messages, ( + "history is appended to, never rewritten" + ) + assert result["script"] is None + assert result["step_count"] == 1 + assert result["tool_calls_used_run"] == 2 and result["tool_budget_exhausted"] is False + + +@pytest.mark.asyncio +async def test_raw_provider_shaped_calls_with_json_string_args_are_decoded(monkeypatch): + """LangChain parses ``args`` into a dict before an AIMessage is built, so a + JSON-string payload only arrives in the raw provider shape (``function.arguments``). + ``model_construct`` bypasses validation to put exactly that on the message.""" + _caps(monkeypatch) + seen = [] + + async def add(a: int, b: int) -> int: + seen.append((a, b)) + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage.model_construct( + content="", + tool_calls=[ + {"id": "c1", "function": {"name": "add", "arguments": '{"a": 5, "b": 6}'}}, + {"id": "c2", "function": {"name": "add", "arguments": "{not json"}}, + ], + invalid_tool_calls=[], + additional_kwargs={}, + response_metadata={}, + ) + + result = await node(_state(last), config={"configurable": {}}) + + m1, m2 = _tool_messages(result) + assert (m1.tool_call_id, m1.content, m1.status) == ("c1", "11", "success") and seen == [(5, 6)] + assert m2.tool_call_id == "c2" and m2.status == "error" and "valid JSON" in m2.content + + +# ── every failure is a ToolMessage, never an exception ─────────────────────── + + +@pytest.mark.asyncio +async def test_unknown_tool_lists_the_known_names(monkeypatch): + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"add": add, "_internal": add})) + result = await node(_state(AIMessage(content="", tool_calls=[_call("subtract", {}, "c1")])), config=None) + + (m,) = _tool_messages(result) + assert m.status == "error" and m.tool_call_id == "c1" + assert "Unknown tool 'subtract'" in m.content and "add" in m.content + assert "_internal" not in m.content, "internal injections are not offered to the model" + + +@pytest.mark.asyncio +async def test_a_failing_call_does_not_abort_its_siblings(monkeypatch): + _caps(monkeypatch) + + async def boom(x: int) -> int: + raise ValueError("nope") + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"boom": boom, "add": add})) + last = AIMessage( + content="", + tool_calls=[ + _call("boom", {"x": 1}, "c1"), + _call("add", {"wrong": 1}, "c2"), + _call("add", {"a": 2, "b": 2}, "c3"), + ], + ) + + result = await node(_state(last), config=None) + + m1, m2, m3 = _tool_messages(result) + assert m1.status == "error" and "nope" in m1.content + assert m2.status == "error" and "rejected these arguments" in m2.content + assert (m3.content, m3.status) == ("4", "success") + + +@pytest.mark.asyncio +async def test_timeout_is_an_error_toolmessage(monkeypatch): + _caps(monkeypatch) + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "tool_call_timeout", 0.05, raising=False) + + async def slow() -> str: + await asyncio.Event().wait() # never set: immune to any asyncio.sleep patch elsewhere in the suite + return "late" + + node = create_tool_exec_node(_Adapter({"slow": slow})) + result = await node(_state(AIMessage(content="", tool_calls=[_call("slow", {}, "c1")])), config=None) + + (m,) = _tool_messages(result) + assert m.status == "error" and "timed out" in m.content and m.tool_call_id == "c1" + + +@pytest.mark.asyncio +async def test_provider_side_invalid_tool_calls_are_answered(monkeypatch): + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage( + content="", + tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")], + invalid_tool_calls=[ + { + "name": "add", + "args": "{broken", + "id": "bad1", + "error": "unterminated", + "type": "invalid_tool_call", + } + ], + ) + + result = await node(_state(last), config=None) + + good, bad = _tool_messages(result) + assert (good.tool_call_id, good.content) == ("c1", "2") + assert bad.tool_call_id == "bad1" and bad.status == "error" and "unterminated" in bad.content + + +# ── guard envelope ─────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_bare_callables_are_still_budget_capped(monkeypatch): + """The callables in _tools_context carry no budget wrapper — the node must add it.""" + _caps(monkeypatch, run=1) + calls = [] + + async def add(a: int, b: int) -> int: + calls.append((a, b)) + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage( + content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] + ) + + result = await node(_state(last), config=None) + + m1, m2 = _tool_messages(result) + assert m1.content == "2" and calls == [(1, 1)], "the second call must be refused, not executed" + assert m2.status == "error" and "Tool call limit reached" in m2.content + assert result["tool_calls_used_run"] == 1, "a refused call never inflates the counter" + assert result["tool_budget_exhausted"] is True, "call_model must see the exhaustion and end the turn" + + +@pytest.mark.asyncio +async def test_tracking_session_wraps_the_calls(monkeypatch): + """A recorded tool lands in result['tool_calls'] when tracking is requested, and + the session is closed afterwards — the receipt cannot leak into the next node.""" + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"add": make_recording_awaitable(add, "add", app_name="calc")})) + last = AIMessage(content="", tool_calls=[_call("add", {"a": 3, "b": 4}, "c1")]) + + result = await node(_state(last), config={"configurable": {"track_tool_calls": True}}) + + assert [(c["name"], c["app_name"], c["result"]) for c in result["tool_calls"]] == [("add", "calc", 7)] + assert ToolCallTracker.is_enabled() is False + + untracked = await node(_state(last), config={"configurable": {}}) + assert untracked["tool_calls"] == [] + + +# ── routing edge cases ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_nothing_executable_still_advances_the_step(monkeypatch): + """A mis-route into tool_exec must not become an infinite no-op loop.""" + _caps(monkeypatch) + node = create_tool_exec_node(_Adapter({})) + + result = await node(_state(AIMessage(content="no calls here"), step_count=4), config=None) + + assert result["step_count"] == 5 and result["script"] is None + assert "chat_messages" not in result, "no ToolMessage is fabricated" + + +@pytest.mark.asyncio +async def test_step_limit_ends_the_run_with_an_error_command(monkeypatch): + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"add": add}, max_steps=3)) + last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) + + result = await node(_state(last, step_count=3), config=None) + + assert isinstance(result, Command) + assert "Maximum step limit" in result.update["error"] + assert result.update["execution_complete"] is True + assert "tool_calls_used_run" in result.update, "budget fields ride every exit" diff --git a/src/cuga/config.py b/src/cuga/config.py index 9dfd57458..d2cf46b53 100644 --- a/src/cuga/config.py +++ b/src/cuga/config.py @@ -220,6 +220,9 @@ def get_all_paths(config, parent_key=""): Validator("advanced_features.cuga_lite_bind_tools_tool_names", default=[]), Validator("advanced_features.cuga_lite_bind_tools_max_count", default=128), Validator("advanced_features.cuga_lite_bind_tools_pad_to_cap", default=False), + # Native function-calling execution mode (default off). + Validator("advanced_features.cuga_lite_execution_mode", default="codeact"), + Validator("advanced_features.cuga_lite_fc_prompt_fragments", default=[]), # Read at prepare_node without a getattr guard; without this a settings.toml # missing the key raises AttributeError mid-run instead of using the default. Validator("advanced_features.shortlisting_tool_threshold", default=35), diff --git a/src/cuga/sdk.py b/src/cuga/sdk.py index 8407ea164..c43be34aa 100644 --- a/src/cuga/sdk.py +++ b/src/cuga/sdk.py @@ -1855,6 +1855,7 @@ def __init__( skills_folder: Optional[str] = None, shortlister: Optional["Shortlister"] = None, final_answer: "Optional[Union[str, Callable[[str], str], FinalAnswerConfig]]" = None, + execution_mode: Optional[str] = None, ): """ Initialize the CUGA Agent. @@ -1877,6 +1878,10 @@ def __init__( shortlister: How to shrink a large tool set before the model sees it, e.g. `Shortlister(strategy="hybrid")`. None = use `[shortlister]` settings (default `"llm"`, i.e. unchanged behavior). Overridable per invoke()/stream(). + execution_mode: How the model drives tools — `"codeact"` (default: the model writes + Python that calls tools in the sandbox) or `"function_calling"` (the model emits + native tool calls, executed one round-trip at a time). None = use + `advanced_features.cuga_lite_execution_mode`. Overridable per invoke()/stream(). final_answer: How the final answer is shaped. A `str` adds guidance for the answer-composing LLM; a callable `(str) -> str` deterministically post-processes the composed answer before delivery (must be pure; @@ -1984,6 +1989,10 @@ def __init__( # is the pre-feature behavior. self._shortlister = shortlister + # Execution mode. None => settings.toml (default "codeact", i.e. the + # pre-feature behavior). + self._execution_mode = execution_mode + # Setup tool provider. ToolGuard is installed immediately as a transparent # provider-level decorator so create-agent-first, add-guard-later flows work. from cuga.backend.cuga_graph.nodes.cuga_lite.providers.langchain import DirectLangChainToolsProvider @@ -2125,6 +2134,23 @@ def _apply_shortlister(self, run_config: dict, shortlister: Optional["Shortliste except Exception as e: logger.warning(f"Applying Shortlister config failed; using default shortlisting: {e}") + def _apply_execution_mode( + self, + run_config: dict, + execution_mode: Optional[str] = None, + ) -> None: + """Merge the execution mode into ``run_config['configurable']``. + + A per-invoke value overrides the constructor default. A raw + ``cuga_lite_execution_mode`` key already set by the caller wins over + both — we ``setdefault`` rather than overwrite. + No-op when nothing is configured, so settings.toml decides. + """ + configurable = run_config["configurable"] + mode = execution_mode if execution_mode is not None else self._execution_mode + if mode is not None: + configurable.setdefault("cuga_lite_execution_mode", mode) + def _apply_callbacks( self, run_config: dict, extra_callbacks: Optional[List[BaseCallbackHandler]] = None ) -> None: @@ -2583,6 +2609,7 @@ async def invoke( track_tool_calls: bool = False, variables: Optional[Dict[str, Any]] = None, shortlister: Optional["Shortlister"] = None, + execution_mode: Optional[str] = None, ) -> InvokeResult: """ Invoke the agent with a message and get the response. @@ -2601,6 +2628,8 @@ async def invoke( Used when delegating from a supervisor to pass context to sub-agents. shortlister: Overrides the agent's shortlister for this call only, e.g. `Shortlister(strategy="hybrid")`. None = use the agent default. + execution_mode: `"codeact"` or `"function_calling"` for this call only. + None = use the agent default. Returns: InvokeResult containing: @@ -2676,6 +2705,7 @@ async def invoke( # Setup config (shallow-copied so we don't mutate the caller's dict) run_config = self._prepare_run_config(config) self._apply_shortlister(run_config, shortlister) + self._apply_execution_mode(run_config, execution_mode) # Pass track_tool_calls flag via configurable run_config["configurable"]["track_tool_calls"] = track_tool_calls @@ -3105,6 +3135,7 @@ async def stream( config: Optional[Dict[str, Any]] = None, action_response: Optional[Any] = None, # ActionResponse for resuming after HITL shortlister: Optional["Shortlister"] = None, + execution_mode: Optional[str] = None, ): """ Stream the agent's execution step by step. @@ -3120,6 +3151,8 @@ async def stream( action_response: Optional ActionResponse for resuming after approval/interruption shortlister: Overrides the agent's shortlister for this call only, e.g. `Shortlister(strategy="hybrid")`. None = use the agent default. + execution_mode: `"codeact"` or `"function_calling"` for this call only. + None = use the agent default. Yields: State updates as the agent executes @@ -3152,6 +3185,7 @@ async def stream( # Setup config (shallow-copied so we don't mutate the caller's dict) run_config = self._prepare_run_config(config) self._apply_shortlister(run_config, shortlister) + self._apply_execution_mode(run_config, execution_mode) # Pass skills configuration via configurable (overrides settings when set) if self._enable_skills is not None: diff --git a/src/cuga/settings.toml b/src/cuga/settings.toml index 912d5ec13..0f3e1326d 100644 --- a/src/cuga/settings.toml +++ b/src/cuga/settings.toml @@ -54,6 +54,13 @@ cuga_lite_bind_tools_include_find_tools = false # Also bind find_tools alongsid cuga_lite_bind_tools_max_count = 128 # Pad shortlister output to fill the cap. Off by default for code-execution mode. cuga_lite_bind_tools_pad_to_cap = false +# Execution mode: codeact (write Python, run it in the sandbox — default) | function_calling (the model +# issues native tool_calls, a tool_exec node runs them and replies with ToolMessages). Aliases: fc, native. +# Per-invoke: config["configurable"]["cuga_lite_execution_mode"]. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_EXECUTION_MODE +cuga_lite_execution_mode = "codeact" +# Opt-in prompt fragments for function_calling mode. Known: "evidence_first" (call a tool before answering; never +# answer from prior knowledge). Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_FC_PROMPT_FRAGMENTS='["evidence_first"]' +cuga_lite_fc_prompt_fragments = [] cuga_lite_nl_auto_continue = false # When model returns NL with no code, LLM-classify interim vs final; if interim, simulate user "continue" and re-call model cuga_lite_blocked_claim_retry = true # When an NL turn claims tools/data are unavailable but no code has executed and tools ARE bound, retry once with a corrective message instead of finalizing the false refusal (issue #610). Only active when cuga_lite_nl_auto_continue is on. path_segment_index = 1 # Which path segment to use for operation naming (1 = first, 2 = second, 3 = third) diff --git a/tests/unit/test_variables_addendum_not_persisted.py b/tests/unit/test_variables_addendum_not_persisted.py index e8f013a4f..431a49e3b 100644 --- a/tests/unit/test_variables_addendum_not_persisted.py +++ b/tests/unit/test_variables_addendum_not_persisted.py @@ -67,6 +67,7 @@ def _make_adapter(captured_outbound: List[List[dict]]) -> MagicMock: a.on_response_processed.return_value = None a.classify_auto_continue = AsyncMock(return_value=False) a.resolve_bind_tools = AsyncMock(return_value=None) + a.execute_call_model_fc = AsyncMock(return_value=None) # base seam: no-op in CodeAct var_manager = MagicMock() var_manager.get_variable_names.return_value = ["my_var"] From cafcca14749ff30a1a8eb898bd1cf2e641a2be20 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 15 Sep 2026 11:15:14 +0300 Subject: [PATCH 02/18] =?UTF-8?q?feat(cuga-lite):=20step=20discipline=20?= =?UTF-8?q?=E2=80=94=20one=20tool=20call=20per=20step,=20both=20modes?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `advanced_features.cuga_lite_step_discipline = "one_tool_per_step"` (default "off"): the model observes one tool result before it makes the next call, whichever execution mode is active. - codeact: the sandbox forces the per-block tool-call cap to 1 for the block (contextvar override on the #560 budget, restored on exit). The second call is refused before it runs — neither executed, recorded nor counted — with a message that says so, and the executor keeps the block's variables on the refusal so the first call's result survives into the next block. - function_calling: tool_exec runs the first tool_call of the turn and answers the rest with a "deferred" ToolMessage, so every id still gets a reply; the FC prompt gains the one-call line. - configurable per invoke (`cuga_lite_step_discipline`), per model profile, in settings.toml, and from the SDK (`step_discipline=` on the constructor, invoke() and stream()). Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/prepare_node.py | 10 +- .../nodes/cuga_lite/adapter/sandbox_node.py | 38 ++++- .../nodes/cuga_lite/adapter/tool_exec_node.py | 17 ++ .../executors/local/local_executor.py | 10 +- .../tests/test_step_discipline_codeact.py | 160 ++++++++++++++++++ .../nodes/cuga_lite/model_runtime_profile.py | 49 +++++- .../nodes/cuga_lite/prompt_utils.py | 3 + .../nodes/cuga_lite/prompts/fc_prompt.jinja2 | 3 +- .../tests/test_execution_mode_config.py | 71 ++++++-- .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 30 ++++ .../nodes/cuga_lite/tests/test_fc_prompt.py | 10 +- .../cuga_lite/tests/test_tool_exec_node.py | 56 +++++- .../nodes/cuga_lite/tracking/tracker.py | 39 ++++- src/cuga/config.py | 3 +- src/cuga/sdk.py | 31 +++- src/cuga/settings.toml | 4 + 16 files changed, 497 insertions(+), 37 deletions(-) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/tests/test_step_discipline_codeact.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index dda9efc69..4f6e694f2 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -35,8 +35,10 @@ ) from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( EXECUTION_MODE_FUNCTION_CALLING, + STEP_DISCIPLINE_ONE_TOOL_PER_STEP, resolve_execution_mode, resolve_fc_prompt_fragments, + resolve_step_discipline, resolved_runtime_model_name, ) from cuga.backend.cuga_graph.nodes.cuga_lite.providers.langchain import DirectLangChainToolsProvider @@ -742,15 +744,21 @@ async def _wrapped(*args, **kwargs): # travel natively via bind_tools, so this prompt is only the behavioural # contract plus the same instructions / special_instructions. _fragments = resolve_fc_prompt_fragments(configurable, _runtime_model_name) + _one_per_step = ( + resolve_step_discipline(configurable, _runtime_model_name) + == STEP_DISCIPLINE_ONE_TOOL_PER_STEP + ) dynamic_prompt = render_fc_prompt( instructions=effective_instructions, special_instructions=special_instructions_final, is_autonomous_subtask=settings.advanced_features.force_autonomous_mode or is_autonomous_subtask, + step_discipline=_one_per_step, fragments=_fragments, ) logger.info( - "Prepared CugaLite function-calling prompt: fragments={} prompt_chars={}", + "Prepared CugaLite function-calling prompt: step_discipline={} fragments={} prompt_chars={}", + _one_per_step, _fragments, len(dynamic_prompt), ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py index 8b94fc5ed..87f124e45 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py @@ -23,6 +23,11 @@ CodeExecutor, is_find_tools_listing_markdown, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( + STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + resolve_step_discipline, + resolved_runtime_model_name, +) from cuga.backend.cuga_graph.nodes.cuga_lite.reflection.pre_execute import ( decide_pre_execute_verify, log_pre_execute_verify, @@ -243,14 +248,33 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): _exec_plan.filesystem_backend, ) logger.debug(f"\n\n------\n\n📝 Generated code:\n\n{state.script}\n\n------\n\n") - output, new_vars = await CodeExecutor.eval_with_tools_async( - code=state.script, - _locals=context, - state=state, # Pass CugaLiteState - it has variables_manager property - thread_id=current_thread_id, - apps_list=current_apps_list, - plan=_exec_plan, + # Step discipline: cap this block at one tool call. The cap rides the + # existing per-block budget, so the second call is refused by + # enforce_call_budget with the one-tool message and the executor + # keeps the block's variables (see LocalExecutor). + one_tool_per_step = ( + resolve_step_discipline( + configurable, + resolved_runtime_model_name( + configurable_llm=configurable.get("llm"), + graph_default_model=getattr(adapter, "_model", None), + ), + ) + == STEP_DISCIPLINE_ONE_TOOL_PER_STEP ) + cap_token = ToolCallTracker.set_block_cap_override(1) if one_tool_per_step else None + try: + output, new_vars = await CodeExecutor.eval_with_tools_async( + code=state.script, + _locals=context, + state=state, # Pass CugaLiteState - it has variables_manager property + thread_id=current_thread_id, + apps_list=current_apps_list, + plan=_exec_plan, + ) + finally: + if cap_token is not None: + ToolCallTracker.reset_block_cap_override(cap_token) adapter._tracker.collect_step(step=Step(name="User_output", data=output)) adapter._tracker.collect_step( diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py index 7c83e6dc9..4e56c2d7f 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py @@ -41,6 +41,11 @@ create_error_command as core_create_error_command, ) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import _budget_updates +from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( + STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + resolve_step_discipline, + resolved_runtime_model_name, +) from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( ToolCallBudgetExceeded, ToolCallTracker, @@ -48,6 +53,11 @@ ) from cuga.config import settings +DEFERRED_CALL_MESSAGE = ( + "Deferred: step discipline is on, so only the first tool call of a turn runs. " + "Read that result, then re-issue this call in your next turn if you still need it." +) + TRUNCATION_MARKER = "\n... [result truncated to {limit} characters]" @@ -107,6 +117,10 @@ async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): configurable = config.get("configurable", {}) if config else {} track_tool_calls = configurable.get("track_tool_calls", False) max_steps = configurable.get("cuga_lite_max_steps") if "cuga_lite_max_steps" in configurable else None + model_name = resolved_runtime_model_name( + configurable_llm=configurable.get("llm"), graph_default_model=getattr(adapter, "_model", None) + ) + one_per_step = resolve_step_discipline(configurable, model_name) == STEP_DISCIPLINE_ONE_TOOL_PER_STEP timeout = getattr(settings.advanced_features, "tool_call_timeout", 30) or None output_limit = int(getattr(settings.advanced_features, "execution_output_max_length", 0) or 0) @@ -146,6 +160,9 @@ async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): if budget_stop: results.append(_error_message(budget_stop, call_id=call_id, name=name)) continue + if one_per_step and executed >= 1: + results.append(_error_message(DEFERRED_CALL_MESSAGE, call_id=call_id, name=name)) + continue fn = adapter._tools_context.get(name) if fn is None or not callable(fn): diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/local/local_executor.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/local/local_executor.py index c54e12126..273db2810 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/local/local_executor.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/local/local_executor.py @@ -11,7 +11,10 @@ from ..common.restricted_environment import RestrictedEnvironment from ..common.security import CodeSyntaxError, SecurityValidator from ..common.benchmark_mode import is_relaxed_execution -from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import BlockToolCallCounter +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( + BlockToolCallCounter, + ToolCallBudgetExceeded, +) class _BlockSystemExit: @@ -180,6 +183,11 @@ async def _run_block(): captured = stdout_buf.getvalue() if captured: e.captured_stdout = captured # type: ignore[attr-defined] + if isinstance(e, ToolCallBudgetExceeded): + # A refused tool call is not a bug in the block: keep what it + # computed before the refusal (same as a timeout / exit()), so + # under step discipline the one call that did run is not lost. + context_locals.update(self._locals_from_frame(e, "_async_main")) raise result = stdout_buf.getvalue() diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/tests/test_step_discipline_codeact.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/tests/test_step_discipline_codeact.py new file mode 100644 index 000000000..b0a1fe362 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/tests/test_step_discipline_codeact.py @@ -0,0 +1,160 @@ +"""Step discipline in CodeAct: one tool call per code block. + +``cuga_lite_step_discipline = "one_tool_per_step"`` rides the per-block budget +from #560 with the cap forced to 1 for the block. What has to hold, and is +only visible on a real graph run: + +- the second call in a block is refused before it runs, with the one-tool + message, and is neither executed nor recorded nor counted; +- the block's variables — including the first call's result — survive the + refusal, so the next block can use them (the executor keeps the frame); +- with discipline off every call runs, exactly as before. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.tools import StructuredTool +from langgraph.checkpoint.memory import MemorySaver + +from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState, create_cuga_lite_graph +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( + BlockToolCallBudgetExceeded, + ToolCallTracker, + tracked_tool, +) + +pytestmark = pytest.mark.unit + +CALLS: list = [] + +TWO_CALL_BLOCK = ( + "```python\nfirst = await echo(value=1)\nsecond = await echo(value=2)\nprint(first, second)\n```" +) +USE_FIRST_BLOCK = "```python\nprint('kept:', first + 10)\n```" +FINAL = "First was 1." + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + self.invocations = 0 + + def bind_tools(self, *args, **kwargs): + return self + + async def ainvoke(self, messages, config=None, **kwargs): + self.invocations += 1 + if not self._responses: + raise AssertionError( + f"model asked for response #{self.invocations} — the run did not end when it should" + ) + return AIMessage(content=self._responses.pop(0)) + + +def _provider(): + # A custom-provider tool records itself via @tracked_tool (direct LangChain + # tools get the same wrapper from prepare). Recording sits *inside* the + # budget wrapper the executor adds, which is what makes "refused, therefore + # not recorded" observable below. + @tracked_tool(app_name="test_app") + async def echo(value: int) -> int: + """Echo a value.""" + CALLS.append(value) + return value + + provider = MagicMock() + provider.get_all_tools = AsyncMock( + return_value=[StructuredTool.from_function(coroutine=echo, name="echo", description="Echo a value.")] + ) + provider.get_apps = AsyncMock(return_value=[]) + provider.get_tools = AsyncMock(return_value=[]) + provider.app_name = "test_app" + return provider + + +@pytest.fixture(autouse=True) +def _clean(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_block", 100, raising=False) + CALLS.clear() + yield + CALLS.clear() + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + tracker_module._block_tool_call_cap_override_context.set(None) + + +def _run(responses, **configurable): + model = _ScriptedModel(responses) + graph = create_cuga_lite_graph( + model=model, tool_provider=_provider(), apps_list=[], thread_id="t" + ).compile(checkpointer=MemorySaver()) + config = { + "configurable": {"thread_id": "sd", "enable_todos": False, "track_tool_calls": True, **configurable} + } + return graph.ainvoke( + CugaLiteState(chat_messages=[HumanMessage(content="echo twice")]), config=config + ), model + + +# ── tracker: the override and its restore ──────────────────────────────────── + + +def test_cap_override_refuses_the_second_call_with_the_one_tool_message_and_restores_on_reset(): + ToolCallTracker.seed_call_budget(0) + token = ToolCallTracker.set_block_cap_override(1) + try: + ToolCallTracker.seed_block_budget() + ToolCallTracker.enforce_call_budget() + with pytest.raises(BlockToolCallBudgetExceeded) as exc: + ToolCallTracker.enforce_call_budget() + assert "One tool call per step" in str(exc.value) + assert "were kept" in str(exc.value) + assert ToolCallTracker.get_run_budget_used() == 1, "a refused call is never counted" + finally: + ToolCallTracker.reset_block_cap_override(token) + + # Back to the settings cap (100 here): the same block budget now admits more calls. + ToolCallTracker.seed_block_budget() + for _ in range(5): + ToolCallTracker.enforce_call_budget() + assert ToolCallTracker.get_run_budget_used() == 6 + + +# ── graph: refuse, keep variables, do not record ───────────────────────────── + + +@pytest.mark.asyncio +async def test_one_tool_per_step_refuses_the_second_call_and_keeps_the_first_result(): + run, model = _run([TWO_CALL_BLOCK, USE_FIRST_BLOCK, FINAL], cuga_lite_step_discipline="one_tool_per_step") + result = await run + + assert CALLS == [1], f"only the first call of the block may run, got {CALLS}" + assert result["final_answer"] == FINAL and model.invocations == 3 + + outputs = [m.content for m in result["chat_messages"] if isinstance(m, HumanMessage)][1:] + assert "One tool call per step" in outputs[0], "the model is told why the block stopped" + assert "kept: 11" in outputs[1], "the first call's result survived the refusal into the next block" + + assert [c["name"] for c in result["tool_calls"]] == ["echo"], "the refused call is not recorded" + assert result["tool_calls_used_run"] == 1, "nor counted against the run budget" + + +@pytest.mark.asyncio +async def test_discipline_off_runs_every_call_as_before(): + run, model = _run([TWO_CALL_BLOCK, FINAL]) + result = await run + + assert CALLS == [1, 2] + assert result["tool_calls_used_run"] == 2 and model.invocations == 2 + assert "One tool call per step" not in " ".join( + m.content for m in result["chat_messages"] if isinstance(m, HumanMessage) + ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py index ef6c98f51..47ef3e71c 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py @@ -116,7 +116,7 @@ def resolve_bind_tools_fields( return mode_s, apps_list, tool_names_list, _bool_coerce(inc) -# ── Execution mode ─────────────────────────────────────────────────────────── +# ── Execution mode / step discipline ───────────────────────────────────────── # # Same three-layer resolution as the bind_tools keys above: ``configurable`` # (per invoke) overrides the per-model runtime profile overrides ``settings``. @@ -141,6 +141,23 @@ def resolve_bind_tools_fields( "toolcalling": EXECUTION_MODE_FUNCTION_CALLING, } +STEP_DISCIPLINE_OFF = "off" +STEP_DISCIPLINE_ONE_TOOL_PER_STEP = "one_tool_per_step" + +_STEP_DISCIPLINE_ALIASES: Dict[str, str] = { + "off": STEP_DISCIPLINE_OFF, + "none": STEP_DISCIPLINE_OFF, + "false": STEP_DISCIPLINE_OFF, + "0": STEP_DISCIPLINE_OFF, + "one_tool_per_step": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "one_tool_per_turn": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "one_tool_per_block": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "stepwise": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "one": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "true": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "1": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, +} + FC_PROMPT_FRAGMENT_EVIDENCE_FIRST = "evidence_first" KNOWN_FC_PROMPT_FRAGMENTS = frozenset({FC_PROMPT_FRAGMENT_EVIDENCE_FIRST}) @@ -200,6 +217,25 @@ def normalize_execution_mode(raw: Any) -> str: return EXECUTION_MODE_CODEACT +def normalize_step_discipline(raw: Any) -> str: + if raw is None: + return STEP_DISCIPLINE_OFF + if isinstance(raw, bool): + return STEP_DISCIPLINE_ONE_TOOL_PER_STEP if raw else STEP_DISCIPLINE_OFF + key = str(raw).strip().lower() + if key in _STEP_DISCIPLINE_ALIASES: + return _STEP_DISCIPLINE_ALIASES[key] + try: + from loguru import logger + + logger.warning( + "cuga_lite_step_discipline={!r} is not recognised; falling back to {!r}", raw, STEP_DISCIPLINE_OFF + ) + except Exception: + pass + return STEP_DISCIPLINE_OFF + + def resolve_execution_mode( configurable: Optional[Dict[str, Any]], model_name: Optional[str] = None, @@ -213,6 +249,17 @@ def resolve_execution_mode( return normalize_execution_mode(raw) +def resolve_step_discipline( + configurable: Optional[Dict[str, Any]], + model_name: Optional[str] = None, + *, + settings_fn: Optional[Callable[[], Any]] = None, +) -> str: + """Resolved step discipline: ``off`` or ``one_tool_per_step``. Applies to both modes.""" + raw = _layered("cuga_lite_step_discipline", configurable, model_name, settings_fn, STEP_DISCIPLINE_OFF) + return normalize_step_discipline(raw) + + def resolve_fc_prompt_fragments( configurable: Optional[Dict[str, Any]], model_name: Optional[str] = None, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py index bf2d7e01b..178d811fa 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py @@ -919,6 +919,7 @@ def render_fc_prompt( special_instructions: Optional[str] = None, base_prompt: Optional[str] = None, is_autonomous_subtask: bool = False, + step_discipline: bool = False, fragments: Optional[List[str]] = None, ) -> str: """System prompt for ``cuga_lite_execution_mode = "function_calling"``. @@ -929,6 +930,7 @@ def render_fc_prompt( rendered as text. The prompt only sets the behavioural contract: - the generic preamble (always), + - the step-discipline line when ``cuga_lite_step_discipline`` is on, - opt-in fragments from ``cuga_lite_fc_prompt_fragments`` (``evidence_first``), - then the agent's ``instructions`` and the caller's ``special_instructions``. @@ -939,6 +941,7 @@ def render_fc_prompt( rendered = _fc_prompt_template().format( base_prompt=base_prompt or "", is_autonomous_subtask=bool(is_autonomous_subtask), + step_discipline=bool(step_discipline), evidence_first="evidence_first" in frags, instructions=(instructions or "").strip(), special_instructions=(special_instructions or "").strip(), diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 index 9e7bc41c8..31e24f8d5 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 @@ -8,7 +8,8 @@ You are Cuga Agent, a helpful assistant that answers the user's request by calli {% endif %} # HOW TO ACT - Use the available tools via native function-calling. Do NOT write Python code, code fences, pseudo-code, or JSON tool-call blocks in your text — issue real tool calls only. -{% if evidence_first %}- You MUST call at least one tool before giving a final answer, and base the answer STRICTLY on values returned by your tool calls — never answer from prior knowledge, memory, or assumption. If you have not yet observed the needed value in a tool result, make another tool call. +{% if step_discipline %}- Make exactly ONE tool call at a time. Read its result, then decide the next call based on what you actually observed. Build the answer one verified hop at a time. +{% endif %}{% if evidence_first %}- You MUST call at least one tool before giving a final answer, and base the answer STRICTLY on values returned by your tool calls — never answer from prior knowledge, memory, or assumption. If you have not yet observed the needed value in a tool result, make another tool call. - Never guess an intermediate value; only use values returned by a previous tool call. {% endif %}- If a tool returns an error, reason about why and try a corrected call. diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py index fb4e6df73..0ec021b8b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py @@ -1,4 +1,4 @@ -"""Configuration surfaces for ``cuga_lite_execution_mode``. +"""Configuration surfaces for ``cuga_lite_execution_mode`` / ``cuga_lite_step_discipline``. One assertion per control plane — settings.toml, validators, ``configurable``, the per-model runtime profile, the SDK constructor / ``invoke`` / ``stream`` — @@ -19,9 +19,13 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( EXECUTION_MODE_CODEACT, EXECUTION_MODE_FUNCTION_CALLING, + STEP_DISCIPLINE_OFF, + STEP_DISCIPLINE_ONE_TOOL_PER_STEP, normalize_execution_mode, + normalize_step_discipline, resolve_execution_mode, resolve_fc_prompt_fragments, + resolve_step_discipline, ) pytestmark = pytest.mark.unit @@ -30,10 +34,11 @@ # ── 1. settings.toml + validators: the shipped default is the old behaviour ── -def test_settings_toml_defaults_to_codeact(): +def test_settings_toml_defaults_to_codeact_and_no_discipline(): from cuga.config import settings assert settings.advanced_features.cuga_lite_execution_mode == "codeact" + assert settings.advanced_features.cuga_lite_step_discipline == "off" assert list(settings.advanced_features.cuga_lite_fc_prompt_fragments) == [] @@ -43,11 +48,13 @@ def test_validators_supply_the_same_defaults_without_the_keys(): declared = {v.names[0]: v.default for v in validators if v.names} assert declared["advanced_features.cuga_lite_execution_mode"] == "codeact" + assert declared["advanced_features.cuga_lite_step_discipline"] == "off" assert declared["advanced_features.cuga_lite_fc_prompt_fragments"] == [] def test_real_settings_resolve_to_the_defaults(): assert resolve_execution_mode({}) == EXECUTION_MODE_CODEACT + assert resolve_step_discipline({}) == STEP_DISCIPLINE_OFF assert resolve_fc_prompt_fragments({}) == [] @@ -77,6 +84,24 @@ def test_unknown_execution_mode_falls_back_to_codeact_without_raising(): assert normalize_execution_mode(42) == EXECUTION_MODE_CODEACT +@pytest.mark.parametrize( + "raw, expected", + [ + (None, STEP_DISCIPLINE_OFF), + ("off", STEP_DISCIPLINE_OFF), + (False, STEP_DISCIPLINE_OFF), + ("0", STEP_DISCIPLINE_OFF), + ("bogus", STEP_DISCIPLINE_OFF), + ("one_tool_per_step", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), + ("One_Tool_Per_Step", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), + (True, STEP_DISCIPLINE_ONE_TOOL_PER_STEP), + ("stepwise", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), + ], +) +def test_step_discipline_aliases(raw, expected): + assert normalize_step_discipline(raw) == expected + + def test_fragments_are_normalised_deduped_and_unknown_dropped(): cfg = {"cuga_lite_fc_prompt_fragments": ["evidence_first", "Evidence_First", "not_a_fragment"]} assert resolve_fc_prompt_fragments(cfg) == ["evidence_first"] @@ -93,6 +118,10 @@ def test_configurable_beats_settings(): resolve_execution_mode({"cuga_lite_execution_mode": "fc"}, settings_mode_fn=lambda: "codeact") == EXECUTION_MODE_FUNCTION_CALLING ) + assert ( + resolve_step_discipline({"cuga_lite_step_discipline": "off"}, settings_fn=lambda: "one_tool_per_step") + == STEP_DISCIPLINE_OFF + ) def test_settings_apply_when_configurable_is_silent_or_blank(): @@ -111,12 +140,19 @@ def test_model_profile_sits_between_configurable_and_settings(monkeypatch): monkeypatch.setattr( mrp, "runtime_defaults_for_model", - lambda name: {"cuga_lite_execution_mode": "fc"} if name == "profiled-model" else {}, + lambda name: ( + {"cuga_lite_execution_mode": "fc", "cuga_lite_step_discipline": "one"} + if name == "profiled-model" + else {} + ), ) # profile beats settings assert resolve_execution_mode({}, "profiled-model", settings_mode_fn=lambda: "codeact") == ( EXECUTION_MODE_FUNCTION_CALLING ) + assert resolve_step_discipline({}, "profiled-model", settings_fn=lambda: "off") == ( + STEP_DISCIPLINE_ONE_TOOL_PER_STEP + ) # configurable beats profile assert ( resolve_execution_mode( @@ -134,46 +170,55 @@ def test_model_profile_sits_between_configurable_and_settings(monkeypatch): # ── 4. SDK: constructor default, per-call override, raw keys win ───────────── -def test_cuga_agent_accepts_execution_mode_everywhere(): +def test_cuga_agent_accepts_execution_mode_and_step_discipline_everywhere(): from cuga.sdk import CugaAgent for method in (CugaAgent.__init__, CugaAgent.invoke, CugaAgent.stream): params = inspect.signature(method).parameters assert "execution_mode" in params, method.__name__ + assert "step_discipline" in params, method.__name__ assert params["execution_mode"].default is None + assert params["step_discipline"].default is None -def test_apply_execution_mode_writes_the_configurable_key(): +def test_apply_execution_mode_writes_the_configurable_keys(): from cuga.sdk import CugaAgent - agent = SimpleNamespace(_execution_mode="function_calling") + agent = SimpleNamespace(_execution_mode="function_calling", _step_discipline="one_tool_per_step") run_config = {"configurable": {}} CugaAgent._apply_execution_mode(agent, run_config) - assert run_config["configurable"] == {"cuga_lite_execution_mode": "function_calling"} + assert run_config["configurable"] == { + "cuga_lite_execution_mode": "function_calling", + "cuga_lite_step_discipline": "one_tool_per_step", + } -def test_per_invoke_value_overrides_the_constructor_default(): +def test_per_invoke_values_override_the_constructor_default(): from cuga.sdk import CugaAgent - agent = SimpleNamespace(_execution_mode="function_calling") + agent = SimpleNamespace(_execution_mode="function_calling", _step_discipline=None) run_config = {"configurable": {}} - CugaAgent._apply_execution_mode(agent, run_config, execution_mode="codeact") + CugaAgent._apply_execution_mode( + agent, run_config, execution_mode="codeact", step_discipline="one_tool_per_step" + ) assert run_config["configurable"]["cuga_lite_execution_mode"] == "codeact" + assert run_config["configurable"]["cuga_lite_step_discipline"] == "one_tool_per_step" def test_explicit_raw_keys_are_never_clobbered(): from cuga.sdk import CugaAgent - agent = SimpleNamespace(_execution_mode="function_calling") - run_config = {"configurable": {"cuga_lite_execution_mode": "codeact"}} + agent = SimpleNamespace(_execution_mode="function_calling", _step_discipline="one_tool_per_step") + run_config = {"configurable": {"cuga_lite_execution_mode": "codeact", "cuga_lite_step_discipline": "off"}} CugaAgent._apply_execution_mode(agent, run_config, execution_mode="fc") assert run_config["configurable"]["cuga_lite_execution_mode"] == "codeact" + assert run_config["configurable"]["cuga_lite_step_discipline"] == "off" def test_apply_execution_mode_is_a_noop_when_unconfigured(): from cuga.sdk import CugaAgent - agent = SimpleNamespace(_execution_mode=None) + agent = SimpleNamespace(_execution_mode=None, _step_discipline=None) run_config = {"configurable": {}} CugaAgent._apply_execution_mode(agent, run_config) assert run_config["configurable"] == {}, "settings.toml must decide when nothing is passed" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index 6b03d05df..7cc3d51a6 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -24,6 +24,7 @@ FC_MODE_VIOLATION_CORRECTION, AgentGraphAdapter, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import DEFERRED_CALL_MESSAGE from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState, create_cuga_lite_graph from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module @@ -100,6 +101,7 @@ def _quiet(monkeypatch): tracker_module._tool_call_budget_context.set(None) tracker_module._thread_tool_call_budget_context.set(None) tracker_module._block_tool_call_budget_context.set(None) + tracker_module._block_tool_call_cap_override_context.set(None) def _run(model, config, *tools): @@ -189,6 +191,34 @@ def final(messages): assert model.invocations == 4 +@pytest.mark.asyncio +async def test_step_discipline_in_fc_mode_runs_one_call_per_turn(): + model = _ScriptedModel( + [ + AIMessage( + content="", tool_calls=[_tc("echo", {"value": 1}, "c1"), _tc("echo", {"value": 2}, "c2")] + ), + AIMessage(content="", tool_calls=[_tc("echo", {"value": 2}, "c3")]), + AIMessage(content="1 and 2."), + ] + ) + config = _config( + "fc-sd", cuga_lite_execution_mode="function_calling", cuga_lite_step_discipline="one_tool_per_step" + ) + + result = await _run(model, config, _echo_tool()) + + assert CALLS == [("echo", 1), ("echo", 2)], "the second call of turn 1 is deferred, then re-issued" + tool_msgs = [m for m in result["chat_messages"] if isinstance(m, ToolMessage)] + assert [(m.tool_call_id, m.content) for m in tool_msgs][:2] == [ + ("c1", "1"), + ("c2", DEFERRED_CALL_MESSAGE), + ] + assert tool_msgs[2].tool_call_id == "c3" and tool_msgs[2].content == "2" + assert "exactly ONE tool call" in result["prepared_prompt"] + assert result["final_answer"] == "1 and 2." + + @pytest.mark.asyncio async def test_code_block_in_fc_mode_is_corrected_not_executed(): model = _ScriptedModel( diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py index 34103cd9e..4f0f06260 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py @@ -11,6 +11,7 @@ pytestmark = pytest.mark.unit _TEMPLATE = Path(__file__).resolve().parents[1] / "prompts" / "fc_prompt.jinja2" +_STEP_LINE = "Make exactly ONE tool call at a time" _EVIDENCE_LINE = "You MUST call at least one tool before giving a final answer" @@ -19,10 +20,14 @@ def test_default_prompt_is_the_bare_contract(): assert text.startswith("# ROLE") assert "native function-calling" in text and "# FINAL ANSWER" in text assert "Do NOT write Python code" in text - assert _EVIDENCE_LINE not in text + assert _STEP_LINE not in text and _EVIDENCE_LINE not in text assert "autonomously" not in text +def test_step_discipline_adds_the_one_call_line(): + assert _STEP_LINE in render_fc_prompt(step_discipline=True) + + def test_fragments_are_opt_in_and_unknown_names_are_ignored(): assert _EVIDENCE_LINE in render_fc_prompt(fragments=["evidence_first"]) assert _EVIDENCE_LINE in render_fc_prompt(fragments=["Evidence_First"]) @@ -35,13 +40,14 @@ def test_autonomous_subtask_line(): def test_instructions_then_special_instructions_close_the_prompt(): text = render_fc_prompt( + step_discipline=True, fragments=["evidence_first"], instructions=" Agent instructions. ", special_instructions="Answer format: one line.", ) body = text.rstrip("\n") assert body.endswith("Agent instructions.\n\nAnswer format: one line.") - assert text.index(_EVIDENCE_LINE) < text.index("# FINAL ANSWER") + assert text.index(_STEP_LINE) < text.index(_EVIDENCE_LINE) < text.index("# FINAL ANSWER") assert "\n\n\n" not in text, "no blank-line runs from unset optional sections" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py index 90649ce3d..2a23fe0a6 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py @@ -5,7 +5,7 @@ call even though the callables in ``_tools_context`` are bare, and every exit carries the budget fields. On top of that, the FC-specific contract: ids are preserved, one failing call never aborts its siblings, provider-side invalid -calls are answered. +calls are answered, and step discipline defers everything but the first call. """ from __future__ import annotations @@ -20,6 +20,7 @@ from langgraph.types import Command from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import ( + DEFERRED_CALL_MESSAGE, create_tool_exec_node, ) from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module @@ -74,6 +75,7 @@ def _reset_budget_contexts(): tracker_module._tool_call_budget_context.set(None) tracker_module._thread_tool_call_budget_context.set(None) tracker_module._block_tool_call_budget_context.set(None) + tracker_module._block_tool_call_cap_override_context.set(None) def _caps(monkeypatch, *, block=0, run=0, thread=0): @@ -290,6 +292,58 @@ async def add(a: int, b: int) -> int: assert untracked["tool_calls"] == [] +# ── step discipline ────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_step_discipline_runs_only_the_first_call_and_defers_the_rest(monkeypatch): + _caps(monkeypatch) + calls = [] + + async def add(a: int, b: int) -> int: + calls.append((a, b)) + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage( + content="", + tool_calls=[ + _call("add", {"a": 1, "b": 1}, "c1"), + _call("add", {"a": 2, "b": 2}, "c2"), + _call("add", {"a": 3, "b": 3}, "c3"), + ], + ) + + result = await node( + _state(last), config={"configurable": {"cuga_lite_step_discipline": "one_tool_per_step"}} + ) + + m1, m2, m3 = _tool_messages(result) + assert calls == [(1, 1)] + assert (m1.tool_call_id, m1.content, m1.status) == ("c1", "2", "success") + assert (m2.tool_call_id, m2.content, m2.status) == ("c2", DEFERRED_CALL_MESSAGE, "error") + assert (m3.tool_call_id, m3.content, m3.status) == ("c3", DEFERRED_CALL_MESSAGE, "error") + assert result["tool_calls_used_run"] == 1 + + +@pytest.mark.asyncio +async def test_step_discipline_off_runs_every_call(monkeypatch): + _caps(monkeypatch) + calls = [] + + async def add(a: int, b: int) -> int: + calls.append((a, b)) + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage( + content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] + ) + + await node(_state(last), config={"configurable": {"cuga_lite_step_discipline": "off"}}) + assert calls == [(1, 1), (2, 2)] + + # ── routing edge cases ─────────────────────────────────────────────────────── diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tracking/tracker.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tracking/tracker.py index d9303321e..c26bbe824 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tracking/tracker.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tracking/tracker.py @@ -73,6 +73,13 @@ async def get_users(limit: int = 10) -> list: "block_tool_call_budget", default=None ) +# Per-execution override of max_tool_calls_per_block. Step discipline +# (cuga_lite_step_discipline = "one_tool_per_step") sets it to 1 around each +# code block; None means settings.toml decides. +_block_tool_call_cap_override_context: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar( + "block_tool_call_cap_override", default=None +) + # Holds a mutable counter dict so the count survives context copies: # ``asyncio.wait_for`` runs each code block in a new Task whose context is a # *copy* of the executor's, but the copy references the SAME dict, so @@ -275,6 +282,21 @@ def seed_block_budget() -> None: """Open a fresh per-block budget. Called once per executed code block.""" _block_tool_call_budget_context.set([0]) + @staticmethod + def set_block_cap_override(cap: Optional[int]) -> contextvars.Token: + """Override ``max_tool_calls_per_block`` for the current execution context. + + Step discipline sets ``cap=1`` so a code block can make exactly one tool + call. Returns the token for :meth:`reset_block_cap_override`, so a + delegated child graph that runs on the caller's Task restores its + caller's cap on the way out instead of clearing it. + """ + return _block_tool_call_cap_override_context.set(cap) + + @staticmethod + def reset_block_cap_override(token: contextvars.Token) -> None: + _block_tool_call_cap_override_context.reset(token) + @staticmethod def get_run_budget_used() -> int: """Tool calls made so far this turn (0 when no budget is active).""" @@ -337,7 +359,12 @@ def enforce_call_budget() -> None: max_tool_calls_per_run = getattr(settings.advanced_features, "max_tool_calls_per_run", 256) max_per_thread = getattr(settings.advanced_features, "max_tool_calls_per_thread", 2000) - max_per_block = getattr(settings.advanced_features, "max_tool_calls_per_block", 100) + cap_override = _block_tool_call_cap_override_context.get() + max_per_block = ( + cap_override + if cap_override is not None + else getattr(settings.advanced_features, "max_tool_calls_per_block", 100) + ) thread_box = _thread_tool_call_budget_context.get() block_box = _block_tool_call_budget_context.get() @@ -355,6 +382,16 @@ def enforce_call_budget() -> None: "Do not call any more tools — produce a final answer from the data already retrieved. " "(Configurable via advanced_features.max_tool_calls_per_run; 0 disables.)" ) + if max_per_block == 1 and block_box is not None and block_box[0] >= 1: + # Step discipline: one verified hop per block. The executor keeps the + # variables computed before the refused call, so the first call's + # result is not lost — the model just has to use it in a new block. + raise BlockToolCallBudgetExceeded( + "One tool call per step: this code block already made its one tool call, so this " + "call was refused and nothing after it ran. The variables computed before it — " + "including the result of the call that did run — were kept and are available in " + "the next block. Read that result, then make your next single tool call in a new block." + ) if max_per_block and block_box is not None and block_box[0] >= max_per_block: raise BlockToolCallBudgetExceeded( f"Tool call limit reached for this code block: it made {max_per_block} tool calls. " diff --git a/src/cuga/config.py b/src/cuga/config.py index d2cf46b53..dd144151a 100644 --- a/src/cuga/config.py +++ b/src/cuga/config.py @@ -220,8 +220,9 @@ def get_all_paths(config, parent_key=""): Validator("advanced_features.cuga_lite_bind_tools_tool_names", default=[]), Validator("advanced_features.cuga_lite_bind_tools_max_count", default=128), Validator("advanced_features.cuga_lite_bind_tools_pad_to_cap", default=False), - # Native function-calling execution mode (default off). + # Native function-calling execution mode + step discipline (both default off). Validator("advanced_features.cuga_lite_execution_mode", default="codeact"), + Validator("advanced_features.cuga_lite_step_discipline", default="off"), Validator("advanced_features.cuga_lite_fc_prompt_fragments", default=[]), # Read at prepare_node without a getattr guard; without this a settings.toml # missing the key raises AttributeError mid-run instead of using the default. diff --git a/src/cuga/sdk.py b/src/cuga/sdk.py index c43be34aa..078b77a14 100644 --- a/src/cuga/sdk.py +++ b/src/cuga/sdk.py @@ -1856,6 +1856,7 @@ def __init__( shortlister: Optional["Shortlister"] = None, final_answer: "Optional[Union[str, Callable[[str], str], FinalAnswerConfig]]" = None, execution_mode: Optional[str] = None, + step_discipline: Optional[str] = None, ): """ Initialize the CUGA Agent. @@ -1882,6 +1883,9 @@ def __init__( Python that calls tools in the sandbox) or `"function_calling"` (the model emits native tool calls, executed one round-trip at a time). None = use `advanced_features.cuga_lite_execution_mode`. Overridable per invoke()/stream(). + step_discipline: `"one_tool_per_step"` lets the model run only one tool per step + (in both execution modes); `"off"` (default) allows any number. None = use + `advanced_features.cuga_lite_step_discipline`. Overridable per invoke()/stream(). final_answer: How the final answer is shaped. A `str` adds guidance for the answer-composing LLM; a callable `(str) -> str` deterministically post-processes the composed answer before delivery (must be pure; @@ -1989,9 +1993,10 @@ def __init__( # is the pre-feature behavior. self._shortlister = shortlister - # Execution mode. None => settings.toml (default "codeact", i.e. the - # pre-feature behavior). + # Execution mode / step discipline. None => settings.toml (defaults: + # "codeact" / "off", i.e. the pre-feature behavior). self._execution_mode = execution_mode + self._step_discipline = step_discipline # Setup tool provider. ToolGuard is installed immediately as a transparent # provider-level decorator so create-agent-first, add-guard-later flows work. @@ -2138,18 +2143,22 @@ def _apply_execution_mode( self, run_config: dict, execution_mode: Optional[str] = None, + step_discipline: Optional[str] = None, ) -> None: - """Merge the execution mode into ``run_config['configurable']``. + """Merge execution mode / step discipline into ``run_config['configurable']``. - A per-invoke value overrides the constructor default. A raw - ``cuga_lite_execution_mode`` key already set by the caller wins over - both — we ``setdefault`` rather than overwrite. + Per-invoke values override the constructor defaults. Raw + ``cuga_lite_execution_mode`` / ``cuga_lite_step_discipline`` keys already + set by the caller win over both — we ``setdefault`` rather than overwrite. No-op when nothing is configured, so settings.toml decides. """ configurable = run_config["configurable"] mode = execution_mode if execution_mode is not None else self._execution_mode if mode is not None: configurable.setdefault("cuga_lite_execution_mode", mode) + discipline = step_discipline if step_discipline is not None else self._step_discipline + if discipline is not None: + configurable.setdefault("cuga_lite_step_discipline", discipline) def _apply_callbacks( self, run_config: dict, extra_callbacks: Optional[List[BaseCallbackHandler]] = None @@ -2610,6 +2619,7 @@ async def invoke( variables: Optional[Dict[str, Any]] = None, shortlister: Optional["Shortlister"] = None, execution_mode: Optional[str] = None, + step_discipline: Optional[str] = None, ) -> InvokeResult: """ Invoke the agent with a message and get the response. @@ -2630,6 +2640,8 @@ async def invoke( `Shortlister(strategy="hybrid")`. None = use the agent default. execution_mode: `"codeact"` or `"function_calling"` for this call only. None = use the agent default. + step_discipline: `"off"` or `"one_tool_per_step"` for this call only. + None = use the agent default. Returns: InvokeResult containing: @@ -2705,7 +2717,7 @@ async def invoke( # Setup config (shallow-copied so we don't mutate the caller's dict) run_config = self._prepare_run_config(config) self._apply_shortlister(run_config, shortlister) - self._apply_execution_mode(run_config, execution_mode) + self._apply_execution_mode(run_config, execution_mode, step_discipline) # Pass track_tool_calls flag via configurable run_config["configurable"]["track_tool_calls"] = track_tool_calls @@ -3136,6 +3148,7 @@ async def stream( action_response: Optional[Any] = None, # ActionResponse for resuming after HITL shortlister: Optional["Shortlister"] = None, execution_mode: Optional[str] = None, + step_discipline: Optional[str] = None, ): """ Stream the agent's execution step by step. @@ -3153,6 +3166,8 @@ async def stream( `Shortlister(strategy="hybrid")`. None = use the agent default. execution_mode: `"codeact"` or `"function_calling"` for this call only. None = use the agent default. + step_discipline: `"off"` or `"one_tool_per_step"` for this call only. + None = use the agent default. Yields: State updates as the agent executes @@ -3185,7 +3200,7 @@ async def stream( # Setup config (shallow-copied so we don't mutate the caller's dict) run_config = self._prepare_run_config(config) self._apply_shortlister(run_config, shortlister) - self._apply_execution_mode(run_config, execution_mode) + self._apply_execution_mode(run_config, execution_mode, step_discipline) # Pass skills configuration via configurable (overrides settings when set) if self._enable_skills is not None: diff --git a/src/cuga/settings.toml b/src/cuga/settings.toml index 0f3e1326d..3e05fdc5a 100644 --- a/src/cuga/settings.toml +++ b/src/cuga/settings.toml @@ -58,6 +58,10 @@ cuga_lite_bind_tools_pad_to_cap = false # issues native tool_calls, a tool_exec node runs them and replies with ToolMessages). Aliases: fc, native. # Per-invoke: config["configurable"]["cuga_lite_execution_mode"]. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_EXECUTION_MODE cuga_lite_execution_mode = "codeact" +# Step discipline, both modes: off (default) | one_tool_per_step — observe one result before the next call. +# codeact: only the first tool call in a generated block runs (per-block cap becomes 1). function_calling: only +# the first tool_call of a turn runs; the rest are answered with a "deferred" ToolMessage. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_STEP_DISCIPLINE +cuga_lite_step_discipline = "off" # Opt-in prompt fragments for function_calling mode. Known: "evidence_first" (call a tool before answering; never # answer from prior knowledge). Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_FC_PROMPT_FRAGMENTS='["evidence_first"]' cuga_lite_fc_prompt_fragments = [] From dd3edea722641179c5140c08a47b29c9b6e0ab2b Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:28:14 +0300 Subject: [PATCH 03/18] =?UTF-8?q?fix(cuga-lite):=20function-calling=20revi?= =?UTF-8?q?ew=20fixes=20=E2=80=94=20valid=20replay=20across=20turns,=20fai?= =?UTF-8?q?l-closed=20approval,=20parity=20gaps?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of #777 found the mode unsafe to switch on as shipped. Reproduced each finding before fixing it: - Multi-turn through the SDK replayed bare BaseMessage shells: state crosses the SDK / server boundary via state.model_dump() against List[BaseMessage], so tool_calls and tool_call_id are dropped and turn 2 sent the provider objects it rejects. History is now normalised into a provider-valid transcript before every model call (shells retyped, lost tool results rendered as text, every tool_calls id answered), and a CodeAct turn on the same thread renders native tool results as text instead of skipping them with a warning. - A step-limit or budget breach could persist an assistant turn whose tool_calls nobody answered; those ids now get synthetic error replies. - Tool-approval policies were silently bypassed (the approval check runs on the code string, after the seam). Function-calling now refuses to start while an enabled tool-approval policy exists, before any model call, with a clear error. - An empty reply gets one corrective turn, as on the CodeAct path. - Bind mode "none" is upgraded to exactly the executable set prepare built (per sub-task / relevant apps), not the registry-wide catalogue, so a delegated sub-task cannot see tools it cannot run. - The bundled CodeAct few-shot demos are withheld in function-calling mode (explicit demos are still replayed); the variables addendum reaches the model; the static-prompt short-circuit warns. - tool_exec: outer guard mirrors the sandbox (a failure after the batch still reports the spent budget); the per-call bound is the block wall clock (sandbox_execution_timeout), tools keep their own inner timeout; one batch shares one execution_output_max_length; langfuse callbacks sync. - Native calls and their results stream to the UI as code steps. - Aliases trimmed to the documented spellings; README section added. CodeRabbit (first pass): tool_exec routes with Command (no static edge, so a terminal error never schedules one more model turn); step discipline defers by position, so an error on the first call is still the one result the model reads next; the skill and agent catalogues are rendered in the function-calling prompt so load_skill / delegate_to_* get valid names. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 35 ++ .../cuga_agent_core/graph/graph_nodes.py | 1 + .../cuga_agent_core/graph/shared_graph.py | 15 +- .../cuga_agent_core/graph/shared_nodes.py | 17 + .../tests/graph/test_call_model_fc_seam.py | 29 +- .../nodes/cuga_lite/adapter/graph_adapter.py | 293 ++++++++++++--- .../nodes/cuga_lite/adapter/prepare_node.py | 30 +- .../nodes/cuga_lite/adapter/tool_exec_node.py | 345 +++++++++++------- .../nodes/cuga_lite/model_runtime_profile.py | 14 +- .../nodes/cuga_lite/prompt_utils.py | 6 + .../nodes/cuga_lite/prompts/fc_prompt.jinja2 | 6 +- .../tests/test_execution_mode_config.py | 19 +- .../tests/test_fc_call_model_hook.py | 187 +++++++++- .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 29 ++ .../nodes/cuga_lite/tests/test_fc_prompt.py | 21 +- .../cuga_lite/tests/test_tool_exec_node.py | 136 ++++++- .../backend/cuga_graph/utils/agent_loop.py | 52 +++ tests/unit/test_agent_loop_fc_events.py | 51 +++ tests/unit/test_fc_sdk_two_turn.py | 104 ++++++ 19 files changed, 1144 insertions(+), 246 deletions(-) create mode 100644 tests/unit/test_agent_loop_fc_events.py create mode 100644 tests/unit/test_fc_sdk_two_turn.py diff --git a/README.md b/README.md index 0d3bc3d12..f44ab4401 100644 --- a/README.md +++ b/README.md @@ -1173,6 +1173,41 @@ See `docs/design/pluggable-shortlister.md` for the full design. +
+🔁 Execution mode: CodeAct or native function calling + +## What it is + +By default CUGA Lite works in **CodeAct**: the model writes a Python block and the sandbox runs it. `cuga_lite_execution_mode = "function_calling"` switches to native tool calls: the model emits `tool_calls`, a `tool_exec` node runs them under the same budgets, tracker and timeout as the sandbox and replies with `ToolMessage`s, and the model reads the results before it calls again or answers. Tools are advertised through `bind_tools`; a short dedicated prompt replaces the CodeAct prompt. + +**Step discipline** (`cuga_lite_step_discipline = "one_tool_per_step"`) works in both modes: only the first tool call of a step runs, so the model reads a result before it decides the next call. CodeAct keeps the block's variables (including that first result); function-calling answers the extra calls with a "deferred" reply. + +## Configuration + +| Key | Default | Values | +|---|---|---| +| `cuga_lite_execution_mode` | `"codeact"` | `codeact` \| `function_calling` (alias `fc`) | +| `cuga_lite_step_discipline` | `"off"` | `off` \| `one_tool_per_step` | +| `cuga_lite_fc_prompt_fragments` | `[]` | Opt-in prompt lines for function-calling mode. `evidence_first`: call a tool before answering, never answer from prior knowledge | + +Precedence, highest first: raw keys in `configurable` → per-invoke `invoke(..., execution_mode=..., step_discipline=...)` → constructor `CugaAgent(execution_mode=..., step_discipline=...)` → per-model runtime profile → `[advanced_features]` in `settings.toml` (also where `DYNACONF_ADVANCED_FEATURES__CUGA_LITE_EXECUTION_MODE` lands). An unknown value falls back to the default with a warning. + +```python +agent = CugaAgent(tools=[...], execution_mode="function_calling", step_discipline="one_tool_per_step") +result = await agent.invoke("...") # native tool calls +result = await agent.invoke("...", execution_mode="codeact") # per-call override +``` + +## Gotchas + +- **Tool-approval policies**: function-calling mode has no approval interrupt yet. If an enabled tool-approval policy exists, the run stops with a clear error before any tool runs — use `codeact` for that agent or disable the policy. +- Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call. Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. +- The bundled CodeAct few-shot demos are not sent in function-calling mode; pass your own through `configurable["mcp_few_shot_examples"]`. +- Tools run in-process and sequentially. The pre-execute VERIFY gate, reflection and the E2B / OpenSandbox executors are CodeAct-only; step discipline caps local-executor blocks only. +- A thread can switch modes between turns: results from function-calling turns are shown to a later CodeAct turn as text. + +
+
📝 Special Instructions Configuration diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py index 10a000b72..589641499 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/graph_nodes.py @@ -146,6 +146,7 @@ async def execute_call_model_fc( modified_messages: list, budget_exhausted: bool, playbook_fired: bool, + variables_addendum: str = "", ) -> Optional[Command]: """Native function-calling seam. diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py index acf1c2f4f..0deaf09cd 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py @@ -5,7 +5,7 @@ START → prepare --Command--> call_model ↔ execute (loop) → END CugaLite may add a fourth node, ``tool_exec`` (native function-calling -mode), which loops back into ``call_model`` the same way. +mode), which routes back into ``call_model`` with a Command. Both CugaLite and CugaSupervisor share this structure. The nodes themselves are provided by the caller (produced by adapter factories), so the graph @@ -48,9 +48,10 @@ def build_agent_graph( execute_node: Async node function for the execute/sandbox step. tool_exec_node: Optional node that executes native ``tool_calls`` and replies with ``ToolMessage``s (CugaLite function-calling mode). When - supplied it is added as ``"tool_exec"`` with a static edge back to - ``call_model``; ``call_model`` only routes to it in that mode, so it - is dormant on every CodeAct run. Omitted by the Supervisor graph. + supplied it is added as ``"tool_exec"``; it routes back to + ``call_model`` (or to END on error) with a Command, and ``call_model`` + only routes to it in that mode, so it is dormant on every CodeAct run. + Omitted by the Supervisor graph. Returns: An uncompiled ``StateGraph``. Call ``.compile(checkpointer=...)`` @@ -68,9 +69,9 @@ def build_agent_graph( graph.add_edge(adapter.execute_node_name, "call_model") if tool_exec_node is not None: + # No static edge back: the node routes with Command (call_model on a normal + # exit, END on an error exit), so a terminal error never schedules one more + # model turn the way a static edge would. graph.add_node("tool_exec", tool_exec_node) - # Same shape as the sandbox edge: the node returns a state update, and the - # loop closes back into call_model for the next model turn. - graph.add_edge("tool_exec", "call_model") return graph diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py index f7abdcaf9..c8ab8096f 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_nodes.py @@ -191,6 +191,22 @@ async def call_model(state: Any, config: RunnableConfig = None) -> Command: ai_content = msg.content if hasattr(msg, "content") else (msg.get("content") or "") messages_for_model.append({"role": "assistant", "content": ai_content}) + elif msg_role == "tool": + # A native tool result left by a function-calling turn on this thread. + # Matched on ``type``: persisted history crosses the SDK boundary as bare + # BaseMessage shells. Rendered like an execution output so a CodeAct + # turn on the same thread still sees what the tool returned. + modified_messages.append(msg) + tool_name = ( + getattr(msg, "name", None) + or (msg.get("name") if isinstance(msg, dict) else None) + or "tool" + ) + tool_content = msg.content if hasattr(msg, "content") else (msg.get("content") or "") + messages_for_model.append( + {"role": "user", "content": f"Tool result ({tool_name}):\n{tool_content}"} + ) + else: modified_messages.append(msg) logger.warning("call_model: skipping message {} with unknown role: {}", i, msg_role) @@ -243,6 +259,7 @@ async def call_model(state: Any, config: RunnableConfig = None) -> Command: modified_messages=modified_messages, budget_exhausted=budget_exhausted, playbook_fired=playbook_fired, + variables_addendum=variables_addendum, ) if fc_command is not None: return fc_command diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py index 316c2bc96..28b84310c 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/tests/graph/test_call_model_fc_seam.py @@ -13,7 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from langchain_core.messages import BaseMessage, HumanMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage from langgraph.graph import END from langgraph.types import Command @@ -136,3 +136,30 @@ async def test_a_command_from_the_hook_short_circuits_the_node(mock_summarize): ) assert kw["configurable"] == {"cuga_lite_max_steps": 7} assert kw["budget_exhausted"] is False and kw["playbook_fired"] is False + + +@pytest.mark.asyncio +@patch(_SUMMARIZE, new_callable=AsyncMock) +async def test_tool_results_left_by_a_function_calling_turn_reach_the_codeact_model(mock_summarize): + """Switching a thread from function-calling to codeact must not hide what the + tools returned. Both shapes occur: real ToolMessages inside one run, and bare + ``type="tool"`` shells after the SDK boundary dropped the subclass.""" + mock_summarize.side_effect = lambda messages, *a, **kw: messages + model = _mock_model("Plain answer.") + state = _make_state() + state.chat_messages = [ + HumanMessage(content="echo 7"), + AIMessage( + content="", tool_calls=[{"name": "echo", "args": {"value": 7}, "id": "c1", "type": "tool_call"}] + ), + ToolMessage(content="7", tool_call_id="c1", name="echo"), + BaseMessage(type="tool", content="8", name="echo"), + HumanMessage(content="what did the tools say?"), + ] + node = create_call_model_node(_BaseAdapter(), model, _mock_settings()) + + await node(state, config=None) + + sent = model.ainvoke.await_args[0][0] + assert {"role": "user", "content": "Tool result (echo):\n7"} in sent + assert {"role": "user", "content": "Tool result (echo):\n8"} in sent diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 943ff0df8..f56c35c21 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -20,8 +20,11 @@ format_task_todos_system_block, ) from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( + EMPTY_RESPONSE_CORRECTION, + EMPTY_RESPONSE_CORRECTION_KEY, EXECUTION_OUTPUT_PREFIX, CoreGraphAdapter, + create_error_command, enforce_step_limit, ) from cuga.backend.cuga_graph.utils.harmony import contains_harmony_tokens, strip_harmony_tokens @@ -61,37 +64,133 @@ "Code is never run here. Issue the tool call natively instead, or give the final answer as plain text." ) +FC_TOOL_APPROVAL_UNSUPPORTED = ( + "Function-calling mode does not support tool-approval policies yet. An enabled tool-approval " + "policy exists, so this run was stopped before any tool ran. Use cuga_lite_execution_mode = " + '"codeact" for this agent, or disable the policy.' +) + +FC_STEP_LIMIT_CALL_REPLY = "Not executed: the step limit was reached before this call could run." +FC_BUDGET_CALL_REPLY = ( + "Not executed: the tool budget for this turn is spent. Answer from the data already retrieved." +) +FC_UNANSWERED_CALL_REPLY = "No result was recorded for this call." + + +def _message_role(m: Any) -> str: + if isinstance(m, dict): + return str(m.get("role") or m.get("type") or "") + return str(getattr(m, "type", "") or "") + + +def _message_text(m: Any) -> str: + content = m.get("content") if isinstance(m, dict) else getattr(m, "content", None) + if content is None: + return "" + return content if isinstance(content, str) else str(content) + + +def _message_name(m: Any) -> Optional[str]: + return m.get("name") if isinstance(m, dict) else getattr(m, "name", None) + + +def _tool_result_as_text(name: Optional[str], content: str) -> str: + return f"Tool result ({name or 'tool'}):\n{content}" + + +def _unanswered_call_replies(calls: List[Any], invalid: List[Any], reason: str) -> List[ToolMessage]: + """One error ``ToolMessage`` per id the model issued when nothing will execute them. + + A persisted assistant turn whose ``tool_calls`` have no replies makes strict + providers reject the next replay of the thread, so every id is answered even + when the run ends here (step limit, spent budget). + """ + out: List[ToolMessage] = [] + for index, call in enumerate(list(calls or []) + list(invalid or [])): + call_id = (str(call.get("id") or "") if isinstance(call, dict) else "") or f"call_{index}" + name = (call.get("name") if isinstance(call, dict) else None) or "unknown" + out.append(ToolMessage(content=reason, tool_call_id=call_id, name=str(name), status="error")) + return out -def _sanitize_for_replay(messages: List[BaseMessage]) -> List[BaseMessage]: - """Copy history for replay to the provider. - Assistant turns are replayed verbatim (``tool_calls`` intact) — but their - ``reasoning_content`` is dropped and harmony framing stripped from string - content: gpt-oss emits both, and strict OpenAI-compatible proxies 400 when - they come back in a later round. Irrelevant in CodeAct, whose history is - flattened to text before it leaves; in function-calling mode the raw turns - are replayed from round 2 onward, so it matters here. +def _normalize_history_for_replay(messages: List[Any]) -> List[BaseMessage]: + """Persisted history as a provider-valid function-calling transcript. + + Two things go wrong with history as persisted. State crosses the SDK and + server boundary through ``state.model_dump()`` against ``List[BaseMessage]``, + so pydantic serialises by the declared type: subclass fields are dropped and + the messages come back as bare ``BaseMessage`` shells — an assistant turn + without its ``tool_calls``, a tool turn without its ``tool_call_id``. And a + turn can end on a call nobody answered. Either makes a strict provider + reject the replay (or raise on the unknown message type), so rebuild it: + + - bare shells become typed messages; a lost tool result is rendered as user + text, and an empty assistant shell (a lost tool-call turn) is dropped; + - every ``tool_calls`` id is followed by its ``ToolMessage`` — a synthetic + error reply when none was recorded; + - a ``ToolMessage`` that answers no open call is rendered as user text; + - assistant reasoning payloads are dropped and harmony framing stripped + (gpt-oss emits both; strict OpenAI-compatible proxies 400 on them). """ out: List[BaseMessage] = [] - for m in messages: - ak = getattr(m, "additional_kwargs", None) or {} - drop_reasoning = any(k in ak for k in _REASONING_KEYS) - content = getattr(m, "content", None) - new_content = ( - strip_harmony_tokens(content) if isinstance(content, str) and "<|" in content else content - ) - if not drop_reasoning and new_content is content: - out.append(m) + pending: Dict[str, str] = {} # call id -> tool name, from the last assistant turn + + def close_pending() -> None: + for call_id, name in pending.items(): + out.append( + ToolMessage( + content=FC_UNANSWERED_CALL_REPLY, + tool_call_id=call_id, + name=name or "unknown", + status="error", + ) + ) + pending.clear() + + for m in messages or []: + if isinstance(m, ToolMessage): + if m.tool_call_id in pending: + pending.pop(m.tool_call_id, None) + out.append(m) + else: + close_pending() + out.append(HumanMessage(content=_tool_result_as_text(m.name, _message_text(m)))) continue - update: Dict[str, Any] = {} - if drop_reasoning: - update["additional_kwargs"] = {k: v for k, v in ak.items() if k not in _REASONING_KEYS} - if new_content is not content: - update["content"] = new_content - try: - out.append(m.model_copy(update=update)) - except Exception: + close_pending() + if isinstance(m, AIMessage): + ak = m.additional_kwargs or {} + update: Dict[str, Any] = {} + if any(k in ak for k in _REASONING_KEYS): + update["additional_kwargs"] = {k: v for k, v in ak.items() if k not in _REASONING_KEYS} + if isinstance(m.content, str) and "<|" in m.content: + update["content"] = strip_harmony_tokens(m.content) + msg = m.model_copy(update=update) if update else m + if msg.tool_calls: + for i, c in enumerate(msg.tool_calls): + pending[str(c.get("id") or f"call_{i}")] = str(c.get("name") or "") + out.append(msg) + elif _message_text(msg).strip(): + out.append(msg) + continue # an empty assistant shell is dropped + if isinstance(m, (HumanMessage, SystemMessage)): out.append(m) + continue + role = _message_role(m) + text = _message_text(m) + if role in ("human", "user"): + out.append(HumanMessage(content=text)) + elif role in ("ai", "assistant"): + if "<|" in text: + text = strip_harmony_tokens(text) + if text.strip(): + out.append(AIMessage(content=text)) + elif role == "tool": + out.append(HumanMessage(content=_tool_result_as_text(_message_name(m), text))) + elif role == "system": + out.append(SystemMessage(content=text)) + elif text.strip(): + out.append(HumanMessage(content=text)) + close_pending() return out @@ -231,8 +330,25 @@ async def resolve_bind_tools( # settings is respected. Additive: no effect in codeact. if self._execution_mode(configurable) == EXECUTION_MODE_FUNCTION_CALLING: if self._resolved_bind_mode(configurable) == "none": - configurable = {**(configurable or {}), "cuga_lite_bind_tools_mode": "all"} - logger.info("[fc] bind_tools mode was 'none'; upgraded to 'all' for function-calling mode") + # Advertise exactly the executable set prepare built (filtered per + # sub-task / relevant apps), not the registry-wide catalogue: a tool + # the model can see but the sandbox could not call is an + # "Unknown tool" reply waiting to happen. + names = list((self._tools_context_ref or {}).get("_lc_bind_tools_executable_names") or []) + if names: + configurable = { + **(configurable or {}), + "cuga_lite_bind_tools_mode": "tools", + "cuga_lite_bind_tools_tool_names": names, + } + logger.info( + "[fc] bind_tools mode was 'none'; advertising the {} executable tool(s)", len(names) + ) + else: + configurable = {**(configurable or {}), "cuga_lite_bind_tools_mode": "all"} + logger.info( + "[fc] bind_tools mode was 'none'; upgraded to 'all' for function-calling mode" + ) try: return await resolve_model_with_bind_tools( active_model, @@ -338,6 +454,27 @@ def _resolved_bind_mode(self, configurable: dict) -> str: return str(candidate).strip().lower() return _bind_tools_mode_from_settings() + async def _tool_approval_policies_exist(self, config: Any) -> bool: + """True when an enabled tool-approval policy is configured. + + Function-calling has no approval interrupt yet — CodeAct's runs on the + generated code string, after the seam — so the mode refuses to start + rather than run a guarded tool unprompted. On a storage error this + warns and proceeds, exactly as CodeAct's own approval check does. + """ + from cuga.backend.cuga_graph.policy.configurable import PolicyConfigurable + from cuga.backend.cuga_graph.policy.models import PolicyType + + try: + policy_system = PolicyConfigurable.from_config(config or {}) + policies = await policy_system.agent.storage.list_policies( + policy_type=PolicyType.TOOL_APPROVAL, enabled_only=True, limit=1 + ) + return bool(policies) + except Exception as exc: + logger.warning("[fc] could not query tool-approval policies ({}); proceeding", exc) + return False + async def execute_call_model_fc( self, *, @@ -351,19 +488,25 @@ async def execute_call_model_fc( modified_messages: list, budget_exhausted: bool, playbook_fired: bool, + variables_addendum: str = "", ) -> Optional[Command]: """Function-calling turn: invoke with real message objects, route on ``tool_calls``. Returns ``None`` in codeact mode so the shared CodeAct path runs untouched. Otherwise builds the outbound list from ``system_content`` (the FC prompt that ``prepare`` selected), the few-shot demos as chat messages, and the - sanitized history — hands them straight to ``bound.ainvoke`` so the shared - dict serializer (which flattens ``tool_calls``) is never involved — then: + history normalised into a provider-valid transcript — hands them straight + to ``bound.ainvoke`` so the shared dict serializer is never involved — then: - ``tool_calls`` present -> ``Command(goto="tool_exec")``, assistant turn kept verbatim - otherwise -> final answer, ``END`` + - an empty reply gets one corrective turn, like the CodeAct path - a fenced code block with no ``tool_calls`` is a mode violation: it is never executed; the model gets one corrective turn instead. + + Refuses to start when an enabled tool-approval policy exists: there is + no approval interrupt on this path yet, and silently running a guarded + tool is worse than not running at all. """ if self._execution_mode(configurable) != EXECUTION_MODE_FUNCTION_CALLING: return None @@ -372,11 +515,31 @@ async def execute_call_model_fc( TOOL_BUDGET_EXHAUSTED_INSTRUCTION, ) + cfg = configurable or {} + history: list = list(modified_messages) + + if settings.policy.enabled and await self._tool_approval_policies_exist(config): + logger.error( + "[fc] refusing to start: an enabled tool-approval policy exists and " + "function-calling mode has no approval interrupt" + ) + return create_error_command( + self, history, AIMessage(content=FC_TOOL_APPROVAL_UNSUPPORTED), state.step_count + ) + msgs: List[BaseMessage] = [SystemMessage(content=system_content)] msgs.extend(_few_shot_to_messages(self.get_few_shot_messages(state))) - msgs.extend(_sanitize_for_replay(list(modified_messages))) + msgs.extend(_normalize_history_for_replay(history)) + if variables_addendum: + # Outbound only, like call_model's CodeAct path (#600): never persisted. + for i in range(len(msgs) - 1, -1, -1): + if isinstance(msgs[i], HumanMessage): + msgs[i] = msgs[i].model_copy( + update={"content": _message_text(msgs[i]) + variables_addendum} + ) + break if budget_exhausted: - # Outbound only, like call_model's CodeAct path: never persisted. + # Outbound only as well. msgs.append(HumanMessage(content=TOOL_BUDGET_EXHAUSTED_INSTRUCTION)) clamp_watsonx_completion_for_messages(bound, msgs) @@ -391,11 +554,18 @@ async def execute_call_model_fc( if not isinstance(response, AIMessage): response = AIMessage(content=content, tool_calls=tool_calls) - if budget_exhausted: - # No tools were bound for the grace turn, so any tool_calls are noise. + max_steps = self.resolve_max_steps(state, cfg.get("cuga_lite_max_steps")) + new_step_count: int = state.step_count + 1 + final_messages: list = history + [response] + + if budget_exhausted and (tool_calls or invalid_tool_calls): + # No tools were bound for the grace turn, so the calls are noise — but + # every id still gets a reply, or the persisted thread cannot be replayed. + final_messages += _unanswered_call_replies(tool_calls, invalid_tool_calls, FC_BUDGET_CALL_REPLY) tool_calls, invalid_tool_calls = [], [] - if tool_calls or invalid_tool_calls: + has_calls = bool(tool_calls or invalid_tool_calls) + if has_calls: try: self._tracker.collect_step( step=Step( @@ -408,25 +578,30 @@ async def execute_call_model_fc( else: self.on_response_processed(state, None, content, reasoning) - final_messages: list = list(modified_messages) + [response] - new_step_count: int = state.step_count + 1 + # Step limit. On a breach with calls pending they are answered first, so + # the persisted transcript never ends on a dangling tool_calls turn. limit_cmd = ( None if budget_exhausted else enforce_step_limit( self, state=state, - messages=final_messages, + messages=final_messages + + ( + _unanswered_call_replies(tool_calls, invalid_tool_calls, FC_STEP_LIMIT_CALL_REPLY) + if has_calls + else [] + ), new_step_count=new_step_count, - limit=self.resolve_max_steps(state, (configurable or {}).get("cuga_lite_max_steps")), + limit=max_steps, ) ) if limit_cmd is not None: return limit_cmd - meta_update = {self.metadata_key: self.build_metadata_update(state, playbook_fired=playbook_fired)} + base_meta = dict(self.build_metadata_update(state, playbook_fired=playbook_fired) or {}) - if tool_calls or invalid_tool_calls: + if has_calls: logger.info("[fc] {} native tool_call(s) -> tool_exec", len(tool_calls) + len(invalid_tool_calls)) return Command( goto="tool_exec", @@ -434,16 +609,34 @@ async def execute_call_model_fc( self.messages_key: final_messages, "script": None, "step_count": new_step_count, - **meta_update, + self.metadata_key: base_meta, + }, + ) + + # Empty reply: one retry, same contract as the CodeAct path (#756). + both_blank = not content.strip() and not (reasoning or "").strip() + already_retried = bool(self.get_metadata(state).get(EMPTY_RESPONSE_CORRECTION_KEY)) + if both_blank and not already_retried and not budget_exhausted and new_step_count < max_steps: + logger.warning( + "[fc] model returned an empty reply (no content, no reasoning, no tool_calls) — retrying once" + ) + retry_meta = {**base_meta, EMPTY_RESPONSE_CORRECTION_KEY: True} + return Command( + goto="call_model", + update={ + self.messages_key: final_messages + [HumanMessage(content=EMPTY_RESPONSE_CORRECTION)], + "script": None, + "final_answer": "", + "execution_complete": False, + "step_count": new_step_count, + self.metadata_key: retry_meta, }, ) if "```" in content and not budget_exhausted: # Mode violation: never execute code here. One corrective turn, charged # as a step so it cannot loop past cuga_lite_max_steps. - meta = dict(self.get_metadata(state)) - violations = int(meta.get("fc_mode_violations", 0) or 0) + 1 - self.set_metadata(state, {**meta, "fc_mode_violations": violations}) + violations = int(base_meta.get("fc_mode_violations", 0) or 0) + 1 logger.warning("[fc] mode violation #{}: code block emitted in function-calling mode", violations) return Command( goto="call_model", @@ -453,7 +646,7 @@ async def execute_call_model_fc( "final_answer": "", "execution_complete": False, "step_count": new_step_count, - self.metadata_key: self.build_metadata_update(state, playbook_fired=playbook_fired), + self.metadata_key: {**base_meta, "fc_mode_violations": violations}, }, ) @@ -461,12 +654,12 @@ async def execute_call_model_fc( if not final_answer.strip() and reasoning and not contains_harmony_tokens(reasoning): final_answer = reasoning.strip() if not final_answer.strip(): - for m in reversed(modified_messages): - if isinstance(m, ToolMessage) and isinstance(m.content, str) and m.content.strip(): - final_answer = m.content + for m in reversed(history): + if _message_role(m) == "tool" and _message_text(m).strip(): + final_answer = _message_text(m) break if not content.strip() and final_answer: - final_messages = list(modified_messages) + [AIMessage(content=final_answer)] + final_messages = history + [AIMessage(content=final_answer)] logger.info("[fc] no tool_calls -> final answer (END)") return Command( @@ -477,7 +670,7 @@ async def execute_call_model_fc( "final_answer": final_answer, "execution_complete": True, "step_count": new_step_count, - **meta_update, + self.metadata_key: base_meta, }, ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index 4f6e694f2..8193dfb1c 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -135,6 +135,7 @@ async def prepare_tools_and_apps(state: Any, config: Optional[RunnableConfig] = configurable, model_name=_runtime_model_name, ) + _execution_mode = resolve_execution_mode(configurable, _runtime_model_name) logger.debug( f"[APPROVAL DEBUG] prepare_tools_and_apps received cuga_lite_metadata: {state.cuga_lite_metadata}" ) @@ -295,6 +296,18 @@ async def prepare_tools_and_apps(state: Any, config: Optional[RunnableConfig] = else: few_shot_examples = [] logger.debug("MCP few-shots disabled (cuga_lite_enable_few_shots=false)") + if ( + _execution_mode == EXECUTION_MODE_FUNCTION_CALLING + and few_shot_examples + and configurable.get("mcp_few_shot_examples") is None + ): + # The bundled demos are CodeAct-shaped (```python turns, "Execution output:" + # replies) and would contradict the function-calling prompt. Only demos the + # caller passes explicitly are replayed in this mode. + logger.info( + "Function-calling mode: bundled CodeAct few-shots withheld ({} turns)", len(few_shot_examples) + ) + few_shot_examples = [] if few_shot_examples: logger.debug(f"MCP few-shot examples: {len(few_shot_examples)} turns") @@ -721,6 +734,15 @@ async def _wrapped(*args, **kwargs): lc_bind_tools_meta["_lc_bind_tools_overlay_structured_tools"] = [ t for t in (tools_for_prompt or []) if getattr(t, "name", None) ] + # Function-calling advertises exactly what generated code could call: the + # executable set (already filtered per sub-task / relevant apps, knowledge + # tools already popped when out of scope) plus the in-graph overlays. + _candidates = [getattr(t, "name", None) for t in (tools_for_execution or [])] + [ + getattr(t, "name", None) for t in (tools_for_prompt or []) + ] + lc_bind_tools_meta["_lc_bind_tools_executable_names"] = [ + n for n in dict.fromkeys(_candidates) if n and n in adapter._tools_context + ] # Use tools_for_execution, not tools_for_prompt: when find_tools shortlisting # is active (the common case once an app has more than a handful of tools), @@ -738,7 +760,6 @@ async def _wrapped(*args, **kwargs): # Create prompt dynamically dynamic_prompt = adapter._static_prompt - _execution_mode = resolve_execution_mode(configurable, _runtime_model_name) if not dynamic_prompt and _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: # Function-calling mode: the CodeAct prompt is never rendered. Tools # travel natively via bind_tools, so this prompt is only the behavioural @@ -755,6 +776,8 @@ async def _wrapped(*args, **kwargs): or is_autonomous_subtask, step_discipline=_one_per_step, fragments=_fragments, + skills_prompt_section=skills_prompt_section if skills_enabled else "", + agents_prompt_section=agents_prompt_section if agents_enabled else "", ) logger.info( "Prepared CugaLite function-calling prompt: step_discipline={} fragments={} prompt_chars={}", @@ -795,6 +818,11 @@ async def _wrapped(*args, **kwargs): len(dynamic_prompt), ) else: + if _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: + logger.warning( + "Function-calling mode with a static prompt: the prompt is used verbatim; " + "make sure it does not instruct the model to write code" + ) logger.info( "Using static CugaLite prompt; dynamic few-shot injection skipped " "(enable_find_tools={} few_shot_turns={})", diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py index 4e56c2d7f..114a3d0ad 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py @@ -32,6 +32,7 @@ from langchain_core.messages import AIMessage, ToolMessage from langchain_core.runnables import RunnableConfig +from langgraph.types import Command from loguru import logger from cuga.backend.activity_tracker.tracker import Step @@ -54,11 +55,11 @@ from cuga.config import settings DEFERRED_CALL_MESSAGE = ( - "Deferred: step discipline is on, so only the first tool call of a turn runs. " + "Deferred: step discipline is on, so only the first tool call of a turn is attempted. " "Read that result, then re-issue this call in your next turn if you still need it." ) -TRUNCATION_MARKER = "\n... [result truncated to {limit} characters]" +TRUNCATION_MARKER = "\n... [result truncated: {limit} characters of output remain for this batch]" def _tool_call_parts(call: Any) -> tuple[str, str, Any]: @@ -95,17 +96,34 @@ def _coerce_args(args: Any) -> tuple[Optional[Dict[str, Any]], Optional[str]]: return None, f"Tool arguments must be an object, got {type(args).__name__}." -def _stringify(result: Any, limit: int) -> str: +def _stringify(result: Any) -> str: if isinstance(result, str): - text = result - else: - try: - text = json.dumps(result, ensure_ascii=False, default=str) - except Exception: - text = str(result) - if limit and len(text) > limit: - return text[:limit] + TRUNCATION_MARKER.format(limit=limit) - return text + return result + try: + return json.dumps(result, ensure_ascii=False, default=str) + except Exception: + return str(result) + + +class _BatchOutputBudget: + """One turn's batch of results shares ``execution_output_max_length``, like one block. + + Per-result truncation alone would let N calls return N x the limit; this keeps + the whole batch under the ceiling the sandbox applies to a block's output. + """ + + def __init__(self, limit: Any): + self.limit = int(limit or 0) + self.remaining = self.limit + + def take(self, text: str) -> str: + if not self.limit: + return text + if len(text) > self.remaining: + cap = max(self.remaining, 0) + text = text[:cap] + TRUNCATION_MARKER.format(limit=cap) + self.remaining = max(self.remaining - len(text), 0) + return text def _error_message(text: str, *, call_id: str, name: str) -> ToolMessage: @@ -115,14 +133,22 @@ def _error_message(text: str, *, call_id: str, name: str) -> ToolMessage: def create_tool_exec_node(adapter: Any) -> Callable: async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): configurable = config.get("configurable", {}) if config else {} + from cuga.backend.cuga_graph.utils.langfuse_tracing import sync_langfuse_callbacks_from_config + + sync_langfuse_callbacks_from_config(config) track_tool_calls = configurable.get("track_tool_calls", False) max_steps = configurable.get("cuga_lite_max_steps") if "cuga_lite_max_steps" in configurable else None model_name = resolved_runtime_model_name( configurable_llm=configurable.get("llm"), graph_default_model=getattr(adapter, "_model", None) ) one_per_step = resolve_step_discipline(configurable, model_name) == STEP_DISCIPLINE_ONE_TOOL_PER_STEP - timeout = getattr(settings.advanced_features, "tool_call_timeout", 30) or None - output_limit = int(getattr(settings.advanced_features, "execution_output_max_length", 0) or 0) + # One batch is the function-calling analogue of one code block, so the + # outer bound per call is the block wall clock. Registry tools keep their + # own inner ``tool_call_timeout``. + timeout = getattr(settings.advanced_features, "sandbox_execution_timeout", 30) or None + output_budget = _BatchOutputBudget( + getattr(settings.advanced_features, "execution_output_max_length", 0) + ) messages = adapter.get_messages(state) last = messages[-1] if messages else None @@ -147,134 +173,191 @@ async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): results: List[ToolMessage] = [] executed = 0 - budget_stop: Optional[str] = None + execution_tool_calls: List[Dict[str, Any]] = [] try: - for index, call in enumerate(calls): - call_id, name, raw_args = _tool_call_parts(call) - call_id = call_id or f"call_{index}" - if not name: - results.append( - _error_message("Tool call carried no tool name.", call_id=call_id, name="") - ) - continue - if budget_stop: - results.append(_error_message(budget_stop, call_id=call_id, name=name)) - continue - if one_per_step and executed >= 1: - results.append(_error_message(DEFERRED_CALL_MESSAGE, call_id=call_id, name=name)) - continue - - fn = adapter._tools_context.get(name) - if fn is None or not callable(fn): - known = ", ".join(sorted(k for k in adapter._tools_context if not k.startswith("_"))) - results.append( - _error_message( - f"Unknown tool '{name}'. Choose one of the provided tools: {known or '(none)'}.", - call_id=call_id, - name=name, - ) - ) - continue - - kwargs, arg_error = _coerce_args(raw_args) - if arg_error: - results.append(_error_message(arg_error, call_id=call_id, name=name)) - continue - - awaitable = fn if inspect.iscoroutinefunction(fn) else make_tool_awaitable(fn) - counted = counted_tool_call(awaitable) - try: - result = await asyncio.wait_for(counted(**kwargs), timeout=timeout) - executed += 1 - results.append( - ToolMessage(content=_stringify(result, output_limit), tool_call_id=call_id, name=name) - ) - except asyncio.TimeoutError: - results.append( - _error_message( - f"Tool '{name}' timed out after {timeout}s. Try a narrower call or a different tool.", - call_id=call_id, - name=name, - ) - ) - except ToolCallBudgetExceeded as exc: - # Every remaining id still gets a reply — without invoking anything. - if exc.scope != "block": - budget_stop = str(exc) - results.append(_error_message(str(exc), call_id=call_id, name=name)) - except TypeError as exc: - results.append( - _error_message( - f"Tool '{name}' rejected these arguments: {exc}. Check the parameter names and types.", - call_id=call_id, - name=name, - ) - ) - except Exception as exc: # one failing call must not abort its siblings - results.append(_error_message(f"Tool '{name}' failed: {exc}", call_id=call_id, name=name)) - - for index, bad in enumerate(invalid): - call_id, name, _ = _tool_call_parts(bad) - reason = (bad.get("error") if isinstance(bad, dict) else None) or "malformed tool call" - results.append( - _error_message( - f"The provider could not parse this tool call ({reason}). Re-issue it with valid arguments.", - call_id=call_id or f"invalid_{index}", - name=name, + try: + executed = await _run_batch( + adapter, + calls, + invalid, + results, + one_per_step=one_per_step, + timeout=timeout, + output_budget=output_budget, + ) + finally: + execution_tool_calls = ToolCallTracker.stop_tracking() + + if not calls and not invalid: + # Nothing executable: advance the step and hand back to the model so a + # mis-route can never become an infinite no-op loop. + logger.warning("tool_exec entered with no tool_calls on the last message") + return Command( + goto="call_model", + update={"step_count": state.step_count + 1, "script": None, **_budget_updates()}, + ) + + try: + adapter._tracker.collect_step( + step=Step( + name="Tool_results", + data=json.dumps( + [ + { + "tool_call_id": m.tool_call_id, + "name": m.name, + "status": m.status, + "content": m.content, + } + for m in results + ], + ensure_ascii=False, + default=str, + ), ) ) - finally: - execution_tool_calls = ToolCallTracker.stop_tracking() + except Exception as exc: + logger.debug(f"tool_exec tracker error: {exc}") - if not calls and not invalid: - # Nothing executable: advance the step and hand back to the model so a - # mis-route can never become an infinite no-op loop. - logger.warning("tool_exec entered with no tool_calls on the last message") - return {"step_count": state.step_count + 1, "script": None, **_budget_updates()} + updated_messages, error_message = core_append_with_step_limit(adapter, state, results, max_steps) + accumulated_tool_calls = (state.tool_calls or []) + ( + execution_tool_calls if track_tool_calls else [] + ) - try: - adapter._tracker.collect_step( - step=Step( - name="Tool_results", - data=json.dumps( - [ - { - "tool_call_id": m.tool_call_id, - "name": m.name, - "status": m.status, - "content": m.content, - } - for m in results - ], - ensure_ascii=False, - default=str, - ), + if error_message: + return core_create_error_command( + adapter, + updated_messages, + error_message, + state.step_count, + additional_updates={"tool_calls": accumulated_tool_calls, **_budget_updates()}, ) + + logger.info( + "[fc] tool_exec: {} call(s) -> {} ToolMessage(s), {} executed", + len(calls), + len(results), + executed, + ) + return Command( + goto="call_model", + update={ + adapter.messages_key: updated_messages, + "script": None, + "step_count": state.step_count + 1, + "tool_calls": accumulated_tool_calls, + **_budget_updates(), + }, ) except Exception as exc: - logger.debug(f"tool_exec tracker error: {exc}") - - updated_messages, error_message = core_append_with_step_limit(adapter, state, results, max_steps) - accumulated_tool_calls = (state.tool_calls or []) + (execution_tool_calls if track_tool_calls else []) - - if error_message: + # Mirrors the sandbox's outer guard: a failure outside the per-call + # handling must still report the budget the batch spent, or the + # checkpoint keeps the pre-batch counts (see _budget_updates). + logger.error(f"tool_exec failed: {exc}") return core_create_error_command( adapter, - updated_messages, - error_message, + adapter.get_messages(state), + AIMessage(content=f"Error during tool execution: {exc}"), state.step_count, - additional_updates={"tool_calls": accumulated_tool_calls, **_budget_updates()}, + additional_updates={ + "tool_calls": list(state.tool_calls or []) + + (execution_tool_calls if track_tool_calls else []), + **_budget_updates(), + }, ) - logger.info( - "[fc] tool_exec: {} call(s) -> {} ToolMessage(s), {} executed", len(calls), len(results), executed - ) - return { - adapter.messages_key: updated_messages, - "script": None, - "step_count": state.step_count + 1, - "tool_calls": accumulated_tool_calls, - **_budget_updates(), - } - return tool_exec + + +async def _run_batch( + adapter: Any, + calls: List[Any], + invalid: List[Any], + results: List[ToolMessage], + *, + one_per_step: bool, + timeout: Optional[float], + output_budget: _BatchOutputBudget, +) -> int: + """Run one turn's calls in order, appending one ``ToolMessage`` per id to ``results``. + + Returns how many calls actually executed. Every outcome is a message, never + an exception, so one failing call cannot take its siblings down. + """ + executed = 0 + budget_stop: Optional[str] = None + for index, call in enumerate(calls): + call_id, name, raw_args = _tool_call_parts(call) + call_id = call_id or f"call_{index}" + if not name: + results.append(_error_message("Tool call carried no tool name.", call_id=call_id, name="")) + continue + if budget_stop: + results.append(_error_message(budget_stop, call_id=call_id, name=name)) + continue + if one_per_step and index > 0: + # By position, not by success: an error on the first call is the result + # the model must read before it decides the next one. + results.append(_error_message(DEFERRED_CALL_MESSAGE, call_id=call_id, name=name)) + continue + + fn = adapter._tools_context.get(name) + if fn is None or not callable(fn): + known = ", ".join(sorted(k for k in adapter._tools_context if not k.startswith("_"))) + results.append( + _error_message( + f"Unknown tool '{name}'. Choose one of the provided tools: {known or '(none)'}.", + call_id=call_id, + name=name, + ) + ) + continue + + kwargs, arg_error = _coerce_args(raw_args) + if arg_error: + results.append(_error_message(arg_error, call_id=call_id, name=name)) + continue + + awaitable = fn if inspect.iscoroutinefunction(fn) else make_tool_awaitable(fn) + counted = counted_tool_call(awaitable) + try: + result = await asyncio.wait_for(counted(**kwargs), timeout=timeout) + executed += 1 + results.append( + ToolMessage(content=output_budget.take(_stringify(result)), tool_call_id=call_id, name=name) + ) + except asyncio.TimeoutError: + results.append( + _error_message( + f"Tool '{name}' timed out after {timeout}s. Try a narrower call or a different tool.", + call_id=call_id, + name=name, + ) + ) + except ToolCallBudgetExceeded as exc: + # Every remaining id still gets a reply — without invoking anything. + if exc.scope != "block": + budget_stop = str(exc) + results.append(_error_message(str(exc), call_id=call_id, name=name)) + except TypeError as exc: + results.append( + _error_message( + f"Tool '{name}' rejected these arguments: {exc}. Check the parameter names and types.", + call_id=call_id, + name=name, + ) + ) + except Exception as exc: # one failing call must not abort its siblings + results.append(_error_message(f"Tool '{name}' failed: {exc}", call_id=call_id, name=name)) + + for index, bad in enumerate(invalid): + call_id, name, _ = _tool_call_parts(bad) + reason = (bad.get("error") if isinstance(bad, dict) else None) or "malformed tool call" + results.append( + _error_message( + f"The provider could not parse this tool call ({reason}). Re-issue it with valid arguments.", + call_id=call_id or f"invalid_{index}", + name=name, + ) + ) + return executed diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py index 47ef3e71c..db3003c90 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/model_runtime_profile.py @@ -129,16 +129,9 @@ def resolve_bind_tools_fields( _EXECUTION_MODE_ALIASES: Dict[str, str] = { "codeact": EXECUTION_MODE_CODEACT, "code_act": EXECUTION_MODE_CODEACT, - "code": EXECUTION_MODE_CODEACT, - "sandbox": EXECUTION_MODE_CODEACT, "function_calling": EXECUTION_MODE_FUNCTION_CALLING, - "functioncalling": EXECUTION_MODE_FUNCTION_CALLING, "function-calling": EXECUTION_MODE_FUNCTION_CALLING, "fc": EXECUTION_MODE_FUNCTION_CALLING, - "native": EXECUTION_MODE_FUNCTION_CALLING, - "native_tool_calling": EXECUTION_MODE_FUNCTION_CALLING, - "tool_calling": EXECUTION_MODE_FUNCTION_CALLING, - "toolcalling": EXECUTION_MODE_FUNCTION_CALLING, } STEP_DISCIPLINE_OFF = "off" @@ -148,14 +141,9 @@ def resolve_bind_tools_fields( "off": STEP_DISCIPLINE_OFF, "none": STEP_DISCIPLINE_OFF, "false": STEP_DISCIPLINE_OFF, - "0": STEP_DISCIPLINE_OFF, "one_tool_per_step": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - "one_tool_per_turn": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - "one_tool_per_block": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - "stepwise": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - "one": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, + "one-tool-per-step": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, "true": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - "1": STEP_DISCIPLINE_ONE_TOOL_PER_STEP, } FC_PROMPT_FRAGMENT_EVIDENCE_FIRST = "evidence_first" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py index 178d811fa..01a9b2ed0 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py @@ -921,6 +921,8 @@ def render_fc_prompt( is_autonomous_subtask: bool = False, step_discipline: bool = False, fragments: Optional[List[str]] = None, + skills_prompt_section: Optional[str] = None, + agents_prompt_section: Optional[str] = None, ) -> str: """System prompt for ``cuga_lite_execution_mode = "function_calling"``. @@ -932,6 +934,8 @@ def render_fc_prompt( - the generic preamble (always), - the step-discipline line when ``cuga_lite_step_discipline`` is on, - opt-in fragments from ``cuga_lite_fc_prompt_fragments`` (``evidence_first``), + - the skill and agent catalogues when those tools are bound (``load_skill`` / + ``delegate_to_*`` take a name; the model has to be able to pick a valid one), - then the agent's ``instructions`` and the caller's ``special_instructions``. Answer-format contracts belong in ``special_instructions``; no @@ -943,6 +947,8 @@ def render_fc_prompt( is_autonomous_subtask=bool(is_autonomous_subtask), step_discipline=bool(step_discipline), evidence_first="evidence_first" in frags, + skills_prompt_section=(skills_prompt_section or "").strip(), + agents_prompt_section=(agents_prompt_section or "").strip(), instructions=(instructions or "").strip(), special_instructions=(special_instructions or "").strip(), ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 index 31e24f8d5..ddf99b083 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 @@ -15,7 +15,11 @@ You are Cuga Agent, a helpful assistant that answers the user's request by calli # FINAL ANSWER - When your tool results contain everything needed, STOP calling tools and give the final answer with no further tool call. -{% if instructions %} +{% if skills_prompt_section %} +{{ skills_prompt_section }} +{% endif %}{% if agents_prompt_section %} +{{ agents_prompt_section }} +{% endif %}{% if instructions %} {{ instructions }} {% endif %}{% if special_instructions %} {{ special_instructions }} diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py index 0ec021b8b..f74bd7f04 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py @@ -70,14 +70,18 @@ def test_real_settings_resolve_to_the_defaults(): ("function_calling", EXECUTION_MODE_FUNCTION_CALLING), (" Function-Calling ", EXECUTION_MODE_FUNCTION_CALLING), ("fc", EXECUTION_MODE_FUNCTION_CALLING), - ("native", EXECUTION_MODE_FUNCTION_CALLING), - ("tool_calling", EXECUTION_MODE_FUNCTION_CALLING), ], ) def test_execution_mode_aliases(raw, expected): assert normalize_execution_mode(raw) == expected +def test_execution_mode_accepts_only_the_documented_spellings(): + """Two canonical values plus ``fc``; anything else is a typo and falls back loudly.""" + for typo in ("native", "tool_calling", "functioncalling", "code", "sandbox", "1", "true"): + assert normalize_execution_mode(typo) == EXECUTION_MODE_CODEACT, typo + + def test_unknown_execution_mode_falls_back_to_codeact_without_raising(): assert normalize_execution_mode("banana") == EXECUTION_MODE_CODEACT assert normalize_execution_mode(None) == EXECUTION_MODE_CODEACT @@ -90,12 +94,13 @@ def test_unknown_execution_mode_falls_back_to_codeact_without_raising(): (None, STEP_DISCIPLINE_OFF), ("off", STEP_DISCIPLINE_OFF), (False, STEP_DISCIPLINE_OFF), - ("0", STEP_DISCIPLINE_OFF), ("bogus", STEP_DISCIPLINE_OFF), + ("1", STEP_DISCIPLINE_OFF), + ("stepwise", STEP_DISCIPLINE_OFF), ("one_tool_per_step", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), ("One_Tool_Per_Step", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), + ("one-tool-per-step", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), (True, STEP_DISCIPLINE_ONE_TOOL_PER_STEP), - ("stepwise", STEP_DISCIPLINE_ONE_TOOL_PER_STEP), ], ) def test_step_discipline_aliases(raw, expected): @@ -141,7 +146,7 @@ def test_model_profile_sits_between_configurable_and_settings(monkeypatch): mrp, "runtime_defaults_for_model", lambda name: ( - {"cuga_lite_execution_mode": "fc", "cuga_lite_step_discipline": "one"} + {"cuga_lite_execution_mode": "fc", "cuga_lite_step_discipline": "one_tool_per_step"} if name == "profiled-model" else {} ), @@ -268,7 +273,7 @@ def test_build_agent_graph_adds_tool_exec_only_when_given(): tool_exec_node=tool_exec, ) assert "tool_exec" in with_node.nodes - assert ("tool_exec", "call_model") in with_node.edges, "tool results always go back to the model" + assert ("tool_exec", "call_model") not in with_node.edges, "it routes with a Command, no static edge" def test_cuga_lite_graph_always_wires_tool_exec_so_mode_can_switch_per_invoke(): @@ -281,4 +286,4 @@ def test_cuga_lite_graph_always_wires_tool_exec_so_mode_can_switch_per_invoke(): graph = create_cuga_lite_graph(model=MagicMock(), tool_provider=provider, apps_list=[]) assert {"prepare", "call_model", "sandbox", "tool_exec"} <= set(graph.nodes) - assert ("tool_exec", "call_model") in graph.edges + assert ("tool_exec", "call_model") not in graph.edges diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index 1fe916fff..03bad6079 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -1,10 +1,12 @@ """CugaLite's ``execute_call_model_fc``: the function-calling turn. Codeact: ``None``, nothing invoked. Function-calling: the bound model gets real -message objects (system prompt, few-shot demos, sanitized history), and the -response routes on ``tool_calls`` — to ``tool_exec`` with the assistant turn -kept verbatim, or to END as the final answer. A fenced code block with no -``tool_calls`` is a mode violation: never executed, one corrective turn. +message objects (system prompt, few-shot demos, history normalised into a +provider-valid transcript), and the response routes on ``tool_calls`` — to +``tool_exec`` with the assistant turn kept verbatim, or to END as the final +answer. Every id the model issued is answered before the run can end, an empty +reply gets one retry, a fenced code block is a mode violation (never executed), +and the mode refuses to start while a tool-approval policy is configured. """ from __future__ import annotations @@ -14,20 +16,29 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage from langgraph.graph import END +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( + EMPTY_RESPONSE_CORRECTION, + EMPTY_RESPONSE_CORRECTION_KEY, +) from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import TOOL_BUDGET_EXHAUSTED_INSTRUCTION from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import graph_adapter as ga from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( + FC_BUDGET_CALL_REPLY, FC_MODE_VIOLATION_CORRECTION, + FC_STEP_LIMIT_CALL_REPLY, + FC_TOOL_APPROVAL_UNSUPPORTED, + FC_UNANSWERED_CALL_REPLY, AgentGraphAdapter, - _sanitize_for_replay, + _normalize_history_for_replay, ) pytestmark = pytest.mark.unit FC = {"cuga_lite_execution_mode": "function_calling"} +_CALL = {"name": "add", "args": {"a": 1, "b": 2}, "id": "c1", "type": "tool_call"} class _Model: @@ -42,27 +53,29 @@ async def ainvoke(self, messages, config=None, **kwargs): return self.response -def _adapter() -> AgentGraphAdapter: +def _adapter(tools_context_ref=None) -> AgentGraphAdapter: return AgentGraphAdapter( tracker=MagicMock(), base_callbacks=[], task_todos_ref=[], - tools_context_ref={}, + tools_context_ref=tools_context_ref if tools_context_ref is not None else {}, base_tool_provider=None, ) -def _state(messages=None, few_shots=None): +def _state(messages=None, few_shots=None, step_count=0, max_steps=None, metadata=None): return SimpleNamespace( chat_messages=messages or [HumanMessage(content="what is 1+2?")], - step_count=0, - cuga_lite_metadata={}, + step_count=step_count, + cuga_lite_metadata=metadata if metadata is not None else {}, mcp_few_shot_messages=few_shots or [], - cuga_lite_max_steps=None, + cuga_lite_max_steps=max_steps, ) -async def _turn(adapter, model, state, configurable, *, budget_exhausted=False, system="SYS"): +async def _turn( + adapter, model, state, configurable, *, budget_exhausted=False, system="SYS", variables_addendum="" +): return await adapter.execute_call_model_fc( state=state, config=None, @@ -74,10 +87,18 @@ async def _turn(adapter, model, state, configurable, *, budget_exhausted=False, modified_messages=list(state.chat_messages), budget_exhausted=budget_exhausted, playbook_fired=False, + variables_addendum=variables_addendum, ) -_CALL = {"name": "add", "args": {"a": 1, "b": 2}, "id": "c1", "type": "tool_call"} +@pytest.fixture(autouse=True) +def _policy_off(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + + +# ── routing ────────────────────────────────────────────────────────────────── @pytest.mark.asyncio @@ -117,13 +138,14 @@ async def test_plain_text_is_the_final_answer(): @pytest.mark.asyncio -async def test_empty_answer_falls_back_to_the_last_tool_result(): +async def test_empty_answer_after_a_retry_falls_back_to_the_last_tool_result(): history = [ HumanMessage(content="q"), AIMessage(content="", tool_calls=[_CALL]), ToolMessage(content="3", tool_call_id="c1", name="add"), ] - cmd = await _turn(_adapter(), _Model(AIMessage(content="")), _state(history), FC) + state = _state(history, metadata={EMPTY_RESPONSE_CORRECTION_KEY: True}) + cmd = await _turn(_adapter(), _Model(AIMessage(content="")), state, FC) assert cmd.goto == END and cmd.update["final_answer"] == "3" @@ -138,6 +160,23 @@ async def test_code_block_without_tool_calls_is_a_mode_violation_never_executed( assert cmd.update["chat_messages"][-1] == HumanMessage(content=FC_MODE_VIOLATION_CORRECTION) assert cmd.update["cuga_lite_metadata"]["fc_mode_violations"] == 1 assert cmd.update["step_count"] == 1, "charged as a step so it cannot loop forever" + assert state.cuga_lite_metadata == {}, "metadata travels in the update, state is not mutated in place" + + +@pytest.mark.asyncio +async def test_empty_reply_is_retried_once_then_finalized(): + cmd = await _turn(_adapter(), _Model(AIMessage(content="")), _state(), FC) + + assert cmd.goto == "call_model" + assert cmd.update["chat_messages"][-1] == HumanMessage(content=EMPTY_RESPONSE_CORRECTION) + assert cmd.update["cuga_lite_metadata"][EMPTY_RESPONSE_CORRECTION_KEY] is True + + retried = _state(metadata={EMPTY_RESPONSE_CORRECTION_KEY: True}) + cmd2 = await _turn(_adapter(), _Model(AIMessage(content="")), retried, FC) + assert cmd2.goto == END, "the marker survives exactly one turn — no retry loop" + + +# ── every issued id is answered before the run can end ─────────────────────── @pytest.mark.asyncio @@ -153,6 +192,33 @@ async def test_budget_exhausted_turn_ignores_tool_calls_and_sends_the_instructio assert all(m.content != TOOL_BUDGET_EXHAUSTED_INSTRUCTION for m in cmd.update["chat_messages"]), ( "never persisted" ) + ai, reply = cmd.update["chat_messages"][-2:] + assert ai.tool_calls == [_CALL] and isinstance(reply, ToolMessage) + assert (reply.tool_call_id, reply.content) == ("c1", FC_BUDGET_CALL_REPLY), ( + "no dangling call is persisted" + ) + + +@pytest.mark.asyncio +async def test_step_limit_breach_answers_pending_calls_before_the_error(): + """The persisted transcript must never end on a tool_calls turn nobody answered — + the next replay of the thread would be a provider 400.""" + model = _Model(AIMessage(content="", tool_calls=[_CALL])) + state = _state(step_count=1, max_steps=1) + + cmd = await _turn(_adapter(), model, state, FC) + + assert cmd.goto == END and "Maximum step limit" in cmd.update["final_answer"] + ai, reply, error = cmd.update["chat_messages"][-3:] + assert ai.tool_calls == [_CALL] + assert isinstance(reply, ToolMessage) and (reply.tool_call_id, reply.content) == ( + "c1", + FC_STEP_LIMIT_CALL_REPLY, + ) + assert isinstance(error, AIMessage) and "Maximum step limit" in error.content + + +# ── what the model is sent ─────────────────────────────────────────────────── @pytest.mark.asyncio @@ -168,6 +234,20 @@ async def test_few_shot_demos_go_out_as_chat_messages_before_history(): assert outbound[3].content == "what is 1+2?" +@pytest.mark.asyncio +async def test_variables_addendum_rides_the_last_human_turn_outbound_only(): + model = _Model(AIMessage(content="ok")) + state = _state([HumanMessage(content="first"), AIMessage(content="sure"), HumanMessage(content="last")]) + addendum = "\n\n## Available Variables\n\nvar_1: ..." + + cmd = await _turn(_adapter(), model, state, FC, variables_addendum=addendum) + + (outbound,) = model.seen + humans = [m for m in outbound if isinstance(m, HumanMessage)] + assert humans[-1].content == "last" + addendum and humans[0].content == "first" + assert all(addendum not in m.content for m in cmd.update["chat_messages"]), "never persisted (#600)" + + def test_replay_sanitizer_drops_reasoning_and_strips_harmony_without_touching_history(monkeypatch): monkeypatch.setattr(ga, "strip_harmony_tokens", lambda s: s.replace("<|channel|>final", "")) kept = HumanMessage(content="plain") @@ -175,13 +255,67 @@ def test_replay_sanitizer_drops_reasoning_and_strips_harmony_without_touching_hi content="<|channel|>finalanswer", additional_kwargs={"reasoning_content": "thinking...", "x": 1} ) - out = _sanitize_for_replay([kept, noisy]) + out = _normalize_history_for_replay([kept, noisy]) assert out[0] is kept, "untouched messages are passed through by identity" assert out[1].content == "answer" and out[1].additional_kwargs == {"x": 1} assert noisy.additional_kwargs["reasoning_content"] == "thinking...", "persisted history is never mutated" +def test_replay_rebuilds_bare_shells_and_closes_dangling_calls(): + """History that crossed the SDK boundary (bare BaseMessage shells, subclass fields + gone) and a turn that ended on an unanswered call both come back provider-valid.""" + shells = [ + BaseMessage(type="human", content="echo 7"), + BaseMessage(type="ai", content=""), # was AIMessage(tool_calls=[...]) before model_dump + BaseMessage(type="tool", content="7", name="echo"), # was ToolMessage(tool_call_id="call_1") + BaseMessage(type="ai", content="The value is 7."), + HumanMessage(content="again"), + AIMessage(content="", tool_calls=[_CALL]), # never answered: step limit / crash + ToolMessage(content="orphan", tool_call_id="nope", name="add"), # answers nothing + ] + + out = _normalize_history_for_replay(shells) + + assert [type(m).__name__ for m in out] == [ + "HumanMessage", + "HumanMessage", + "AIMessage", + "HumanMessage", + "AIMessage", + "ToolMessage", + "HumanMessage", + ], [type(m).__name__ for m in out] + assert out[1].content == "Tool result (echo):\n7", "a lost tool result is replayed as text" + assert out[4].tool_calls == [_CALL] and (out[5].tool_call_id, out[5].content) == ( + "c1", + FC_UNANSWERED_CALL_REPLY, + ) + assert out[6].content == "Tool result (add):\norphan" + assert all(type(m) is not BaseMessage for m in out), "no bare shell reaches the provider" + + +# ── guards and binding ─────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_refuses_to_start_when_a_tool_approval_policy_exists(monkeypatch): + """No approval interrupt exists on this path yet: fail closed, before any model call.""" + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", True, raising=False) + model = _Model(AIMessage(content="", tool_calls=[_CALL])) + + with patch.object(AgentGraphAdapter, "_tool_approval_policies_exist", new=AsyncMock(return_value=True)): + cmd = await _turn(_adapter(), model, _state(), FC) + assert cmd.goto == END and cmd.update["error"] == FC_TOOL_APPROVAL_UNSUPPORTED + assert model.seen == [], "refused before the model was invoked" + + with patch.object(AgentGraphAdapter, "_tool_approval_policies_exist", new=AsyncMock(return_value=False)): + cmd = await _turn(_adapter(), model, _state(), FC) + assert cmd.goto == "tool_exec" and len(model.seen) == 1, "policies without approval rules do not block" + + @pytest.mark.asyncio async def test_bind_mode_is_upgraded_from_none_only_in_function_calling_mode(): captured = [] @@ -202,3 +336,22 @@ async def fake_resolve(active_model, **kwargs): assert fc_default["cuga_lite_bind_tools_mode"] == "all", "FC is inert without advertised tools" assert fc_explicit["cuga_lite_bind_tools_mode"] == "find_tools", "an explicit choice is respected" assert "cuga_lite_bind_tools_mode" not in codeact, "codeact is byte-identical" + + +@pytest.mark.asyncio +async def test_bind_advertises_the_executable_set_prepare_recorded(): + captured = [] + + async def fake_resolve(active_model, **kwargs): + captured.append(kwargs["configurable"]) + return "BOUND" + + with patch.object(ga, "resolve_model_with_bind_tools", new=AsyncMock(side_effect=fake_resolve)): + scoped = _adapter(tools_context_ref={"_lc_bind_tools_executable_names": ["echo", "add"]}) + await scoped.resolve_bind_tools(_state(), object(), dict(FC), None) + + (cfg,) = captured + assert cfg["cuga_lite_bind_tools_mode"] == "tools" + assert cfg["cuga_lite_bind_tools_tool_names"] == ["echo", "add"], ( + "bind what the sandbox could call, not the catalogue" + ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index 7cc3d51a6..df60ae2cc 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -264,3 +264,32 @@ async def forbidden(state, config=None): "CodeAct still serialises turns through the dict path" ) assert "```python" in result["prepared_prompt"], "the CodeAct prompt is unchanged" + + +@pytest.mark.asyncio +async def test_bundled_codeact_few_shots_are_withheld_in_fc_mode(): + """With find_tools active the CodeAct path replays bundled ```python demos; in + function-calling mode they would contradict the prompt, so only demos passed + explicitly are sent.""" + codeact = _ScriptedModel([AIMessage(content="done")]) + await _run(codeact, _config("fs-codeact", shortlisting_tool_threshold=0), _echo_tool()) + assert any("```python" in (m.get("content") or "") for m in codeact.seen[0] if isinstance(m, dict)), ( + "precondition: the bundled CodeAct demos are in play" + ) + + fc = _ScriptedModel([AIMessage(content="done")]) + await _run( + fc, + _config("fs-fc", cuga_lite_execution_mode="function_calling", shortlisting_tool_threshold=0), + _echo_tool(), + ) + assert not any("```python" in getattr(m, "content", "") for m in fc.seen[0]), "CodeAct demos withheld" + + explicit = _ScriptedModel([AIMessage(content="done")]) + demos = [{"role": "user", "content": "demo q"}, {"role": "assistant", "content": "demo a"}] + await _run( + explicit, + _config("fs-explicit", cuga_lite_execution_mode="function_calling", mcp_few_shot_examples=demos), + _echo_tool(), + ) + assert [m.content for m in explicit.seen[0][1:3]] == ["demo q", "demo a"], "explicit demos still replayed" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py index 4f0f06260..3bfa987f5 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py @@ -51,9 +51,20 @@ def test_instructions_then_special_instructions_close_the_prompt(): assert "\n\n\n" not in text, "no blank-line runs from unset optional sections" -def test_template_ships_no_benchmark_or_deployment_wording(): +def test_template_ships_no_placeholder_or_deployment_wording(): + """Answer-format contracts belong in special_instructions, never in the shipped template.""" text = _TEMPLATE.read_text() - assert "[[" not in text - assert "vakra" not in text.lower() - for module in ("adapter/tool_exec_node.py", "adapter/graph_adapter.py", "prompt_utils.py"): - assert "vakra" not in (_TEMPLATE.parents[1] / module).read_text().lower(), module + assert "[[" not in text and "]]" not in text + assert "{{ special_instructions }}" in text and "{{ instructions }}" in text + + +def test_skill_and_agent_catalogues_are_rendered_when_given(): + """``load_skill`` / ``delegate_to_*`` take a name; the model needs the catalogue to pick one.""" + text = render_fc_prompt( + skills_prompt_section="## Available skills\n- deploy: ship it", + agents_prompt_section="## Available agents\n- researcher", + instructions="Be brief.", + ) + assert "- deploy: ship it" in text and "- researcher" in text + assert text.index("- deploy") < text.index("- researcher") < text.index("Be brief.") + assert "Available skills" not in render_fc_prompt(), "absent when nothing is bound" diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py index 2a23fe0a6..ab854d52d 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py @@ -65,8 +65,14 @@ def _call(name: str, args: dict, call_id: str) -> dict: return {"name": name, "args": args, "id": call_id, "type": "tool_call"} +def _update(result) -> dict: + """Normal exits are ``Command(goto="call_model")``; error exits ``Command(goto=END)``.""" + assert isinstance(result, Command), type(result) + return result.update + + def _tool_messages(result) -> List[ToolMessage]: - return [m for m in result["chat_messages"] if isinstance(m, ToolMessage)] + return [m for m in _update(result)["chat_messages"] if isinstance(m, ToolMessage)] @pytest.fixture(autouse=True) @@ -111,12 +117,12 @@ def greet(name: str) -> str: # sync tools are made awaitable, like the sandbox ("c1", "add", "3", "success"), ("c2", "greet", "hi bob", "success"), ] - assert result["chat_messages"][:2] == _state(last).chat_messages, ( + assert _update(result)["chat_messages"][:2] == _state(last).chat_messages, ( "history is appended to, never rewritten" ) - assert result["script"] is None - assert result["step_count"] == 1 - assert result["tool_calls_used_run"] == 2 and result["tool_budget_exhausted"] is False + assert _update(result)["script"] is None + assert _update(result)["step_count"] == 1 + assert _update(result)["tool_calls_used_run"] == 2 and _update(result)["tool_budget_exhausted"] is False @pytest.mark.asyncio @@ -202,7 +208,7 @@ async def test_timeout_is_an_error_toolmessage(monkeypatch): _caps(monkeypatch) from cuga.config import settings - monkeypatch.setattr(settings.advanced_features, "tool_call_timeout", 0.05, raising=False) + monkeypatch.setattr(settings.advanced_features, "sandbox_execution_timeout", 0.05, raising=False) async def slow() -> str: await asyncio.Event().wait() # never set: immune to any asyncio.sleep patch elsewhere in the suite @@ -267,8 +273,10 @@ async def add(a: int, b: int) -> int: m1, m2 = _tool_messages(result) assert m1.content == "2" and calls == [(1, 1)], "the second call must be refused, not executed" assert m2.status == "error" and "Tool call limit reached" in m2.content - assert result["tool_calls_used_run"] == 1, "a refused call never inflates the counter" - assert result["tool_budget_exhausted"] is True, "call_model must see the exhaustion and end the turn" + assert _update(result)["tool_calls_used_run"] == 1, "a refused call never inflates the counter" + assert _update(result)["tool_budget_exhausted"] is True, ( + "call_model must see the exhaustion and end the turn" + ) @pytest.mark.asyncio @@ -285,11 +293,13 @@ async def add(a: int, b: int) -> int: result = await node(_state(last), config={"configurable": {"track_tool_calls": True}}) - assert [(c["name"], c["app_name"], c["result"]) for c in result["tool_calls"]] == [("add", "calc", 7)] + assert [(c["name"], c["app_name"], c["result"]) for c in _update(result)["tool_calls"]] == [ + ("add", "calc", 7) + ] assert ToolCallTracker.is_enabled() is False untracked = await node(_state(last), config={"configurable": {}}) - assert untracked["tool_calls"] == [] + assert _update(untracked)["tool_calls"] == [] # ── step discipline ────────────────────────────────────────────────────────── @@ -323,7 +333,7 @@ async def add(a: int, b: int) -> int: assert (m1.tool_call_id, m1.content, m1.status) == ("c1", "2", "success") assert (m2.tool_call_id, m2.content, m2.status) == ("c2", DEFERRED_CALL_MESSAGE, "error") assert (m3.tool_call_id, m3.content, m3.status) == ("c3", DEFERRED_CALL_MESSAGE, "error") - assert result["tool_calls_used_run"] == 1 + assert _update(result)["tool_calls_used_run"] == 1 @pytest.mark.asyncio @@ -355,8 +365,9 @@ async def test_nothing_executable_still_advances_the_step(monkeypatch): result = await node(_state(AIMessage(content="no calls here"), step_count=4), config=None) - assert result["step_count"] == 5 and result["script"] is None - assert "chat_messages" not in result, "no ToolMessage is fabricated" + assert result.goto == "call_model" + assert _update(result)["step_count"] == 5 and _update(result)["script"] is None + assert "chat_messages" not in _update(result), "no ToolMessage is fabricated" @pytest.mark.asyncio @@ -375,3 +386,102 @@ async def add(a: int, b: int) -> int: assert "Maximum step limit" in result.update["error"] assert result.update["execution_complete"] is True assert "tool_calls_used_run" in result.update, "budget fields ride every exit" + + +# ── outer guard + batch output budget ──────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_unexpected_failure_after_the_batch_still_reports_the_budget(monkeypatch): + """Same contract as the sandbox: whatever breaks after a call ran, the update + carries the spent budget, or the checkpoint keeps the pre-batch counts.""" + _caps(monkeypatch) + from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import tool_exec_node as mod + + async def add(a: int, b: int) -> int: + return a + b + + def boom(*args, **kwargs): + raise RuntimeError("history store unavailable") + + monkeypatch.setattr(mod, "core_append_with_step_limit", boom) + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) + + result = await node(_state(last), config=None) + + assert isinstance(result, Command) + assert "history store unavailable" in result.update["error"] + assert result.update["tool_calls_used_run"] == 1, "the call that ran must still be counted" + assert ToolCallTracker.is_enabled() is False + + +@pytest.mark.asyncio +async def test_one_batch_shares_one_output_limit(monkeypatch): + """N results share execution_output_max_length, like the N prints of one block.""" + _caps(monkeypatch) + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "execution_output_max_length", 10, raising=False) + + async def text(n: int) -> str: + return "x" * n + + node = create_tool_exec_node(_Adapter({"text": text})) + last = AIMessage( + content="", + tool_calls=[ + _call("text", {"n": 8}, "c1"), + _call("text", {"n": 8}, "c2"), + _call("text", {"n": 3}, "c3"), + ], + ) + + result = await node(_state(last), config=None) + + m1, m2, m3 = _tool_messages(result) + assert m1.content == "x" * 8, "under the limit: untouched" + assert m2.content.startswith("xx") and "truncated" in m2.content, "only the remaining budget is spent" + assert "truncated" in m3.content and not m3.content.startswith("x"), "nothing left for the third result" + + +@pytest.mark.asyncio +async def test_normal_exit_routes_to_call_model_and_error_exit_to_end(monkeypatch): + """One routing mechanism: no static edge, so a terminal error never schedules + one more model turn.""" + _caps(monkeypatch) + + async def add(a: int, b: int) -> int: + return a + b + + node = create_tool_exec_node(_Adapter({"add": add}, max_steps=3)) + last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) + + assert (await node(_state(last), config=None)).goto == "call_model" + from langgraph.graph import END + + assert (await node(_state(last, step_count=3), config=None)).goto == END + + +@pytest.mark.asyncio +async def test_step_discipline_defers_by_position_even_when_the_first_call_fails(monkeypatch): + """An error on the first call is the result the model must read before the next + call — the second call must not run just because the first one failed.""" + _caps(monkeypatch) + calls = [] + + async def add(a: int, b: int) -> int: + calls.append((a, b)) + return a + b + + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage(content="", tool_calls=[_call("nope", {}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")]) + + result = await node( + _state(last), config={"configurable": {"cuga_lite_step_discipline": "one_tool_per_step"}} + ) + + m1, m2 = _tool_messages(result) + assert m1.status == "error" and "Unknown tool" in m1.content + assert (m2.tool_call_id, m2.content, m2.status) == ("c2", DEFERRED_CALL_MESSAGE, "error") + assert calls == [] diff --git a/src/cuga/backend/cuga_graph/utils/agent_loop.py b/src/cuga/backend/cuga_graph/utils/agent_loop.py index 705415751..b22ff9001 100644 --- a/src/cuga/backend/cuga_graph/utils/agent_loop.py +++ b/src/cuga/backend/cuga_graph/utils/agent_loop.py @@ -28,6 +28,28 @@ from cuga.backend.observability.openlit_init import set_session_attribute +def _native_tool_calls(messages: Any) -> list: + """``tool_calls`` on the last message of a function-calling turn (object or dict), else [].""" + if not messages: + return [] + last = messages[-1] + calls = last.get("tool_calls") if isinstance(last, dict) else getattr(last, "tool_calls", None) + return list(calls or []) + + +def _tool_exec_step_output(messages: Any) -> str: + """Render the trailing run of tool results (one per native call) as execution output.""" + lines: list = [] + for msg in reversed(messages or []): + role = msg.get("type") if isinstance(msg, dict) else getattr(msg, "type", None) + if role != "tool": + break + name = (msg.get("name") if isinstance(msg, dict) else getattr(msg, "name", None)) or "tool" + content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "") + lines.append(f"{name}: {content}") + return "\n".join(reversed(lines)) + + class OutputFormat(str, Enum): WXO = "wxo" DEFAULT = "default" @@ -394,6 +416,21 @@ def get_event_message(self, event) -> StreamEvent: "variables": state_data.get("variables_storage", {}), } return StreamEvent(name="CodeAgent", data=json.dumps(output)) + elif _native_tool_calls(subgraph_messages): + # Function-calling mode: show the native calls the way a code block is shown. + calls_text = "\n".join( + f"{c.get('name')}({json.dumps(c.get('args') or {}, ensure_ascii=False, default=str)})" + for c in _native_tool_calls(subgraph_messages) + if isinstance(c, dict) + ) + output = { + "code": calls_text, + "execution_output": "", + "steps_summary": [], + "summary": "Tool calls issued, preparing to execute", + "variables": state_data.get("variables_storage", {}), + } + return StreamEvent(name="CodeAgent", data=json.dumps(output)) else: # Text/reasoning output - only when last chat turn is a non-empty assistant message logger.info("call_model generated text response (no code)") @@ -445,6 +482,21 @@ def get_event_message(self, event) -> StreamEvent: logger.debug("Skipping empty sandbox event") return StreamEvent(name="", data="") + # Function-calling mode: the step appended one ToolMessage per native call. + elif node_name == "tool_exec": + tool_output = _tool_exec_step_output(subgraph_messages or []) + if tool_output.strip(): + output = { + "code": "", + "execution_output": tool_output, + "steps_summary": [], + "summary": "Tool calls completed", + "variables": state_data.get("variables_storage", {}), + } + return StreamEvent(name="CodeAgent", data=json.dumps(output)) + logger.debug("Skipping empty tool_exec event") + return StreamEvent(name="", data="") + # Default handling for other subgraph nodes logger.debug(f"Unhandled subgraph node: {node_name}") return StreamEvent(name=str(node_name), data="") diff --git a/tests/unit/test_agent_loop_fc_events.py b/tests/unit/test_agent_loop_fc_events.py new file mode 100644 index 000000000..60c6308f7 --- /dev/null +++ b/tests/unit/test_agent_loop_fc_events.py @@ -0,0 +1,51 @@ +"""Function-calling steps stream to the UI like code steps. + +``get_event_message`` maps subgraph nodes to StreamEvents; without these branches a +native tool-call turn and its results were an empty ``tool_exec`` event and the UI +showed nothing between the question and the answer. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from cuga.backend.cuga_graph.utils.agent_loop import AgentLoop + +pytestmark = pytest.mark.unit + +_CALL = {"name": "echo", "args": {"value": 7}, "id": "c1", "type": "tool_call"} + + +def _event(node, messages): + return ( + ("CugaLiteSubgraph",), + {node: {"chat_messages": messages, "script": None, "variables_storage": {}}}, + ) + + +def test_native_tool_calls_stream_as_a_code_step(): + messages = [HumanMessage(content="q"), AIMessage(content="", tool_calls=[_CALL])] + ev = AgentLoop.get_event_message(MagicMock(), _event("call_model", messages)) + assert ev.name == "CodeAgent" + assert json.loads(ev.data)["code"] == 'echo({"value": 7})' + + +def test_tool_results_stream_as_execution_output(): + messages = [ + HumanMessage(content="q"), + AIMessage(content="", tool_calls=[_CALL, {**_CALL, "id": "c2", "args": {"value": 8}}]), + ToolMessage(content="7", tool_call_id="c1", name="echo"), + ToolMessage(content="8", tool_call_id="c2", name="echo"), + ] + ev = AgentLoop.get_event_message(MagicMock(), _event("tool_exec", messages)) + assert ev.name == "CodeAgent" + assert json.loads(ev.data)["execution_output"] == "echo: 7\necho: 8" + + +def test_tool_exec_without_results_is_an_empty_event(): + ev = AgentLoop.get_event_message(MagicMock(), _event("tool_exec", [HumanMessage(content="q")])) + assert ev.name == "" and ev.data == "" diff --git a/tests/unit/test_fc_sdk_two_turn.py b/tests/unit/test_fc_sdk_two_turn.py new file mode 100644 index 000000000..84e4dc91f --- /dev/null +++ b/tests/unit/test_fc_sdk_two_turn.py @@ -0,0 +1,104 @@ +"""Function-calling mode across two SDK turns on one thread. + +The SDK wrapper hands state between the CugaLite subgraph and the outer +``AgentState`` graph through ``state.model_dump()``. ``AgentState.chat_messages`` +is declared ``List[BaseMessage]``, so pydantic serialises every message by the +*declared* type: ``AIMessage.tool_calls`` and ``ToolMessage.tool_call_id`` are +dropped and revalidation yields bare ``BaseMessage`` objects. CodeAct never +notices (its serializer keys on ``msg.type``); function-calling replays history +as message objects, so turn 2 must still send the provider a valid transcript. +""" + +from __future__ import annotations + + +import pytest +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.tools import StructuredTool + +pytestmark = pytest.mark.unit + +CALLS: list = [] + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + self.seen: list = [] + + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages, config=None, **kwargs): + self.seen.append(list(messages)) + if not self._responses: + raise AssertionError("model asked for more responses than scripted") + return self._responses.pop(0) + + +def _echo(): + async def echo(value: int) -> int: + """Echo a value.""" + CALLS.append(value) + return value + + return StructuredTool.from_function(coroutine=echo, name="echo", description="Echo a value.") + + +def _tc(value, call_id): + return {"name": "echo", "args": {"value": value}, "id": call_id, "type": "tool_call"} + + +@pytest.fixture(autouse=True) +def _quiet(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + CALLS.clear() + yield + CALLS.clear() + + +def _describe(m): + kind = type(m).__name__ + extra = "" + if getattr(m, "tool_calls", None): + extra = f" tool_calls={[c['id'] for c in m.tool_calls]}" + if getattr(m, "tool_call_id", None): + extra = f" tool_call_id={m.tool_call_id}" + return f"{kind}({getattr(m, 'type', '?')}){extra}" + + +@pytest.mark.asyncio +async def test_second_turn_replays_a_valid_function_calling_transcript(): + from cuga.sdk import CugaAgent + + model = _ScriptedModel( + [ + AIMessage(content="", tool_calls=[_tc(7, "call_1")]), + AIMessage(content="The value is 7."), + AIMessage(content="", tool_calls=[_tc(8, "call_2")]), + AIMessage(content="And now 8."), + ] + ) + agent = CugaAgent(tools=[_echo()], model=model, execution_mode="function_calling") + + turn1 = await agent.invoke("echo 7", thread_id="fc-two-turn") + assert turn1.answer == "The value is 7." and CALLS == [7] + + turn2 = await agent.invoke("now echo 8", thread_id="fc-two-turn") + assert turn2.answer == "And now 8." and CALLS == [7, 8] + + outbound = model.seen[2] # first model call of turn 2 + shapes = [_describe(m) for m in outbound] + print("\nTURN 2 outbound:", shapes) + + assert isinstance(outbound[0], SystemMessage) + assert all(type(m) is not BaseMessage for m in outbound), ( + f"bare BaseMessage replayed to the provider: {shapes}" + ) + ai_with_calls = [m for m in outbound if isinstance(m, AIMessage) and m.tool_calls] + tool_replies = [m for m in outbound if isinstance(m, ToolMessage)] + dangling = {c["id"] for m in ai_with_calls for c in m.tool_calls} - {t.tool_call_id for t in tool_replies} + assert not dangling, f"tool_calls without a ToolMessage reply: {dangling} in {shapes}" + assert isinstance(outbound[-1], HumanMessage) and outbound[-1].content == "now echo 8" From 948350b92c1d45f20593a7341953e195e0773e72 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:43:13 +0300 Subject: [PATCH 04/18] fix(cuga-lite): app catalogue in the FC prompt, real-path fail-closed test, wording - find_tools(query, app_name) requires an app name: when it is bound in function-calling mode the prompt now lists the available apps. - Fail-closed approval guard covered through the real SDK path: a tool-approval policy added via agent.policies refuses the FC turn before any model call, while CodeAct on the same agent still reaches the approval prompt. - Step discipline wording aligned everywhere: the first call is attempted (success or error), the rest are deferred. - SDK tests disable the Evolve registry round-trips. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 2 +- .../nodes/cuga_lite/adapter/prepare_node.py | 1 + .../nodes/cuga_lite/prompt_utils.py | 5 +- .../nodes/cuga_lite/prompts/fc_prompt.jinja2 | 6 +- .../nodes/cuga_lite/tests/test_fc_prompt.py | 11 +++ src/cuga/settings.toml | 2 +- .../unit/test_fc_sdk_fail_closed_approval.py | 93 +++++++++++++++++++ tests/unit/test_fc_sdk_two_turn.py | 1 + 8 files changed, 117 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_fc_sdk_fail_closed_approval.py diff --git a/README.md b/README.md index f44ab4401..1c109aa49 100644 --- a/README.md +++ b/README.md @@ -1180,7 +1180,7 @@ See `docs/design/pluggable-shortlister.md` for the full design. By default CUGA Lite works in **CodeAct**: the model writes a Python block and the sandbox runs it. `cuga_lite_execution_mode = "function_calling"` switches to native tool calls: the model emits `tool_calls`, a `tool_exec` node runs them under the same budgets, tracker and timeout as the sandbox and replies with `ToolMessage`s, and the model reads the results before it calls again or answers. Tools are advertised through `bind_tools`; a short dedicated prompt replaces the CodeAct prompt. -**Step discipline** (`cuga_lite_step_discipline = "one_tool_per_step"`) works in both modes: only the first tool call of a step runs, so the model reads a result before it decides the next call. CodeAct keeps the block's variables (including that first result); function-calling answers the extra calls with a "deferred" reply. +**Step discipline** (`cuga_lite_step_discipline = "one_tool_per_step"`) works in both modes: only the first tool call of a step is attempted — success or error, that is the result the model reads before it decides the next call. CodeAct keeps the block's variables (including that first result); function-calling answers the extra calls with a "deferred" reply. ## Configuration diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index 8193dfb1c..a439c47f1 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -778,6 +778,7 @@ async def _wrapped(*args, **kwargs): fragments=_fragments, skills_prompt_section=skills_prompt_section if skills_enabled else "", agents_prompt_section=agents_prompt_section if agents_enabled else "", + apps=apps_for_prompt if enable_find_tools else None, ) logger.info( "Prepared CugaLite function-calling prompt: step_discipline={} fragments={} prompt_chars={}", diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py index 01a9b2ed0..922a94f69 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompt_utils.py @@ -923,6 +923,7 @@ def render_fc_prompt( fragments: Optional[List[str]] = None, skills_prompt_section: Optional[str] = None, agents_prompt_section: Optional[str] = None, + apps: Optional[list] = None, ) -> str: """System prompt for ``cuga_lite_execution_mode = "function_calling"``. @@ -934,7 +935,8 @@ def render_fc_prompt( - the generic preamble (always), - the step-discipline line when ``cuga_lite_step_discipline`` is on, - opt-in fragments from ``cuga_lite_fc_prompt_fragments`` (``evidence_first``), - - the skill and agent catalogues when those tools are bound (``load_skill`` / + - the app catalogue when ``find_tools`` is bound (it takes an ``app_name``), and + the skill / agent catalogues when those tools are bound (``load_skill`` / ``delegate_to_*`` take a name; the model has to be able to pick a valid one), - then the agent's ``instructions`` and the caller's ``special_instructions``. @@ -947,6 +949,7 @@ def render_fc_prompt( is_autonomous_subtask=bool(is_autonomous_subtask), step_discipline=bool(step_discipline), evidence_first="evidence_first" in frags, + apps=format_apps_for_prompt(apps) if apps else [], skills_prompt_section=(skills_prompt_section or "").strip(), agents_prompt_section=(agents_prompt_section or "").strip(), instructions=(instructions or "").strip(), diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 index ddf99b083..4b99e9ab4 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/prompts/fc_prompt.jinja2 @@ -15,7 +15,11 @@ You are Cuga Agent, a helpful assistant that answers the user's request by calli # FINAL ANSWER - When your tool results contain everything needed, STOP calling tools and give the final answer with no further tool call. -{% if skills_prompt_section %} +{% if apps %} +# AVAILABLE APPS +Use the exact app name as `app_name` when calling `find_tools`. +{% for app in apps %}- {{ app.name }} ({{ app.type }}): {{ app.description }} +{% endfor %}{% endif %}{% if skills_prompt_section %} {{ skills_prompt_section }} {% endif %}{% if agents_prompt_section %} {{ agents_prompt_section }} diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py index 3bfa987f5..143e042e3 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_prompt.py @@ -68,3 +68,14 @@ def test_skill_and_agent_catalogues_are_rendered_when_given(): assert "- deploy: ship it" in text and "- researcher" in text assert text.index("- deploy") < text.index("- researcher") < text.index("Be brief.") assert "Available skills" not in render_fc_prompt(), "absent when nothing is bound" + + +def test_app_catalogue_is_rendered_when_find_tools_is_bound(): + """``find_tools(query, app_name)`` requires an app name; the model needs the list.""" + from types import SimpleNamespace + + apps = [SimpleNamespace(name="crm", type="api", description="Customers and deals")] + text = render_fc_prompt(apps=apps) + assert "# AVAILABLE APPS" in text and "- crm (api): Customers and deals" in text + assert "`app_name`" in text + assert "AVAILABLE APPS" not in render_fc_prompt(), "absent when find_tools is not bound" diff --git a/src/cuga/settings.toml b/src/cuga/settings.toml index 3e05fdc5a..bd6831789 100644 --- a/src/cuga/settings.toml +++ b/src/cuga/settings.toml @@ -60,7 +60,7 @@ cuga_lite_bind_tools_pad_to_cap = false cuga_lite_execution_mode = "codeact" # Step discipline, both modes: off (default) | one_tool_per_step — observe one result before the next call. # codeact: only the first tool call in a generated block runs (per-block cap becomes 1). function_calling: only -# the first tool_call of a turn runs; the rest are answered with a "deferred" ToolMessage. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_STEP_DISCIPLINE +# the first tool_call of a turn is attempted; the rest are answered with a "deferred" ToolMessage. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_STEP_DISCIPLINE cuga_lite_step_discipline = "off" # Opt-in prompt fragments for function_calling mode. Known: "evidence_first" (call a tool before answering; never # answer from prior knowledge). Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_FC_PROMPT_FRAGMENTS='["evidence_first"]' diff --git a/tests/unit/test_fc_sdk_fail_closed_approval.py b/tests/unit/test_fc_sdk_fail_closed_approval.py new file mode 100644 index 000000000..8d16dbddc --- /dev/null +++ b/tests/unit/test_fc_sdk_fail_closed_approval.py @@ -0,0 +1,93 @@ +"""Function-calling mode refuses to start while a tool-approval policy is stored. + +Through the real SDK path: a policy added with ``agent.policies.add_tool_approval`` +lands in policy storage, and the function-calling turn queries that storage before +its first model call. CodeAct on the same agent still goes through the approval +prompt, which is the behaviour the refusal protects. +""" + +from __future__ import annotations + +import pytest +from langchain_core.messages import AIMessage +from langchain_core.tools import StructuredTool + +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import FC_TOOL_APPROVAL_UNSUPPORTED + +pytestmark = pytest.mark.unit + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + self.invocations = 0 + + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages, config=None, **kwargs): + self.invocations += 1 + if not self._responses: + raise AssertionError("model asked for more responses than scripted") + return self._responses.pop(0) + + +def _delete_record(): + async def delete_record(record_id: int) -> str: + """Delete a record.""" + raise AssertionError("a guarded tool ran without approval") + + return StructuredTool.from_function( + coroutine=delete_record, name="delete_record", description="Delete a record." + ) + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", True, raising=False) + monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) + + +@pytest.mark.asyncio +async def test_fc_refuses_with_a_stored_tool_approval_policy_and_codeact_still_asks(): + from cuga.sdk import CugaAgent + + fc_model = _ScriptedModel( + [ + AIMessage( + content="", + tool_calls=[ + {"name": "delete_record", "args": {"record_id": 1}, "id": "c1", "type": "tool_call"} + ], + ) + ] + ) + agent = CugaAgent( + tools=[_delete_record()], model=fc_model, execution_mode="function_calling", reset_policy_storage=True + ) + await agent.policies.add_tool_approval( + name="Approve deletions", + required_tools=["delete_record"], + approval_message="Deleting needs approval.", + ) + + result = await agent.invoke("delete record 1", thread_id="fc-approval") + + assert fc_model.invocations == 0, "refused before the model was invoked" + assert ( + FC_TOOL_APPROVAL_UNSUPPORTED in (result.error or "") or FC_TOOL_APPROVAL_UNSUPPORTED in result.answer + ) + + # Same agent, CodeAct for this call: the guarded code goes to the approval prompt. + codeact_model = _ScriptedModel( + [AIMessage(content="```python\nprint(await delete_record(record_id=1))\n```")] + ) + agent._model = codeact_model + agent._graph = None + agent._compiled_graph = None + result = await agent.invoke("delete record 1", thread_id="codeact-approval", execution_mode="codeact") + + assert codeact_model.invocations == 1 + assert "approval" in result.answer.lower(), result.answer diff --git a/tests/unit/test_fc_sdk_two_turn.py b/tests/unit/test_fc_sdk_two_turn.py index 84e4dc91f..0aeaba1cb 100644 --- a/tests/unit/test_fc_sdk_two_turn.py +++ b/tests/unit/test_fc_sdk_two_turn.py @@ -54,6 +54,7 @@ def _quiet(monkeypatch): from cuga.config import settings monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) # no registry round-trips CALLS.clear() yield CALLS.clear() From a66aca4dd0db935f591ef35fdfe4b8bf7f7ded2b Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 15 Sep 2026 13:05:35 +0300 Subject: [PATCH 05/18] fix(cuga-lite): fail the function-calling approval guard closed on verification errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guard returned False when the policy query raised, so a storage failure let function-calling proceed. It is a safety check and a refusal is cheap, so it is now stricter than CodeAct's own check: any failure to verify refuses the turn with a distinct message. PolicyStorage.list_policies swallows backend errors into an empty list, which the guard could not tell apart from "no policies" — it takes a strict=True variant now that re-raises them, so both the initialisation and the backend failure paths are closed. Also from CodeRabbit: the function-calling prompt advertises the app catalogue find_tools was actually built with (runtime apps appended later are not in its app_to_tools_map), and the real-path approval test now proves the guarded tool never ran and asserts the policy's own approval message. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 2 +- .../nodes/cuga_lite/adapter/graph_adapter.py | 56 +++++++++++-------- .../nodes/cuga_lite/adapter/prepare_node.py | 6 +- .../tests/test_fc_call_model_hook.py | 39 +++++++++++++ src/cuga/backend/cuga_graph/policy/storage.py | 10 ++++ .../policy/tests/test_storage_strict_list.py | 34 +++++++++++ .../unit/test_fc_sdk_fail_closed_approval.py | 12 +++- 7 files changed, 132 insertions(+), 27 deletions(-) create mode 100644 src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py diff --git a/README.md b/README.md index 1c109aa49..ace7044a5 100644 --- a/README.md +++ b/README.md @@ -1200,7 +1200,7 @@ result = await agent.invoke("...", execution_mode="codeact") # per-call overri ## Gotchas -- **Tool-approval policies**: function-calling mode has no approval interrupt yet. If an enabled tool-approval policy exists, the run stops with a clear error before any tool runs — use `codeact` for that agent or disable the policy. +- **Tool-approval policies**: function-calling mode has no approval interrupt yet. If an enabled tool-approval policy exists — or policy storage cannot be checked — the run stops with a clear error before any tool runs (fail closed); use `codeact` for that agent or disable the policy. - Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call. Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. - The bundled CodeAct few-shot demos are not sent in function-calling mode; pass your own through `configurable["mcp_few_shot_examples"]`. - Tools run in-process and sequentially. The pre-execute VERIFY gate, reflection and the E2B / OpenSandbox executors are CodeAct-only; step discipline caps local-executor blocks only. diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index f56c35c21..104c6bba1 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -70,6 +70,11 @@ '"codeact" for this agent, or disable the policy.' ) +FC_TOOL_APPROVAL_UNVERIFIED = ( + "Function-calling mode could not verify whether a tool-approval policy exists ({error}). " + "It fails closed: no tool ran. Fix policy storage, or use cuga_lite_execution_mode = \"codeact\"." +) + FC_STEP_LIMIT_CALL_REPLY = "Not executed: the step limit was reached before this call could run." FC_BUDGET_CALL_REPLY = ( "Not executed: the tool budget for this turn is spent. Answer from the data already retrieved." @@ -455,25 +460,23 @@ def _resolved_bind_mode(self, configurable: dict) -> str: return _bind_tools_mode_from_settings() async def _tool_approval_policies_exist(self, config: Any) -> bool: - """True when an enabled tool-approval policy is configured. + """True when an enabled tool-approval policy is stored. Function-calling has no approval interrupt yet — CodeAct's runs on the generated code string, after the seam — so the mode refuses to start - rather than run a guarded tool unprompted. On a storage error this - warns and proceeds, exactly as CodeAct's own approval check does. + rather than run a guarded tool unprompted. Raises when the answer + cannot be established (policy system not initialised, storage backend + down): ``strict=True`` keeps the storage layer from turning a backend + failure into an empty list, and the caller fails closed on any error. """ from cuga.backend.cuga_graph.policy.configurable import PolicyConfigurable from cuga.backend.cuga_graph.policy.models import PolicyType - try: - policy_system = PolicyConfigurable.from_config(config or {}) - policies = await policy_system.agent.storage.list_policies( - policy_type=PolicyType.TOOL_APPROVAL, enabled_only=True, limit=1 - ) - return bool(policies) - except Exception as exc: - logger.warning("[fc] could not query tool-approval policies ({}); proceeding", exc) - return False + policy_system = PolicyConfigurable.from_config(config or {}) + policies = await policy_system.agent.storage.list_policies( + policy_type=PolicyType.TOOL_APPROVAL, enabled_only=True, limit=1, strict=True + ) + return bool(policies) async def execute_call_model_fc( self, @@ -504,9 +507,9 @@ async def execute_call_model_fc( - a fenced code block with no ``tool_calls`` is a mode violation: it is never executed; the model gets one corrective turn instead. - Refuses to start when an enabled tool-approval policy exists: there is - no approval interrupt on this path yet, and silently running a guarded - tool is worse than not running at all. + Refuses to start when an enabled tool-approval policy exists — or when + that cannot be verified: there is no approval interrupt on this path + yet, and silently running a guarded tool is worse than not running. """ if self._execution_mode(configurable) != EXECUTION_MODE_FUNCTION_CALLING: return None @@ -518,14 +521,21 @@ async def execute_call_model_fc( cfg = configurable or {} history: list = list(modified_messages) - if settings.policy.enabled and await self._tool_approval_policies_exist(config): - logger.error( - "[fc] refusing to start: an enabled tool-approval policy exists and " - "function-calling mode has no approval interrupt" - ) - return create_error_command( - self, history, AIMessage(content=FC_TOOL_APPROVAL_UNSUPPORTED), state.step_count - ) + if settings.policy.enabled: + # Fail closed, stricter than CodeAct's own check: a refusal here is + # cheap and explicit, a guarded tool running unprompted is not. + try: + blocked_reason = ( + FC_TOOL_APPROVAL_UNSUPPORTED if await self._tool_approval_policies_exist(config) else None + ) + except Exception as exc: + logger.error("[fc] could not verify tool-approval policies ({}); refusing to start", exc) + blocked_reason = FC_TOOL_APPROVAL_UNVERIFIED.format(error=exc) + if blocked_reason: + logger.error("[fc] refusing to start: {}", blocked_reason) + return create_error_command( + self, history, AIMessage(content=blocked_reason), state.step_count + ) msgs: List[BaseMessage] = [SystemMessage(content=system_content)] msgs.extend(_few_shot_to_messages(self.get_few_shot_messages(state))) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index a439c47f1..2fb56bd0b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -488,6 +488,10 @@ async def prepare_tools_and_apps(state: Any, config: Optional[RunnableConfig] = _runtime_bundle = build_runtime_tools(thread_id=runtime_thread_id, backends=_runtime_backends) adapter._tools_context.update(_runtime_bundle.execution_callables) tools_for_prompt.extend(_runtime_bundle.prompt_tools) + # The catalogue find_tools was built with (above). Runtime apps appended + # below are not in its app_to_tools_map, so advertising them as + # ``app_name`` targets would only produce empty searches. + _find_tools_apps = list(apps_for_prompt) if apps_for_prompt is not None else None if _runtime_bundle.app_definitions and apps_for_prompt is not None: apps_for_prompt = list(apps_for_prompt) + _runtime_bundle.app_definitions @@ -778,7 +782,7 @@ async def _wrapped(*args, **kwargs): fragments=_fragments, skills_prompt_section=skills_prompt_section if skills_enabled else "", agents_prompt_section=agents_prompt_section if agents_enabled else "", - apps=apps_for_prompt if enable_find_tools else None, + apps=_find_tools_apps if enable_find_tools else None, ) logger.info( "Prepared CugaLite function-calling prompt: step_discipline={} fragments={} prompt_chars={}", diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index 03bad6079..18c27ab0e 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -30,6 +30,7 @@ FC_MODE_VIOLATION_CORRECTION, FC_STEP_LIMIT_CALL_REPLY, FC_TOOL_APPROVAL_UNSUPPORTED, + FC_TOOL_APPROVAL_UNVERIFIED, FC_UNANSWERED_CALL_REPLY, AgentGraphAdapter, _normalize_history_for_replay, @@ -355,3 +356,41 @@ async def fake_resolve(active_model, **kwargs): assert cfg["cuga_lite_bind_tools_tool_names"] == ["echo", "add"], ( "bind what the sandbox could call, not the catalogue" ) + + +@pytest.mark.asyncio +async def test_refuses_when_the_policy_query_fails(monkeypatch): + """Fail closed on infrastructure failure too: 'could not verify' is not 'no policies'.""" + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", True, raising=False) + model = _Model(AIMessage(content="", tool_calls=[_CALL])) + + with patch.object( + AgentGraphAdapter, + "_tool_approval_policies_exist", + new=AsyncMock(side_effect=RuntimeError("backend down")), + ): + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == END and model.seen == [] + assert cmd.update["error"] == FC_TOOL_APPROVAL_UNVERIFIED.format(error="backend down") + + +@pytest.mark.asyncio +async def test_guard_asks_storage_strictly_so_a_backend_failure_is_not_an_empty_list(): + """PolicyStorage.list_policies swallows backend errors by default; the guard must not.""" + from types import SimpleNamespace + + calls = [] + + async def list_policies(**kwargs): + calls.append(kwargs) + return [] + + fake_system = SimpleNamespace(agent=SimpleNamespace(storage=SimpleNamespace(list_policies=list_policies))) + with patch( + "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.from_config", return_value=fake_system + ): + assert await _adapter()._tool_approval_policies_exist({}) is False + assert calls and calls[0]["strict"] is True diff --git a/src/cuga/backend/cuga_graph/policy/storage.py b/src/cuga/backend/cuga_graph/policy/storage.py index 4a8d25f98..38f342bf6 100644 --- a/src/cuga/backend/cuga_graph/policy/storage.py +++ b/src/cuga/backend/cuga_graph/policy/storage.py @@ -417,7 +417,15 @@ async def list_policies( policy_type: Optional[PolicyType] = None, enabled_only: bool = True, limit: int = 100, + *, + strict: bool = False, ) -> List[Policy]: + """List stored policies. + + ``strict=True`` re-raises a backend failure instead of returning ``[]``. + A safety guard that must fail closed cannot tell "no policies" from + "the backend is down" otherwise. + """ if not self._connected: await self.initialize_async() try: @@ -426,6 +434,8 @@ async def list_policies( policies.sort(key=lambda x: x.priority, reverse=True) return policies except Exception as e: + if strict: + raise logger.error(f"Failed to list policies: {e}") return [] diff --git a/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py b/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py new file mode 100644 index 000000000..0eb93115d --- /dev/null +++ b/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py @@ -0,0 +1,34 @@ +"""``PolicyStorage.list_policies(strict=True)`` surfaces backend failures. + +The default swallows them and returns ``[]`` — fine for lookups, wrong for a +guard that must fail closed and cannot tell "no policies" from "backend down". +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from cuga.backend.cuga_graph.policy.storage import PolicyStorage + +pytestmark = pytest.mark.unit + + +def _storage_with_broken_backend() -> PolicyStorage: + storage = PolicyStorage.__new__(PolicyStorage) + storage._connected = True + storage._backend = MagicMock() + storage._backend.list_policies = AsyncMock(side_effect=RuntimeError("backend down")) + return storage + + +@pytest.mark.asyncio +async def test_default_swallows_backend_errors(): + assert await _storage_with_broken_backend().list_policies() == [] + + +@pytest.mark.asyncio +async def test_strict_raises_backend_errors(): + with pytest.raises(RuntimeError, match="backend down"): + await _storage_with_broken_backend().list_policies(strict=True) diff --git a/tests/unit/test_fc_sdk_fail_closed_approval.py b/tests/unit/test_fc_sdk_fail_closed_approval.py index 8d16dbddc..768ae7785 100644 --- a/tests/unit/test_fc_sdk_fail_closed_approval.py +++ b/tests/unit/test_fc_sdk_fail_closed_approval.py @@ -32,10 +32,14 @@ async def ainvoke(self, messages, config=None, **kwargs): return self._responses.pop(0) +RAN: list = [] + + def _delete_record(): async def delete_record(record_id: int) -> str: """Delete a record.""" - raise AssertionError("a guarded tool ran without approval") + RAN.append(record_id) + return "deleted" return StructuredTool.from_function( coroutine=delete_record, name="delete_record", description="Delete a record." @@ -48,6 +52,9 @@ def _settings(monkeypatch): monkeypatch.setattr(settings.policy, "enabled", True, raising=False) monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) + RAN.clear() + yield + RAN.clear() @pytest.mark.asyncio @@ -90,4 +97,5 @@ async def test_fc_refuses_with_a_stored_tool_approval_policy_and_codeact_still_a result = await agent.invoke("delete record 1", thread_id="codeact-approval", execution_mode="codeact") assert codeact_model.invocations == 1 - assert "approval" in result.answer.lower(), result.answer + assert "Deleting needs approval." in result.answer, result.answer + assert RAN == [], "the guarded tool must not run on either path" From da9a5848061ac5ab363f538e9053dcfddc4215b1 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Sun, 27 Sep 2026 16:17:58 +0300 Subject: [PATCH 06/18] =?UTF-8?q?fix(cuga-lite):=20address=20haroldship's?= =?UTF-8?q?=20review=20=E2=80=94=20native=20transcript=20across=20turns,?= =?UTF-8?q?=20fail-closed=20binding,=20one-shot=20violations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Blocking items, each reproduced before the fix: - Function-calling replay is a real transcript again. AgentState dumps chat_messages by runtime type (SerializeAsAny) and revives the message classes on validation, so AIMessage.tool_calls and ToolMessage.tool_call_id survive every Command(update=state.model_dump()) hop instead of collapsing to bare BaseMessage shells; legacy shells from old checkpoints still validate as before. The SDK two-turn test now asserts turn 2 carries the original tool_calls id and its matching ToolMessage. The flatten path stays as defense for history that predates this. - bind_tools failure / unsupported model / nothing to bind: fail closed with a clear error instead of degrading to an unbound model (the grace turn, which withholds tools on purpose, is exempt). An empty executable set binds nothing rather than opening the registry-wide catalogue. - Mode violation: only Python fences count, and there is exactly one corrective turn — a second fence is delivered as the answer, so a fence-happy model can no longer loop to cuga_lite_max_steps. - tool_exec outer guard keeps the replies of calls that ran and answers every id the batch never reached. - settings.toml no longer advertises "native" as an alias. Non-blocking: a static prompt is ignored in function-calling mode (the FC prompt is rendered, with a warning); README states the timeout is per call. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 3 +- .../nodes/cuga_lite/adapter/graph_adapter.py | 54 ++++++++++++--- .../nodes/cuga_lite/adapter/prepare_node.py | 15 ++-- .../nodes/cuga_lite/adapter/tool_exec_node.py | 18 ++++- .../tests/test_fc_call_model_hook.py | 69 +++++++++++++++++-- .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 29 +++++++- .../cuga_lite/tests/test_tool_exec_node.py | 29 ++++++++ .../backend/cuga_graph/state/agent_state.py | 55 +++++++++++++-- src/cuga/settings.toml | 2 +- .../test_agent_state_keeps_message_objects.py | 47 +++++++++++++ .../unit/test_fc_sdk_fail_closed_approval.py | 12 +++- tests/unit/test_fc_sdk_two_turn.py | 17 ++++- 12 files changed, 319 insertions(+), 31 deletions(-) create mode 100644 tests/unit/test_agent_state_keeps_message_objects.py diff --git a/README.md b/README.md index ace7044a5..ac23faee3 100644 --- a/README.md +++ b/README.md @@ -1203,7 +1203,8 @@ result = await agent.invoke("...", execution_mode="codeact") # per-call overri - **Tool-approval policies**: function-calling mode has no approval interrupt yet. If an enabled tool-approval policy exists — or policy storage cannot be checked — the run stops with a clear error before any tool runs (fail closed); use `codeact` for that agent or disable the policy. - Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call. Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. - The bundled CodeAct few-shot demos are not sent in function-calling mode; pass your own through `configurable["mcp_few_shot_examples"]`. -- Tools run in-process and sequentially. The pre-execute VERIFY gate, reflection and the E2B / OpenSandbox executors are CodeAct-only; step discipline caps local-executor blocks only. +- Tools run in-process and sequentially. Each call is bounded by `sandbox_execution_timeout` (the timeout is per call, so a turn with N calls can take up to N × that). The pre-execute VERIFY gate, reflection and the E2B / OpenSandbox executors are CodeAct-only; step discipline caps local-executor blocks only. +- Function-calling needs tools bound natively: if `bind_tools` fails or the model does not support it, or the agent has no executable tools, the run stops with a clear error (fail closed) instead of silently answering without tools. A static prompt is ignored in this mode. - A thread can switch modes between turns: results from function-calling turns are shown to a later CodeAct turn as text.
diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 104c6bba1..aa2c5b060 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import re from typing import Any, Dict, List, Optional, Tuple from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage @@ -75,6 +76,12 @@ "It fails closed: no tool ran. Fix policy storage, or use cuga_lite_execution_mode = \"codeact\"." ) +FC_BIND_FAILED = ( + "Function-calling mode could not advertise any tool natively (bind_tools failed or is unsupported " + "by this model, or there is no executable tool to bind), so no native tool call is possible. " + 'It fails closed: no model call was made. Use cuga_lite_execution_mode = "codeact" for this model.' +) + FC_STEP_LIMIT_CALL_REPLY = "Not executed: the step limit was reached before this call could run." FC_BUDGET_CALL_REPLY = ( "Not executed: the tool budget for this turn is spent. Answer from the data already retrieved." @@ -199,6 +206,22 @@ def close_pending() -> None: return out +_PYTHON_FENCE = re.compile(r"```(?:python|py)\b", re.IGNORECASE) +_UNTAGGED_FENCE = re.compile(r"```[ \t]*\n(.*?)```", re.DOTALL) + + +def _looks_like_python_block(content: str) -> bool: + """A fenced block the CodeAct sandbox would have executed. + + A ``python``/``py`` fence, or an untagged fence whose body awaits or calls + something. Fences carrying JSON, text or shell output in a final answer are + not violations. + """ + if _PYTHON_FENCE.search(content): + return True + return any(re.search(r"\bawait\b|\w+\(", body) for body in _UNTAGGED_FENCE.findall(content)) + + def _few_shot_to_messages(few_shot: List[Any]) -> List[BaseMessage]: """The prepare node's normalized ``{role, content}`` demos, as real chat messages.""" msgs: List[BaseMessage] = [] @@ -350,10 +373,11 @@ async def resolve_bind_tools( "[fc] bind_tools mode was 'none'; advertising the {} executable tool(s)", len(names) ) else: - configurable = {**(configurable or {}), "cuga_lite_bind_tools_mode": "all"} - logger.info( - "[fc] bind_tools mode was 'none'; upgraded to 'all' for function-calling mode" - ) + # Nothing executable was recorded: advertising the full catalogue + # would offer tools tool_exec cannot run. Bind nothing; the turn + # then fails closed in execute_call_model_fc. + logger.error("[fc] no executable tools recorded by prepare; binding nothing") + return None try: return await resolve_model_with_bind_tools( active_model, @@ -537,6 +561,13 @@ async def execute_call_model_fc( self, history, AIMessage(content=blocked_reason), state.step_count ) + if not budget_exhausted and bound is active_model: + # bind_tools failed, is unsupported by the model, or had nothing to bind — + # call_model fell back to the unbound model. That run cannot make a native + # tool call, so stop with a clear error instead of degrading to text. + logger.error("[fc] refusing to start: no tools are bound to the model") + return create_error_command(self, history, AIMessage(content=FC_BIND_FAILED), state.step_count) + msgs: List[BaseMessage] = [SystemMessage(content=system_content)] msgs.extend(_few_shot_to_messages(self.get_few_shot_messages(state))) msgs.extend(_normalize_history_for_replay(history)) @@ -643,11 +674,16 @@ async def execute_call_model_fc( }, ) - if "```" in content and not budget_exhausted: - # Mode violation: never execute code here. One corrective turn, charged - # as a step so it cannot loop past cuga_lite_max_steps. - violations = int(base_meta.get("fc_mode_violations", 0) or 0) + 1 - logger.warning("[fc] mode violation #{}: code block emitted in function-calling mode", violations) + prior_violations = int(base_meta.get("fc_mode_violations", 0) or 0) + if _looks_like_python_block(content) and not budget_exhausted and prior_violations == 0: + # Mode violation: never execute code here. Exactly one corrective turn + # (charged as a step); a second fence is delivered as the answer, so a + # fence-happy model cannot loop to cuga_lite_max_steps. Only Python + # blocks count — a ```json or ```text snippet in an answer is fine. + violations = prior_violations + 1 + logger.warning( + "[fc] mode violation: python block emitted in function-calling mode — one corrective turn" + ) return Command( goto="call_model", update={ diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index 2fb56bd0b..01ce06ef5 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -764,10 +764,16 @@ async def _wrapped(*args, **kwargs): # Create prompt dynamically dynamic_prompt = adapter._static_prompt - if not dynamic_prompt and _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: + if _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: # Function-calling mode: the CodeAct prompt is never rendered. Tools # travel natively via bind_tools, so this prompt is only the behavioural - # contract plus the same instructions / special_instructions. + # contract plus the same instructions / special_instructions. A static + # prompt is ignored here: it is CodeAct-shaped and would tell the model + # to write the very fences this mode treats as violations. + if dynamic_prompt: + logger.warning( + "Function-calling mode ignores the static prompt and renders the FC prompt instead" + ) _fragments = resolve_fc_prompt_fragments(configurable, _runtime_model_name) _one_per_step = ( resolve_step_discipline(configurable, _runtime_model_name) @@ -823,11 +829,6 @@ async def _wrapped(*args, **kwargs): len(dynamic_prompt), ) else: - if _execution_mode == EXECUTION_MODE_FUNCTION_CALLING: - logger.warning( - "Function-calling mode with a static prompt: the prompt is used verbatim; " - "make sure it does not instruct the model to write code" - ) logger.info( "Using static CugaLite prompt; dynamic few-shot injection skipped " "(enable_find_tools={} few_shot_turns={})", diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py index 114a3d0ad..43c2cecd3 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py @@ -254,9 +254,25 @@ async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): # handling must still report the budget the batch spent, or the # checkpoint keeps the pre-batch counts (see _budget_updates). logger.error(f"tool_exec failed: {exc}") + # Tools may already have run: keep their replies, and answer every id + # the batch never reached, so the persisted turn has no dangling call. + answered = {m.tool_call_id for m in results} + replies = list(results) + for index, call in enumerate(list(calls) + list(invalid)): + call_id, name, _ = _tool_call_parts(call) + call_id = call_id or f"call_{index}" + if call_id not in answered: + answered.add(call_id) + replies.append( + _error_message( + f"Not executed: tool execution failed before this call ran ({exc}).", + call_id=call_id, + name=name, + ) + ) return core_create_error_command( adapter, - adapter.get_messages(state), + adapter.get_messages(state) + replies, AIMessage(content=f"Error during tool execution: {exc}"), state.step_count, additional_updates={ diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index 18c27ab0e..cb7e539f1 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -26,6 +26,7 @@ from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import TOOL_BUDGET_EXHAUSTED_INSTRUCTION from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import graph_adapter as ga from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( + FC_BIND_FAILED, FC_BUDGET_CALL_REPLY, FC_MODE_VIOLATION_CORRECTION, FC_STEP_LIMIT_CALL_REPLY, @@ -39,6 +40,7 @@ pytestmark = pytest.mark.unit FC = {"cuga_lite_execution_mode": "function_calling"} +_UNBOUND = object() _CALL = {"name": "add", "args": {"a": 1, "b": 2}, "id": "c1", "type": "tool_call"} @@ -81,7 +83,7 @@ async def _turn( state=state, config=None, configurable=configurable, - active_model=model, + active_model=_UNBOUND, # what call_model passes when bind_tools succeeded: bound is not the model bound=model, invoke_config={}, system_content=system, @@ -327,14 +329,15 @@ async def fake_resolve(active_model, **kwargs): with patch.object(ga, "resolve_model_with_bind_tools", new=AsyncMock(side_effect=fake_resolve)): adapter = _adapter() - assert await adapter.resolve_bind_tools(_state(), object(), dict(FC), None) == "BOUND" + # No executable set recorded: bind nothing (the turn then fails closed) rather + # than open the registry-wide catalogue to tools tool_exec cannot run. + assert await adapter.resolve_bind_tools(_state(), object(), dict(FC), None) is None await adapter.resolve_bind_tools( _state(), object(), {**FC, "cuga_lite_bind_tools_mode": "find_tools"}, None ) await adapter.resolve_bind_tools(_state(), object(), {}, None) - fc_default, fc_explicit, codeact = captured - assert fc_default["cuga_lite_bind_tools_mode"] == "all", "FC is inert without advertised tools" + fc_explicit, codeact = captured assert fc_explicit["cuga_lite_bind_tools_mode"] == "find_tools", "an explicit choice is respected" assert "cuga_lite_bind_tools_mode" not in codeact, "codeact is byte-identical" @@ -394,3 +397,61 @@ async def list_policies(**kwargs): ): assert await _adapter()._tool_approval_policies_exist({}) is False assert calls and calls[0]["strict"] is True + + +@pytest.mark.asyncio +async def test_refuses_when_no_tools_are_bound(): + """bind_tools failed / unsupported / nothing to bind: call_model hands the unbound model + through. Native tool calling is impossible, so stop — do not degrade to text.""" + model = _Model(AIMessage(content="", tool_calls=[_CALL])) + cmd = await _adapter().execute_call_model_fc( + state=_state(), + config=None, + configurable=FC, + active_model=model, + bound=model, # `resolve_bind_tools(...) or active_model` fell back + invoke_config={}, + system_content="SYS", + modified_messages=[HumanMessage(content="q")], + budget_exhausted=False, + playbook_fired=False, + ) + assert cmd.goto == END and cmd.update["error"] == FC_BIND_FAILED and model.seen == [] + + +@pytest.mark.asyncio +async def test_budget_exhausted_grace_turn_runs_unbound_on_purpose(): + """The grace turn deliberately withholds tools, so an unbound model is not a failure there.""" + model = _Model(AIMessage(content="From what I have: 3.")) + cmd = await _adapter().execute_call_model_fc( + state=_state(), + config=None, + configurable=FC, + active_model=model, + bound=model, + invoke_config={}, + system_content="SYS", + modified_messages=[HumanMessage(content="q")], + budget_exhausted=True, + playbook_fired=False, + ) + assert cmd.goto == END and cmd.update["final_answer"] == "From what I have: 3." + + +@pytest.mark.asyncio +async def test_mode_violation_is_corrected_once_then_the_fence_is_delivered(): + fence = "```python\nawait add(a=1, b=2)\n```" + first = await _turn(_adapter(), _Model(AIMessage(content=fence)), _state(), FC) + assert first.goto == "call_model" and first.update["cuga_lite_metadata"]["fc_mode_violations"] == 1 + + again = _state(metadata={"fc_mode_violations": 1}) + second = await _turn(_adapter(), _Model(AIMessage(content=fence)), again, FC) + assert second.goto == END, "one correction only — never a loop to cuga_lite_max_steps" + assert second.update["final_answer"] == fence + + +@pytest.mark.asyncio +async def test_non_python_fences_in_an_answer_are_not_violations(): + answer = "Here is the record:\n```json\n{\"id\": 42}\n```" + cmd = await _turn(_adapter(), _Model(AIMessage(content=answer)), _state(), FC) + assert cmd.goto == END and cmd.update["final_answer"] == answer diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index df60ae2cc..a86a6d8a8 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -33,6 +33,14 @@ CALLS: list = [] +class _Bound: + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + class _ScriptedModel: """Queued responses; an entry may be a callable that receives the outbound messages.""" @@ -44,7 +52,7 @@ def __init__(self, responses): def bind_tools(self, tools, **kwargs): self.bound_tool_names = [getattr(t, "name", str(t)) for t in tools] - return self + return _Bound(self) # a real bind returns a new runnable, never the model itself async def ainvoke(self, messages, config=None, **kwargs): self.invocations += 1 @@ -293,3 +301,22 @@ async def test_bundled_codeact_few_shots_are_withheld_in_fc_mode(): _echo_tool(), ) assert [m.content for m in explicit.seen[0][1:3]] == ["demo q", "demo a"], "explicit demos still replayed" + + +@pytest.mark.asyncio +async def test_static_prompt_is_ignored_in_fc_mode(): + """A static prompt is CodeAct-shaped; used verbatim it would ask for the fences FC treats as violations.""" + model = _ScriptedModel([AIMessage(content="done")]) + graph = create_cuga_lite_graph( + model=model, + prompt="You are a coder. Always answer with a ```python block.", + tool_provider=_provider(_echo_tool()), + apps_list=[], + thread_id="t", + ).compile(checkpointer=MemorySaver()) + result = await graph.ainvoke( + CugaLiteState(chat_messages=[HumanMessage(content="hi")]), + config=_config("fc-static", cuga_lite_execution_mode="function_calling"), + ) + assert "native function-calling" in result["prepared_prompt"] + assert "Always answer with a ```python block" not in result["prepared_prompt"] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py index ab854d52d..003a766e5 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py @@ -414,6 +414,35 @@ def boom(*args, **kwargs): assert "history store unavailable" in result.update["error"] assert result.update["tool_calls_used_run"] == 1, "the call that ran must still be counted" assert ToolCallTracker.is_enabled() is False + replies = [m for m in result.update["chat_messages"] if isinstance(m, ToolMessage)] + assert [(m.tool_call_id, m.content) for m in replies] == [("c1", "2")], ( + "the executed call's reply is kept" + ) + + +@pytest.mark.asyncio +async def test_outer_guard_answers_the_calls_the_batch_never_reached(monkeypatch): + _caps(monkeypatch) + from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import tool_exec_node as mod + + async def add(a: int, b: int) -> int: + return a + b + + async def broken_batch(adapter, calls, invalid, results, **kwargs): + results.append(ToolMessage(content="2", tool_call_id="c1", name="add")) + raise RuntimeError("executor died") + + monkeypatch.setattr(mod, "_run_batch", broken_batch) + node = create_tool_exec_node(_Adapter({"add": add})) + last = AIMessage( + content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] + ) + + result = await node(_state(last), config=None) + + replies = {m.tool_call_id: m for m in result.update["chat_messages"] if isinstance(m, ToolMessage)} + assert replies["c1"].content == "2" + assert replies["c2"].status == "error" and "executor died" in replies["c2"].content, "no dangling id" @pytest.mark.asyncio diff --git a/src/cuga/backend/cuga_graph/state/agent_state.py b/src/cuga/backend/cuga_graph/state/agent_state.py index 6763c53fd..dba1dc45b 100644 --- a/src/cuga/backend/cuga_graph/state/agent_state.py +++ b/src/cuga/backend/cuga_graph/state/agent_state.py @@ -5,8 +5,15 @@ from datetime import datetime from pathlib import Path -from langchain_core.messages import AIMessage, BaseMessage -from pydantic import BaseModel, Field +from langchain_core.messages import ( + AIMessage, + BaseMessage, + FunctionMessage, + HumanMessage, + SystemMessage, + ToolMessage, +) +from pydantic import BaseModel, Field, SerializeAsAny, field_validator from loguru import logger from cuga.backend.cuga_graph.nodes.cuga_agent_core.schemas.api_models import ApiDescription @@ -968,6 +975,15 @@ class AnalyzeTaskAppsOutput(BaseModel): type: Literal['api', 'web'] = 'web' +_MESSAGE_TYPES = { + "human": HumanMessage, + "ai": AIMessage, + "tool": ToolMessage, + "system": SystemMessage, + "function": FunctionMessage, +} + + class AgentState(BaseModel): # pages: Annotated[Sequence[str], operator.add] # List of pages traversed # page: Page # The Playwright web page lets us interact with the web environment @@ -996,11 +1012,40 @@ class AgentState(BaseModel): current_app_description: Optional[str] = None api_last_step: Optional[str] = None guidance: Optional[str] = None - chat_messages: Optional[List[BaseMessage]] = Field(default_factory=list) - chat_agent_messages: Optional[List[BaseMessage]] = Field(default_factory=list) - supervisor_chat_messages: Optional[List[BaseMessage]] = Field( + # SerializeAsAny: nodes pass state around as ``Command(update=state.model_dump())``, + # and pydantic serialises a ``List[BaseMessage]`` by the *declared* type — so an + # AIMessage lost its ``tool_calls`` and a ToolMessage its ``tool_call_id`` at every + # hop, and the next turn revalidated them into bare BaseMessage shells. Dumping by + # the runtime type keeps those fields; ``_revive_message_types`` below turns the + # dicts back into the right message class on the way in. + chat_messages: Optional[List[SerializeAsAny[BaseMessage]]] = Field(default_factory=list) + chat_agent_messages: Optional[List[SerializeAsAny[BaseMessage]]] = Field(default_factory=list) + supervisor_chat_messages: Optional[List[SerializeAsAny[BaseMessage]]] = Field( default_factory=list ) # Supervisor's conversation history + + @field_validator("chat_messages", "chat_agent_messages", "supervisor_chat_messages", mode="before") + @classmethod + def _revive_message_types(cls, value: Any) -> Any: + """Rebuild message subclasses from dumped dicts (``type`` decides the class). + + A dict that cannot be rebuilt — e.g. a ``tool`` shell from a checkpoint + written before the fields were preserved, with no ``tool_call_id`` — is + left for the default ``BaseMessage`` validation, exactly as before. + """ + if not value: + return value + revived = [] + for m in value: + if isinstance(m, dict) and m.get("type") in _MESSAGE_TYPES: + try: + revived.append(_MESSAGE_TYPES[m["type"]](**{k: v for k, v in m.items() if k != "type"})) + continue + except Exception: + pass + revived.append(m) + return revived + # Mirrors CugaSupervisorState.supervisor_metadata so CugaSupervisorNode can read/set approval # flags (e.g. plan_approved) on the outer AgentState across the HITL interrupt boundary — the # subgraph's own state does not survive a WaitForResponse pause/resume. diff --git a/src/cuga/settings.toml b/src/cuga/settings.toml index bd6831789..f689500d9 100644 --- a/src/cuga/settings.toml +++ b/src/cuga/settings.toml @@ -55,7 +55,7 @@ cuga_lite_bind_tools_max_count = 128 # Pad shortlister output to fill the cap. Off by default for code-execution mode. cuga_lite_bind_tools_pad_to_cap = false # Execution mode: codeact (write Python, run it in the sandbox — default) | function_calling (the model -# issues native tool_calls, a tool_exec node runs them and replies with ToolMessages). Aliases: fc, native. +# issues native tool_calls, a tool_exec node runs them and replies with ToolMessages). Alias: fc. # Per-invoke: config["configurable"]["cuga_lite_execution_mode"]. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_EXECUTION_MODE cuga_lite_execution_mode = "codeact" # Step discipline, both modes: off (default) | one_tool_per_step — observe one result before the next call. diff --git a/tests/unit/test_agent_state_keeps_message_objects.py b/tests/unit/test_agent_state_keeps_message_objects.py new file mode 100644 index 000000000..9f648a938 --- /dev/null +++ b/tests/unit/test_agent_state_keeps_message_objects.py @@ -0,0 +1,47 @@ +"""``AgentState.model_dump()`` keeps message subclass fields; validation revives the classes. + +Nodes pass state around as ``Command(update=state.model_dump())``. Serialising a +``List[BaseMessage]`` by the declared type dropped ``tool_calls`` and +``tool_call_id`` at every hop, so the next turn revalidated them into bare +``BaseMessage`` shells — the root cause of the flattened function-calling replay. +""" + +from __future__ import annotations + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from cuga.backend.cuga_graph.state.agent_state import AgentState + +pytestmark = pytest.mark.unit + +_CALL = {"name": "echo", "args": {"value": 7}, "id": "call_1", "type": "tool_call"} + + +def test_model_dump_keeps_tool_calls_and_tool_call_ids(): + state = AgentState( + chat_messages=[ + HumanMessage(content="echo 7"), + AIMessage(content="", tool_calls=[_CALL]), + ToolMessage(content="7", tool_call_id="call_1", name="echo"), + ] + ) + + dumped = state.model_dump()["chat_messages"] + + assert dumped[1]["type"] == "ai" and dumped[1]["tool_calls"] == [_CALL], "dumped by runtime type" + assert dumped[2]["type"] == "tool" and dumped[2]["tool_call_id"] == "call_1" + + # and they revalidate as themselves on the next hop + again = AgentState(**state.model_dump()) + assert type(again.chat_messages[2]) is ToolMessage and again.chat_messages[2].tool_call_id == "call_1" + assert type(again.chat_messages[1]) is AIMessage and again.chat_messages[1].tool_calls == [_CALL] + + +def test_legacy_shells_from_old_checkpoints_still_validate(): + """A tool shell written before the fields were preserved has no tool_call_id: keep it + as the bare BaseMessage it always was rather than failing the whole state.""" + state = AgentState( + chat_messages=[{"type": "tool", "content": "7", "additional_kwargs": {}, "response_metadata": {}}] + ) + assert state.chat_messages[0].type == "tool" and state.chat_messages[0].content == "7" diff --git a/tests/unit/test_fc_sdk_fail_closed_approval.py b/tests/unit/test_fc_sdk_fail_closed_approval.py index 768ae7785..70f680712 100644 --- a/tests/unit/test_fc_sdk_fail_closed_approval.py +++ b/tests/unit/test_fc_sdk_fail_closed_approval.py @@ -17,13 +17,23 @@ pytestmark = pytest.mark.unit +class _Bound: + """What ``bind_tools`` returns: a runnable wrapping the model.""" + + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + class _ScriptedModel: def __init__(self, responses): self._responses = list(responses) self.invocations = 0 def bind_tools(self, tools, **kwargs): - return self + return _Bound(self) # a real bind returns a new runnable, never the model itself async def ainvoke(self, messages, config=None, **kwargs): self.invocations += 1 diff --git a/tests/unit/test_fc_sdk_two_turn.py b/tests/unit/test_fc_sdk_two_turn.py index 0aeaba1cb..3434b539b 100644 --- a/tests/unit/test_fc_sdk_two_turn.py +++ b/tests/unit/test_fc_sdk_two_turn.py @@ -21,13 +21,23 @@ CALLS: list = [] +class _Bound: + """What ``bind_tools`` returns: a runnable wrapping the model.""" + + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + class _ScriptedModel: def __init__(self, responses): self._responses = list(responses) self.seen: list = [] def bind_tools(self, tools, **kwargs): - return self + return _Bound(self) # a real bind returns a new runnable, never the model itself async def ainvoke(self, messages, config=None, **kwargs): self.seen.append(list(messages)) @@ -98,6 +108,11 @@ async def test_second_turn_replays_a_valid_function_calling_transcript(): assert all(type(m) is not BaseMessage for m in outbound), ( f"bare BaseMessage replayed to the provider: {shapes}" ) + # Native transcript, not a flattened one: turn 1's call and its reply survive the SDK boundary as-is. + ai_turn1 = [m for m in outbound if isinstance(m, AIMessage) and m.tool_calls] + assert ai_turn1 and ai_turn1[0].tool_calls[0]["id"] == "call_1", shapes + tool_turn1 = [m for m in outbound if isinstance(m, ToolMessage)] + assert tool_turn1 and tool_turn1[0].tool_call_id == "call_1" and tool_turn1[0].content == "7", shapes ai_with_calls = [m for m in outbound if isinstance(m, AIMessage) and m.tool_calls] tool_replies = [m for m in outbound if isinstance(m, ToolMessage)] dangling = {c["id"] for m in ai_with_calls for c in m.tool_calls} - {t.tool_call_id for t in tool_replies} From b1d13d0d7189e9a66ed7586e07e8562b0914a8d2 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Mon, 28 Sep 2026 07:41:44 +0300 Subject: [PATCH 07/18] fix(cuga-lite): function-calling correction cap is per turn; fences only count when they run something MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The one-correction cap for a Python fence lived in cuga_lite_metadata, which the SDK and server carry into the next turn on the same thread (only policy_decisions was reset). After one correction anywhere in a thread, every later turn's first fence was delivered as the answer. prepare — which runs exactly once per turn on both paths — now drops the counter. Reproduced with two invoke() calls on one thread, each scripted as [fence, answer]: turn 2 returned the raw fence. Untagged fences count as a violation only when they `await` or call a bound tool by name, so notation such as ```f(x) = 2``` in an answer no longer costs a corrective turn. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/graph_adapter.py | 22 ++++-- .../nodes/cuga_lite/adapter/prepare_node.py | 7 +- .../tests/test_fc_call_model_hook.py | 11 +++ tests/unit/test_fc_sdk_violation_per_turn.py | 73 +++++++++++++++++++ 4 files changed, 106 insertions(+), 7 deletions(-) create mode 100644 tests/unit/test_fc_sdk_violation_per_turn.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index aa2c5b060..4aff26862 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -210,16 +210,22 @@ def close_pending() -> None: _UNTAGGED_FENCE = re.compile(r"```[ \t]*\n(.*?)```", re.DOTALL) -def _looks_like_python_block(content: str) -> bool: +def _looks_like_python_block(content: str, tool_names: Any = ()) -> bool: """A fenced block the CodeAct sandbox would have executed. - A ``python``/``py`` fence, or an untagged fence whose body awaits or calls - something. Fences carrying JSON, text or shell output in a final answer are - not violations. + A ``python``/``py`` fence, or an untagged fence whose body ``await``s or + calls one of the bound tools by name. Fences carrying JSON, text, shell + output or plain notation (``f(x) = 2``) in a final answer are not violations. """ if _PYTHON_FENCE.search(content): return True - return any(re.search(r"\bawait\b|\w+\(", body) for body in _UNTAGGED_FENCE.findall(content)) + names = {n for n in (tool_names or ()) if n and not str(n).startswith("_")} + for body in _UNTAGGED_FENCE.findall(content): + if re.search(r"\bawait\b", body): + return True + if names and any(re.search(rf"\b{re.escape(n)}\s*\(", body) for n in names): + return True + return False def _few_shot_to_messages(few_shot: List[Any]) -> List[BaseMessage]: @@ -675,7 +681,11 @@ async def execute_call_model_fc( ) prior_violations = int(base_meta.get("fc_mode_violations", 0) or 0) - if _looks_like_python_block(content) and not budget_exhausted and prior_violations == 0: + if ( + _looks_like_python_block(content, self._tools_context) + and not budget_exhausted + and prior_violations == 0 + ): # Mode violation: never execute code here. Exactly one corrective turn # (charged as a step); a second fence is delivered as the answer, so a # fence-happy model cannot loop to cuga_lite_max_steps. Only Python diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index 01ce06ef5..b2168db7a 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -838,11 +838,16 @@ async def _wrapped(*args, **kwargs): reflection_apps_snapshot = format_apps_for_prompt(apps_for_prompt or []) + # The function-calling one-correction cap is per turn. Metadata is carried + # across turns by the SDK / server (only policy_decisions is reset there), so + # drop the counter here — prepare runs exactly once per turn on both paths. + _turn_metadata = dict(state.cuga_lite_metadata or {}) + _turn_metadata.pop("fc_mode_violations", None) update_payload: dict[str, Any] = { "tools_prepared": True, "prepared_prompt": dynamic_prompt, "step_count": 0, - "cuga_lite_metadata": state.cuga_lite_metadata, + "cuga_lite_metadata": _turn_metadata, "reflection_apps": reflection_apps_snapshot, "reflection_enable_find_tools": enable_find_tools, "reflection_skills_enabled": skills_enabled, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index cb7e539f1..2d9c3f849 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -34,6 +34,7 @@ FC_TOOL_APPROVAL_UNVERIFIED, FC_UNANSWERED_CALL_REPLY, AgentGraphAdapter, + _looks_like_python_block, _normalize_history_for_replay, ) @@ -455,3 +456,13 @@ async def test_non_python_fences_in_an_answer_are_not_violations(): answer = "Here is the record:\n```json\n{\"id\": 42}\n```" cmd = await _turn(_adapter(), _Model(AIMessage(content=answer)), _state(), FC) assert cmd.goto == END and cmd.update["final_answer"] == answer + + +def test_untagged_fences_are_violations_only_when_they_run_something(): + assert _looks_like_python_block("```python\nx = 1\n```") + assert _looks_like_python_block("```\nawait add(a=1, b=2)\n```") + assert _looks_like_python_block("```\nadd(a=1, b=2)\n```", {"add"}), "a bound tool called by name" + assert not _looks_like_python_block("The function is\n```\nf(x) = 2x + 1\n```", {"add"}), ( + "notation, not code" + ) + assert not _looks_like_python_block("```json\n{\"a\": 1}\n```", {"add"}) diff --git a/tests/unit/test_fc_sdk_violation_per_turn.py b/tests/unit/test_fc_sdk_violation_per_turn.py new file mode 100644 index 000000000..523047c1c --- /dev/null +++ b/tests/unit/test_fc_sdk_violation_per_turn.py @@ -0,0 +1,73 @@ +"""The function-calling one-correction cap is per turn, not per thread. + +``fc_mode_violations`` lives in ``cuga_lite_metadata``, which the SDK carries into +the next turn on the same thread. Without a per-turn reset, one correction anywhere +in a thread made every later turn's first Python fence the delivered answer. +""" + +from __future__ import annotations + +import pytest +from langchain_core.messages import AIMessage +from langchain_core.tools import StructuredTool + +pytestmark = pytest.mark.unit + + +class _Bound: + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + + def bind_tools(self, tools, **kwargs): + return _Bound(self) + + async def ainvoke(self, messages, config=None, **kwargs): + if not self._responses: + raise AssertionError("model asked for more responses than scripted") + return self._responses.pop(0) + + +def _echo(): + async def echo(value: int) -> int: + """Echo a value.""" + return value + + return StructuredTool.from_function(coroutine=echo, name="echo", description="Echo a value.") + + +@pytest.fixture(autouse=True) +def _quiet(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) + + +@pytest.mark.asyncio +async def test_each_turn_gets_its_own_correction(): + from cuga.sdk import CugaAgent + + fence = "```python\nawait echo(value=1)\n```" + model = _ScriptedModel( + [ + AIMessage(content=fence), + AIMessage(content="one"), + AIMessage(content=fence), + AIMessage(content="two"), + ] + ) + agent = CugaAgent(tools=[_echo()], model=model, execution_mode="function_calling") + + turn1 = await agent.invoke("first", thread_id="fc-violations") + turn2 = await agent.invoke("second", thread_id="fc-violations") + + assert turn1.answer == "one" + assert turn2.answer == "two", "turn 2 must be corrected too, not handed the raw fence" From b91abeafbce181f8304549836204b820e759e8f4 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Thu, 1 Oct 2026 10:34:13 +0300 Subject: [PATCH 08/18] fix(cuga-lite): read WatsonX's reasoning key on the function-calling path too #796 taught normalize_response to accept either additional_kwargs spelling ("reasoning_content" or "reasoning", the one WatsonX reports). The function-calling turn read only the first, so on WatsonX a reply with empty content and reasoning-only text lost its empty-content fallback. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/graph_adapter.py | 6 +++--- .../nodes/cuga_lite/tests/test_fc_call_model_hook.py | 12 ++++++++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 4aff26862..6b4034220 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -595,9 +595,9 @@ async def execute_call_model_fc( tool_calls = list(getattr(response, "tool_calls", None) or []) invalid_tool_calls = list(getattr(response, "invalid_tool_calls", None) or []) content = strip_harmony_tokens(normalize_assistant_text(getattr(response, "content", "")) or "") - reasoning = normalize_assistant_text( - (getattr(response, "additional_kwargs", None) or {}).get("reasoning_content") - ) + _ak = getattr(response, "additional_kwargs", None) or {} + # Both spellings, same as normalize_response: WatsonX reports "reasoning" (#796). + reasoning = normalize_assistant_text(_ak.get("reasoning_content") or _ak.get("reasoning")) if not isinstance(response, AIMessage): response = AIMessage(content=content, tool_calls=tool_calls) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index 2d9c3f849..d7d4c8f69 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -466,3 +466,15 @@ def test_untagged_fences_are_violations_only_when_they_run_something(): "notation, not code" ) assert not _looks_like_python_block("```json\n{\"a\": 1}\n```", {"add"}) + + +@pytest.mark.asyncio +async def test_watsonx_reasoning_key_is_read_like_the_codeact_path(): + """#796: WatsonX reports the reasoning channel as ``reasoning``, not ``reasoning_content``. + The FC turn uses it for the empty-content fallback, so it must read both spellings.""" + response = AIMessage(content="", additional_kwargs={"reasoning": "The total is 42."}) + state = _state(metadata={EMPTY_RESPONSE_CORRECTION_KEY: True}) # past the empty-reply retry + + cmd = await _turn(_adapter(), _Model(response), state, FC) + + assert cmd.goto == END and cmd.update["final_answer"] == "The total is 42." From a9ea3ec6b8a71b7218f4e4e3cfb19af10081590c Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Thu, 1 Oct 2026 20:12:31 +0300 Subject: [PATCH 09/18] fix(cuga-lite): provider-safe tool names for native tool binding - OpenAI-compatible APIs reject the whole request when a tool name is outside ^[A-Za-z0-9_-]{1,64}$ (Azure: HTTP 422); registry names (_) can be longer. - _safe_bind binds such tools as copies named after a 55-character stem, "_" and 8 hex digits of the SHA-1 of the real name. Legal names are bound unchanged; two tools that would share a bound name raise. - A tool call returned under an alias is mapped back to the real name before code is generated (native tool_calls and the tool_use_failed recovery), so the approval check and the sandbox see the real name. No alias keys are added to the execution context. Fixes #879 Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/graph_adapter.py | 6 +- .../nodes/cuga_lite/adapter/response_utils.py | 25 +- .../nodes/cuga_lite/bind_tools/__init__.py | 11 +- .../nodes/cuga_lite/bind_tools/tool_names.py | 128 ++++++ .../nodes/cuga_lite/helpers/bind_tools.py | 5 + tests/unit/test_provider_safe_tool_names.py | 421 ++++++++++++++++++ 6 files changed, 589 insertions(+), 7 deletions(-) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py create mode 100644 tests/unit/test_provider_safe_tool_names.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 5ba3fecd0..41ce0b7d4 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -24,6 +24,7 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.prepare_node import create_prepare_tools_and_apps_node from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.response_utils import ( clean_empty_response_retry_meta, + extract_code_from_failed_tool_call, extract_code_from_response_tool_calls, ) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import create_sandbox_node @@ -36,7 +37,6 @@ normalize_assistant_text, ) from cuga.backend.cuga_graph.utils.token_counter import clamp_watsonx_completion_for_messages -from cuga.backend.llm.errors import extract_code_from_tool_use_failed from cuga.config import settings @@ -132,7 +132,7 @@ async def ainvoke_model(self, bound: Any, messages: list, invoke_config: dict) - clamp_watsonx_completion_for_messages(bound, messages) return await bound.ainvoke(messages, config=invoke_config) except Exception as exc: - code = extract_code_from_tool_use_failed(exc) + code = extract_code_from_failed_tool_call(exc, self._tools_context) if code: logger.warning( "Model attempted tool call without tools bound (tool_use_failed). " @@ -174,7 +174,7 @@ def normalize_response(self, response: Any) -> Tuple[str, Optional[str]]: # downstream surface inherits clean text (see the base implementation). content = strip_harmony_tokens(normalize_assistant_text(response.content)) if not content: - tool_code = extract_code_from_response_tool_calls(response) + tool_code = extract_code_from_response_tool_calls(response, self._tools_context) if tool_code: logger.warning("Empty content with tool_calls detected; recovering tool call as Python code") content = tool_code diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py index 220c74100..ba4fafecc 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -from typing import Any, Dict, Optional +from typing import Any, Collection, Dict, Optional from langchain_core.messages import HumanMessage @@ -12,7 +12,9 @@ EMPTY_RESPONSE_CORRECTION_KEY, EXECUTION_OUTPUT_PREFIX, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools import resolve_tool_name from cuga.backend.cuga_graph.nodes.cuga_lite.reflection.verify_result import VERIFY_BLOCKED_PREFIX +from cuga.backend.llm.errors import failed_gen_to_code, parse_tool_use_failed_generation def clean_empty_response_retry_meta(meta: Optional[Dict[str, Any]]) -> Dict[str, Any]: @@ -45,8 +47,13 @@ def tool_call_kwarg_literal(value: Any) -> str: return repr(value) -def extract_code_from_response_tool_calls(response: Any) -> Optional[str]: - """Recover fenced Python from AIMessage.tool_calls when content is empty.""" +def extract_code_from_response_tool_calls(response: Any, tool_names: Collection[str] = ()) -> Optional[str]: + """Recover fenced Python from AIMessage.tool_calls when content is empty. + + A provider-safe alias the model was bound with is mapped back to its real + name in ``tool_names``, so the code (and the approval check that reads it) + names the real tool. + """ tool_calls = getattr(response, "tool_calls", None) or ( getattr(response, "additional_kwargs", None) or {} ).get("tool_calls") @@ -67,8 +74,20 @@ def extract_code_from_response_tool_calls(response: Any) -> Optional[str]: if not name: return None + if isinstance(name, str): + name = resolve_tool_name(name, tool_names) args_str = ", ".join( f"{k}={tool_call_kwarg_literal(v)}" for k, v in (args if isinstance(args, dict) else {}).items() ) return f"```python\nresult = await {name}({args_str})\nprint(result)\n```" + + +def extract_code_from_failed_tool_call(err: Any, tool_names: Collection[str] = ()) -> Optional[str]: + """``llm.errors.extract_code_from_tool_use_failed``, with an alias mapped back as above.""" + failed_gen = parse_tool_use_failed_generation(err) + if not isinstance(failed_gen, dict): + return None + if isinstance(failed_gen.get("name"), str): + failed_gen = {**failed_gen, "name": resolve_tool_name(failed_gen["name"], tool_names)} + return failed_gen_to_code(failed_gen) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py index 811051ad0..a2c85b10c 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py @@ -1,7 +1,8 @@ """bind_tools cap + shortlist machinery for cuga_lite. Keeps cuga_lite_graph.py focused on orchestration. See :mod:`.cap` for the -provider-safe cap and shortlister flow. +provider-safe cap and shortlister flow, and :mod:`.tool_names` for +provider-safe tool names. """ from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.cap import ( @@ -10,10 +11,18 @@ bind_tools_max_count_from_settings, bind_tools_pad_to_cap_from_settings, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import ( + provider_safe_tool_name, + provider_safe_tools, + resolve_tool_name, +) __all__ = [ "BindToolsUnsupportedError", "apply_bind_tools_cap_and_merge", "bind_tools_max_count_from_settings", "bind_tools_pad_to_cap_from_settings", + "provider_safe_tool_name", + "provider_safe_tools", + "resolve_tool_name", ] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py new file mode 100644 index 000000000..1f43b61cb --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py @@ -0,0 +1,128 @@ +"""Provider-safe tool names for native ``bind_tools``. + +OpenAI-compatible endpoints (OpenAI, Azure OpenAI, LiteLLM proxies in front of +them) require every function name to match ``^[A-Za-z0-9_-]{1,64}$`` and reject +the whole request when one name does not, so a single long registry name +(``_``) makes every bound tool unusable. + +Such tools are bound under a legal alias that lives only in the request: a +``tool_call`` returned under it is mapped back to the real name before CUGA +turns it into code, so the approval check and the sandbox see the real name. + +- :func:`provider_safe_tool_name` — the alias (legal names come back unchanged). +- :func:`provider_safe_tools` — the list ``_safe_bind`` hands to ``bind_tools``. +- :func:`resolve_tool_name` — maps a returned alias back to the real name. +""" + +from __future__ import annotations + +import hashlib +import re +from typing import Any, Collection, Dict, Sequence, Set, Tuple + +from langchain_core.tools import BaseTool +from loguru import logger + +__all__ = [ + "ORIGINAL_TOOL_NAME_KEY", + "PROVIDER_TOOL_NAME_RE", + "provider_safe_tool_name", + "provider_safe_tools", + "resolve_tool_name", +] + +PROVIDER_TOOL_NAME_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$") +_INVALID = re.compile(r"[^A-Za-z0-9_-]") +_ALIAS_SUFFIX = re.compile(r"_[0-9a-f]{8}$") +_MAX_LEN = 64 +_DIGEST_LEN = 8 + +# Key in an aliased tool copy's ``metadata`` that holds the real tool name. +ORIGINAL_TOOL_NAME_KEY = "cuga_original_tool_name" + +# ``_safe_bind`` runs on every model call: log each (real, alias) pair once per process. +_logged_aliases: Set[Tuple[str, str]] = set() + + +def provider_safe_tool_name(name: str) -> str: + """Legal names come back unchanged; anything else gets a deterministic legal alias. + + The alias is the name with illegal characters replaced by ``_``, cut to 55 + characters, then ``_`` and the first 8 hex digits of the SHA-1 of the + *original* name, so ``a.b`` and ``a_b`` stay distinct. Plain truncation is + not safe: long names of one family share their prefixes, and a merged name + makes the model call the wrong tool without any error. + """ + if PROVIDER_TOOL_NAME_RE.fullmatch(name): + return name + raw = name.encode("utf-8", "surrogatepass") # a lone surrogate must not raise here + digest = hashlib.sha1(raw, usedforsecurity=False).hexdigest()[:_DIGEST_LEN] + stem = _INVALID.sub("_", name)[: _MAX_LEN - 1 - _DIGEST_LEN] or "tool" + return f"{stem}_{digest}" + + +def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: + """Return ``tools`` with every illegal tool name replaced by its alias. + + Aliased entries are copies (``model_copy``) that carry the real name in + ``metadata["cuga_original_tool_name"]``; the shared tool objects are never + mutated, so prompts, policies and tracking keep seeing real names. When no + name needs an alias, ``tools`` itself is returned. + + Raises ``RuntimeError`` when tools with different names would be bound + under the same name: the model would call one believing it is the other. + """ + aliases: Dict[int, str] = {} + for i, tool in enumerate(tools): + if isinstance(tool, BaseTool): + alias = provider_safe_tool_name(tool.name) + if alias != tool.name: + aliases[i] = alias + if not aliases: + return tools + + bound_as: Dict[str, str] = {} + safe_tools = [] + for i, tool in enumerate(tools): + name = getattr(tool, "name", None) + final = aliases.get(i, name) + if isinstance(final, str): + other = bound_as.setdefault(final, name) + if other != name: + raise RuntimeError( + f"bind_tools: tools {other!r} and {name!r} would both be bound as {final!r}. " + "Rename one of them." + ) + if i in aliases: + tool = tool.model_copy( + update={ + "name": final, + "metadata": {**(tool.metadata or {}), ORIGINAL_TOOL_NAME_KEY: name}, + } + ) + if (name, final) not in _logged_aliases: + _logged_aliases.add((name, final)) + logger.info("bind_tools: binding tool {!r} as {!r} (provider-safe name)", name, final) + safe_tools.append(tool) + return safe_tools + + +def resolve_tool_name(name: str, tool_names: Collection[str]) -> str: + """Map a tool name the model returned to the tool it stands for. + + ``tool_names`` are the real names that can run (the execution context's + keys). A name that is the alias of one of them comes back as that real + name; any other name comes back unchanged. + + Raises ``RuntimeError`` when the name is ambiguous: it is the alias of more + than one tool, or of one tool and the real name of another. + """ + if not _ALIAS_SUFFIX.search(name): + return name # every alias ends in "_" + 8 hex digits + owners = [n for n in tool_names if n != name and provider_safe_tool_name(n) == name] + if not owners: + return name + if len(owners) > 1 or name in tool_names: + clash = owners if len(owners) > 1 else [name, *owners] + raise RuntimeError(f"Tool name {name!r} is ambiguous: it matches tools {sorted(clash)!r}.") + return owners[0] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py index bcd7090e2..22bb74f79 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py @@ -19,6 +19,7 @@ BindToolsUnsupportedError, apply_bind_tools_cap_and_merge, bind_tools_max_count_from_settings, + provider_safe_tools, ) @@ -48,7 +49,11 @@ def _safe_bind(model: BaseChatModel, tools: List[StructuredTool]) -> Runnable: guard it would be caught by the cap's deliberate ``except RuntimeError: raise`` in ``resolve_model_with_bind_tools`` and crash ``call_model`` instead of degrading to the unbound (code-act) model. + + Names OpenAI-compatible providers reject are bound under an alias + (:func:`provider_safe_tools`); the response side maps it back. """ + tools = provider_safe_tools(tools) try: return model.bind_tools(tools) except NotImplementedError: diff --git a/tests/unit/test_provider_safe_tool_names.py b/tests/unit/test_provider_safe_tool_names.py new file mode 100644 index 000000000..26d037320 --- /dev/null +++ b/tests/unit/test_provider_safe_tool_names.py @@ -0,0 +1,421 @@ +"""Provider-safe tool names for native tool binding (#879). + +OpenAI-compatible providers reject a request whose tools array has a name +outside ``^[A-Za-z0-9_-]{1,64}$``. ``_safe_bind`` binds such tools under an +alias that lives only in the request: a tool call returned under it is mapped +back to the real name before CUGA turns it into code, so the approval check +and the sandbox see the real name. +""" + +from __future__ import annotations + +import hashlib +from types import SimpleNamespace +from typing import Any, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.tools import StructuredTool +from langgraph.checkpoint.memory import MemorySaver +from loguru import logger +from pydantic import BaseModel + +from cuga.backend.cuga_graph.nodes.cuga_agent_core.execution.code_extraction import ( + extract_code_from_model_response, +) +from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import EXECUTION_OUTPUT_PREFIX +from cuga.backend.cuga_graph.nodes.cuga_agent_core.policy.tool_approval_handler import ToolApprovalHandler +from cuga.backend.cuga_graph.nodes.cuga_lite import cuga_lite_graph +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import AgentGraphAdapter +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools import tool_names +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import ( + PROVIDER_TOOL_NAME_RE, + provider_safe_tool_name, + provider_safe_tools, + resolve_tool_name, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import ( + _safe_bind, + resolve_model_with_bind_tools, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.providers.langchain import DirectLangChainToolsProvider +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module +from cuga.backend.cuga_graph.policy.agent import PolicyAgent +from cuga.backend.cuga_graph.policy.configurable import PolicyConfigurable +from cuga.backend.cuga_graph.policy.models import ToolApproval + +pytestmark = pytest.mark.unit + +# Synthetic names only. LONG has 76 characters, a length registry names +# (``_``) reach. +LONG = "acme_inventory_service_reconcile_warehouse_stock_levels_with_supplier_ledger" +ALIAS = provider_safe_tool_name(LONG) + + +class _Args(BaseModel): + value: int + + +async def _reconcile(value: int) -> dict: + return {"reconciled": value} + + +def _tool(name: str, **kwargs: Any) -> StructuredTool: + return StructuredTool( + name=name, description="Reconcile stock levels.", args_schema=_Args, coroutine=_reconcile, **kwargs + ) + + +def _adapter(tools_context: dict) -> AgentGraphAdapter: + return AgentGraphAdapter( + tracker=MagicMock(), + base_callbacks=[], + task_todos_ref=[], + tools_context_ref={}, + base_tool_provider=None, + tools_context=tools_context, + ) + + +class _BindModel(BaseChatModel): + """Chat model that records what ``bind_tools`` receives.""" + + bound: Optional[List[Any]] = None + + @property + def _llm_type(self) -> str: + return "bind-recorder" + + def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult: + return ChatResult(generations=[ChatGeneration(message=AIMessage(content="ok"))]) + + def bind_tools(self, tools: Any, **kwargs: Any): + return self.model_copy(update={"bound": list(tools)}) + + +@pytest.fixture(autouse=True) +def _pin_bind_tools_cap(): + # As in test_bind_tools_safe_fallback.py: keep an ambient DYNACONF cap out of these tests. + with patch( + "cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools.bind_tools_max_count_from_settings", + return_value=128, + ): + yield + + +# ── The alias ────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("name", ["a", "x" * 64, "get-weather", "A1_b-2"]) +def test_legal_names_are_unchanged(name): + assert provider_safe_tool_name(name) == name + + +def test_long_name_gets_a_deterministic_64_char_alias(): + name = "app_" + "b" * 61 # 65 characters + expected = f"{name[:55]}_{hashlib.sha1(name.encode()).hexdigest()[:8]}" + assert provider_safe_tool_name(name) == expected + assert provider_safe_tool_name(name) == expected + assert len(expected) == 64 and PROVIDER_TOOL_NAME_RE.fullmatch(expected) + + +def test_names_sharing_a_long_prefix_get_distinct_aliases(): + # Truncation alone would merge each pair, and the model would call the wrong tool. + prefix = "geo_get_percentage_of_countries_grouped_by_population_band_and_by_" + region, continent = prefix + "region", prefix + "continent" + assert region[:64] == continent[:64] + assert provider_safe_tool_name(region) != provider_safe_tool_name(continent) + + alpha, beta = "x" * 57 + "_alpha_tool", "x" * 57 + "_beta_tool" + assert provider_safe_tool_name(alpha) != provider_safe_tool_name(beta) + + +def test_illegal_characters_get_an_alias_of_the_original_name(): + assert ( + provider_safe_tool_name("my.tool:v2") == "my_tool_v2_" + hashlib.sha1(b"my.tool:v2").hexdigest()[:8] + ) + # The digest is of the original name, so "a.b" does not collapse into the legal "a_b". + assert provider_safe_tool_name("a_b") == "a_b" + assert provider_safe_tool_name("a.b") != "a_b" + + +@pytest.mark.parametrize("name", [LONG, "my.tool:v2", "", "x" * 64]) +def test_alias_is_idempotent(name): + alias = provider_safe_tool_name(name) + assert provider_safe_tool_name(alias) == alias + + +def test_lone_surrogate_in_a_name_does_not_break_aliasing(): + # A JSON "\ud800" escape can put one in an MCP tool name. + name = "tool_\ud800_name" + alias = provider_safe_tool_name(name) + assert PROVIDER_TOOL_NAME_RE.fullmatch(alias) + assert resolve_tool_name(alias, {name: _reconcile}) == name + + +# ── Binding ──────────────────────────────────────────────────────────────── + + +def test_safe_bind_binds_an_aliased_copy_and_leaves_the_tool_alone(): + tool = _tool(LONG, metadata={"owner": "acme"}) + [bound] = _safe_bind(_BindModel(), [tool]).bound + + assert bound.name == ALIAS + assert bound.description == tool.description + assert bound.args_schema is tool.args_schema + assert bound.coroutine is tool.coroutine + assert bound.metadata == {"owner": "acme", "cuga_original_tool_name": LONG} + assert (tool.name, tool.metadata) == (LONG, {"owner": "acme"}) + + +def test_provider_request_carries_only_legal_names(): + from langchain_openai import ChatOpenAI + + model = ChatOpenAI(model="gpt-4o-mini", api_key="test-key") # nothing is sent + bound = _safe_bind(model, [_tool(LONG), _tool("get_weather")]) + + names = [tool["function"]["name"] for tool in bound.kwargs["tools"]] + assert names == [ALIAS, "get_weather"] + assert all(PROVIDER_TOOL_NAME_RE.fullmatch(name) for name in names) + + +def test_legal_names_are_bound_as_is_and_not_logged(): + tools = [_tool("get_weather"), _tool("get-forecast")] + logs: list[str] = [] + handler_id = logger.add(lambda message: logs.append(str(message)), level="INFO") + try: + bound = _safe_bind(_BindModel(), tools).bound + finally: + logger.remove(handler_id) + + assert provider_safe_tools(tools) is tools + assert all(b is t for b, t in zip(bound, tools, strict=True)) + assert not [line for line in logs if "provider-safe" in line] + + +def test_each_alias_is_logged_once_per_process(monkeypatch): + monkeypatch.setattr(tool_names, "_logged_aliases", set()) + logs: list[str] = [] + handler_id = logger.add(lambda message: logs.append(str(message)), level="INFO") + try: + for _ in range(3): + _safe_bind(_BindModel(), [_tool(LONG)]) + finally: + logger.remove(handler_id) + + assert len([line for line in logs if ALIAS in line]) == 1 + + +@pytest.mark.asyncio +async def test_two_tools_bound_under_one_name_fail_loudly(): + tools = [_tool(LONG), _tool(ALIAS)] # the alias of one is the real name of the other + with pytest.raises(RuntimeError, match="would both be bound") as excinfo: + _safe_bind(_BindModel(), tools) + assert LONG in str(excinfo.value) and ALIAS in str(excinfo.value) + + provider = AsyncMock() + provider.get_all_tools = AsyncMock(return_value=tools) + with pytest.raises(RuntimeError, match="would both be bound"): # not degraded to the unbound model + await resolve_model_with_bind_tools( + _BindModel(), + configurable={ + "cuga_lite_bind_tools_mode": "tools", + "cuga_lite_bind_tools_tool_names": [LONG, ALIAS], + }, + tools_context_ref={}, + tool_provider=provider, + ) + + +# ── Mapping a returned alias back ────────────────────────────────────────── + + +def test_resolve_tool_name_maps_an_alias_back(): + ctx = {LONG: _reconcile, "get_weather": _reconcile} + assert resolve_tool_name(ALIAS, ctx) == LONG + assert resolve_tool_name("get_weather", ctx) == "get_weather" + assert resolve_tool_name("not_a_tool", ctx) == "not_a_tool" + + +def test_resolve_tool_name_refuses_ambiguous_names(monkeypatch): + with pytest.raises(RuntimeError, match="ambiguous"): + resolve_tool_name(ALIAS, {LONG: _reconcile, ALIAS: _reconcile}) + + monkeypatch.setattr(tool_names, "provider_safe_tool_name", lambda name: "shared_0123abcd") + with pytest.raises(RuntimeError, match="ambiguous"): + resolve_tool_name("shared_0123abcd", {"tool_one": _reconcile, "tool_two": _reconcile}) + + +@pytest.mark.parametrize( + "response", + [ + AIMessage(content="", tool_calls=[{"name": ALIAS, "args": {"value": 7}, "id": "call_1"}]), + SimpleNamespace( # raw OpenAI shape + content="", + tool_calls=None, + additional_kwargs={ + "tool_calls": [{"id": "call_1", "function": {"name": ALIAS, "arguments": '{"value": 7}'}}] + }, + ), + ], + ids=["tool_calls", "additional_kwargs"], +) +def test_tool_call_under_the_alias_becomes_code_that_names_the_real_tool(response): + content, _ = _adapter({LONG: _reconcile}).normalize_response(response) + assert f"await {LONG}(value=7)" in content + assert ALIAS not in content + + +@pytest.mark.asyncio +async def test_tool_use_failed_recovery_maps_the_alias_back(): + error = Exception( + "Error code: 400 - {'error': {'message': 'Failed to call a function. tool_use_failed', " + f"'failed_generation': '{{\"name\": \"{ALIAS}\", \"arguments\": {{\"value\": 7}}}}'}}}}" + ) + + class _Rejecting: + async def ainvoke(self, messages, config=None): + raise error + + response = await _adapter({LONG: _reconcile}).ainvoke_model(_Rejecting(), [], {}) + assert f"await {LONG}(value=7)" in response.content + assert ALIAS not in response.content + + +@pytest.mark.asyncio +async def test_approval_policy_on_the_real_name_stops_a_call_made_under_the_alias(): + """The call runs as the real tool, so a ToolApproval on the real name must still stop it. + + ToolApproval matches real names in the code text: code naming the alias would not match. + """ + adapter = _adapter({LONG: _reconcile}) + content, reasoning = adapter.normalize_response( + AIMessage(content="", tool_calls=[{"name": ALIAS, "args": {"value": 7}, "id": "call_1"}]) + ) + code = extract_code_from_model_response(content, reasoning) + + policy = ToolApproval( + id="approve-reconcile", + name="Reconcile needs approval", + description="Approval gate.", + required_tools=[LONG], + ) + policy_agent = PolicyAgent.__new__(PolicyAgent) + policy_agent.storage = SimpleNamespace(list_policies=AsyncMock(return_value=[policy])) + state = SimpleNamespace( + chat_messages=[HumanMessage(content="Reconcile.")], step_count=0, cuga_lite_metadata={} + ) + + with ( + patch.object(PolicyConfigurable, "from_config", return_value=SimpleNamespace(agent=policy_agent)), + patch.object( + PolicyConfigurable, + "create_context_from_state", + return_value=SimpleNamespace(user_input="Reconcile."), + ), + ): + command = await ToolApprovalHandler.check_and_create_approval_interrupt( + adapter, state, code, content, {} + ) + + # None means "no policy matched", or any error on the way: the handler fails open (#880). + assert command is not None + assert command.update["hitl_action"] is not None + assert command.update["script"] == code + + +# ── End to end ───────────────────────────────────────────────────────────── + + +class _ProviderLikeModel: + """Scripted model that, like a provider, only knows the names it was bound with. + + It calls the tool for a request and answers once it sees the tool's output, + so an extra model call cannot put it out of step. + """ + + def __init__(self): + self.bound_names: list[list[str]] = [] + + def bind_tools(self, tools, **kwargs): + self.bound_names.append([tool.name for tool in tools]) + return self + + async def ainvoke(self, messages, config=None, **kwargs): + last = messages[-1] + text = last.get("content") if isinstance(last, dict) else getattr(last, "content", "") + if str(text).startswith(EXECUTION_OUTPUT_PREFIX): + return AIMessage(content="Reconciled.") + call = {"name": self.bound_names[-1][0], "args": {"value": 7}, "id": "call_1"} + return AIMessage(content="", tool_calls=[call]) + + +@pytest.fixture +def _reset_budgets(): + yield + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + + +@pytest.mark.asyncio +async def test_native_call_under_the_alias_runs_the_real_tool_on_every_run(monkeypatch, _reset_budgets): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + monkeypatch.setattr(settings.advanced_features, "cuga_lite_nl_auto_continue", False, raising=False) + adapters: list[AgentGraphAdapter] = [] + + class _RecordingAdapter(cuga_lite_graph.AgentGraphAdapter): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + adapters.append(self) + + monkeypatch.setattr(cuga_lite_graph, "AgentGraphAdapter", _RecordingAdapter) + + executed: list[int] = [] + + async def reconcile(value: int) -> dict: + executed.append(value) + return {"reconciled": value} + + tool = StructuredTool( + name=LONG, description="Reconcile stock levels.", args_schema=_Args, coroutine=reconcile + ) + model = _ProviderLikeModel() + graph = cuga_lite_graph.create_cuga_lite_graph( + model=model, + tool_provider=DirectLangChainToolsProvider(tools=[tool], app_name="acme"), + apps_list=[], + thread_id="alias-e2e", + ).compile(checkpointer=MemorySaver()) + config = { + "configurable": { + "thread_id": "alias-e2e", + "enable_todos": False, + "track_tool_calls": True, + "reflection_enabled": False, + "pre_execute_verify_enabled": False, + "cuga_lite_bind_tools_mode": "tools", + "cuga_lite_bind_tools_tool_names": [LONG], + } + } + + # The execution context outlives a run, so run the same graph twice. + for run in (1, 2): + result = await graph.ainvoke( + cuga_lite_graph.CugaLiteState(chat_messages=[HumanMessage(content="Reconcile the stock.")]), + config=config, + ) + assert executed == [7] * run + assert result["final_answer"] == "Reconciled." + assert [call["name"] for call in result["tool_calls"]] == [LONG] + transcript = " ".join(str(message.content) for message in result["chat_messages"]) + assert f"await {LONG}(value=7)" in transcript and ALIAS not in transcript + + # The alias reached the provider and nothing else. + assert model.bound_names and all(names == [ALIAS] for names in model.bound_names) + assert LONG in adapters[0]._tools_context and ALIAS not in adapters[0]._tools_context From e7a25c60175a23e0edfe571e74cb3f81d7bf5a1f Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Mon, 5 Oct 2026 10:22:01 +0300 Subject: [PATCH 10/18] fix(cuga-lite): map tool aliases back at the decode boundary - Map provider-safe aliases back to real names once, in AgentGraphAdapter.normalize_response, instead of inside two code builders. The one seam covers native tool_calls, the tool_use_failed recovery, code the model writes itself and its reasoning, so a model that types an alias still runs the real tool, behind the approval check. - response_utils.py and ainvoke_model are back to main, and the aliased copy no longer carries the unread cuga_original_tool_name metadata. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/graph_adapter.py | 12 ++- .../nodes/cuga_lite/adapter/response_utils.py | 25 +---- .../nodes/cuga_lite/bind_tools/__init__.py | 4 +- .../nodes/cuga_lite/bind_tools/tool_names.py | 77 ++++++++-------- tests/unit/test_harmony_decode_boundary.py | 1 + tests/unit/test_provider_safe_tool_names.py | 92 ++++++++++++------- 6 files changed, 110 insertions(+), 101 deletions(-) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 41ce0b7d4..3891b4c79 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -24,12 +24,12 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.prepare_node import create_prepare_tools_and_apps_node from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.response_utils import ( clean_empty_response_retry_meta, - extract_code_from_failed_tool_call, extract_code_from_response_tool_calls, ) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import create_sandbox_node from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import resolve_model_with_bind_tools from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.find_tools import _first_user_message_text +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools import resolve_tool_names from cuga.backend.cuga_graph.nodes.cuga_lite.nl_auto_continue_classifier import ( BLOCKED_CLAIM_CORRECTION, BlockedClaimEvidence, @@ -37,6 +37,7 @@ normalize_assistant_text, ) from cuga.backend.cuga_graph.utils.token_counter import clamp_watsonx_completion_for_messages +from cuga.backend.llm.errors import extract_code_from_tool_use_failed from cuga.config import settings @@ -132,7 +133,7 @@ async def ainvoke_model(self, bound: Any, messages: list, invoke_config: dict) - clamp_watsonx_completion_for_messages(bound, messages) return await bound.ainvoke(messages, config=invoke_config) except Exception as exc: - code = extract_code_from_failed_tool_call(exc, self._tools_context) + code = extract_code_from_tool_use_failed(exc) if code: logger.warning( "Model attempted tool call without tools bound (tool_use_failed). " @@ -172,9 +173,11 @@ async def resolve_bind_tools( def normalize_response(self, response: Any) -> Tuple[str, Optional[str]]: # Harmony framing is removed here, at the decode boundary, so every # downstream surface inherits clean text (see the base implementation). + # Provider-safe tool aliases (bind_tools/tool_names.py) are mapped back + # to real names here for the same reason. content = strip_harmony_tokens(normalize_assistant_text(response.content)) if not content: - tool_code = extract_code_from_response_tool_calls(response, self._tools_context) + tool_code = extract_code_from_response_tool_calls(response) if tool_code: logger.warning("Empty content with tool_calls detected; recovering tool call as Python code") content = tool_code @@ -182,7 +185,8 @@ def normalize_response(self, response: Any) -> Tuple[str, Optional[str]]: reasoning = normalize_assistant_text( additional_kwargs.get("reasoning_content") or additional_kwargs.get("reasoning") ) - return content, reasoning + tools = self._tools_context + return resolve_tool_names(content, tools), resolve_tool_names(reasoning, tools) def on_response_processed( self, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py index ba4fafecc..220c74100 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/response_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -from typing import Any, Collection, Dict, Optional +from typing import Any, Dict, Optional from langchain_core.messages import HumanMessage @@ -12,9 +12,7 @@ EMPTY_RESPONSE_CORRECTION_KEY, EXECUTION_OUTPUT_PREFIX, ) -from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools import resolve_tool_name from cuga.backend.cuga_graph.nodes.cuga_lite.reflection.verify_result import VERIFY_BLOCKED_PREFIX -from cuga.backend.llm.errors import failed_gen_to_code, parse_tool_use_failed_generation def clean_empty_response_retry_meta(meta: Optional[Dict[str, Any]]) -> Dict[str, Any]: @@ -47,13 +45,8 @@ def tool_call_kwarg_literal(value: Any) -> str: return repr(value) -def extract_code_from_response_tool_calls(response: Any, tool_names: Collection[str] = ()) -> Optional[str]: - """Recover fenced Python from AIMessage.tool_calls when content is empty. - - A provider-safe alias the model was bound with is mapped back to its real - name in ``tool_names``, so the code (and the approval check that reads it) - names the real tool. - """ +def extract_code_from_response_tool_calls(response: Any) -> Optional[str]: + """Recover fenced Python from AIMessage.tool_calls when content is empty.""" tool_calls = getattr(response, "tool_calls", None) or ( getattr(response, "additional_kwargs", None) or {} ).get("tool_calls") @@ -74,20 +67,8 @@ def extract_code_from_response_tool_calls(response: Any, tool_names: Collection[ if not name: return None - if isinstance(name, str): - name = resolve_tool_name(name, tool_names) args_str = ", ".join( f"{k}={tool_call_kwarg_literal(v)}" for k, v in (args if isinstance(args, dict) else {}).items() ) return f"```python\nresult = await {name}({args_str})\nprint(result)\n```" - - -def extract_code_from_failed_tool_call(err: Any, tool_names: Collection[str] = ()) -> Optional[str]: - """``llm.errors.extract_code_from_tool_use_failed``, with an alias mapped back as above.""" - failed_gen = parse_tool_use_failed_generation(err) - if not isinstance(failed_gen, dict): - return None - if isinstance(failed_gen.get("name"), str): - failed_gen = {**failed_gen, "name": resolve_tool_name(failed_gen["name"], tool_names)} - return failed_gen_to_code(failed_gen) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py index a2c85b10c..39c8c1d6e 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py @@ -14,7 +14,7 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import ( provider_safe_tool_name, provider_safe_tools, - resolve_tool_name, + resolve_tool_names, ) __all__ = [ @@ -24,5 +24,5 @@ "bind_tools_pad_to_cap_from_settings", "provider_safe_tool_name", "provider_safe_tools", - "resolve_tool_name", + "resolve_tool_names", ] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py index 1f43b61cb..65e355f5b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py @@ -5,40 +5,35 @@ the whole request when one name does not, so a single long registry name (``_``) makes every bound tool unusable. -Such tools are bound under a legal alias that lives only in the request: a -``tool_call`` returned under it is mapped back to the real name before CUGA -turns it into code, so the approval check and the sandbox see the real name. - -- :func:`provider_safe_tool_name` — the alias (legal names come back unchanged). -- :func:`provider_safe_tools` — the list ``_safe_bind`` hands to ``bind_tools``. -- :func:`resolve_tool_name` — maps a returned alias back to the real name. +Such tools are bound under a legal alias that lives only in the request: +``_safe_bind`` binds :func:`provider_safe_tools`, and the decode boundary +(``AgentGraphAdapter.normalize_response``) maps every alias in a reply back +with :func:`resolve_tool_names`, so code, the approval check and the sandbox +only see real names. """ from __future__ import annotations import hashlib import re -from typing import Any, Collection, Dict, Sequence, Set, Tuple +from typing import Any, Collection, Dict, List, Sequence, Set, Tuple from langchain_core.tools import BaseTool from loguru import logger __all__ = [ - "ORIGINAL_TOOL_NAME_KEY", "PROVIDER_TOOL_NAME_RE", "provider_safe_tool_name", "provider_safe_tools", - "resolve_tool_name", + "resolve_tool_names", ] PROVIDER_TOOL_NAME_RE = re.compile(r"^[A-Za-z0-9_-]{1,64}$") _INVALID = re.compile(r"[^A-Za-z0-9_-]") -_ALIAS_SUFFIX = re.compile(r"_[0-9a-f]{8}$") _MAX_LEN = 64 _DIGEST_LEN = 8 - -# Key in an aliased tool copy's ``metadata`` that holds the real tool name. -ORIGINAL_TOOL_NAME_KEY = "cuga_original_tool_name" +# A whole token shaped like an alias: up to 55 legal characters, "_", 8 hex digits. +_ALIAS = re.compile(r"(? str: def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: """Return ``tools`` with every illegal tool name replaced by its alias. - Aliased entries are copies (``model_copy``) that carry the real name in - ``metadata["cuga_original_tool_name"]``; the shared tool objects are never - mutated, so prompts, policies and tracking keep seeing real names. When no - name needs an alias, ``tools`` itself is returned. + Aliased entries are copies (``model_copy``); the shared tool objects are + never mutated, so prompts, policies and tracking keep seeing real names. + When no name needs an alias, ``tools`` itself is returned. Raises ``RuntimeError`` when tools with different names would be bound under the same name: the model would call one believing it is the other. @@ -94,12 +88,7 @@ def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: "Rename one of them." ) if i in aliases: - tool = tool.model_copy( - update={ - "name": final, - "metadata": {**(tool.metadata or {}), ORIGINAL_TOOL_NAME_KEY: name}, - } - ) + tool = tool.model_copy(update={"name": final}) if (name, final) not in _logged_aliases: _logged_aliases.add((name, final)) logger.info("bind_tools: binding tool {!r} as {!r} (provider-safe name)", name, final) @@ -107,22 +96,30 @@ def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: return safe_tools -def resolve_tool_name(name: str, tool_names: Collection[str]) -> str: - """Map a tool name the model returned to the tool it stands for. +def resolve_tool_names(text: str, tool_names: Collection[str]) -> str: + """Replace every alias in ``text`` with the real tool name it stands for. - ``tool_names`` are the real names that can run (the execution context's - keys). A name that is the alias of one of them comes back as that real - name; any other name comes back unchanged. + ``text`` is a model reply, code or a single tool name; ``tool_names`` are + the real names that can run (the execution context's keys). Anything that + is not the alias of one of them, real names included, is left as it is. - Raises ``RuntimeError`` when the name is ambiguous: it is the alias of more - than one tool, or of one tool and the real name of another. + Raises ``RuntimeError`` when an alias is ambiguous: it stands for more than + one tool, or for one tool while being the real name of another. """ - if not _ALIAS_SUFFIX.search(name): - return name # every alias ends in "_" + 8 hex digits - owners = [n for n in tool_names if n != name and provider_safe_tool_name(n) == name] - if not owners: - return name - if len(owners) > 1 or name in tool_names: - clash = owners if len(owners) > 1 else [name, *owners] - raise RuntimeError(f"Tool name {name!r} is ambiguous: it matches tools {sorted(clash)!r}.") - return owners[0] + if not text or not _ALIAS.search(text): + return text + owners: Dict[str, List[str]] = {} + for real in tool_names: + alias = provider_safe_tool_name(real) + if alias != real: + owners.setdefault(alias, []).append(real) + + def _real(match: re.Match) -> str: + alias = match.group(0) + reals = owners.get(alias, []) + clash = [*reals, alias] if reals and alias in tool_names else reals + if len(clash) > 1: + raise RuntimeError(f"Tool name {alias!r} is ambiguous: it matches tools {sorted(clash)!r}.") + return reals[0] if reals else alias + + return _ALIAS.sub(_real, text) diff --git a/tests/unit/test_harmony_decode_boundary.py b/tests/unit/test_harmony_decode_boundary.py index 3ef4f96c7..f334298f0 100644 --- a/tests/unit/test_harmony_decode_boundary.py +++ b/tests/unit/test_harmony_decode_boundary.py @@ -12,6 +12,7 @@ def test_lite_boundary_sanitizes_content(): from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import AgentGraphAdapter adapter = object.__new__(AgentGraphAdapter) + adapter._tools_context = {} content, _ = AgentGraphAdapter.normalize_response(adapter, AIMessage(content=RAW)) assert "<|" not in content assert "The total is 42" in content diff --git a/tests/unit/test_provider_safe_tool_names.py b/tests/unit/test_provider_safe_tool_names.py index 26d037320..667b55f9a 100644 --- a/tests/unit/test_provider_safe_tool_names.py +++ b/tests/unit/test_provider_safe_tool_names.py @@ -2,9 +2,9 @@ OpenAI-compatible providers reject a request whose tools array has a name outside ``^[A-Za-z0-9_-]{1,64}$``. ``_safe_bind`` binds such tools under an -alias that lives only in the request: a tool call returned under it is mapped -back to the real name before CUGA turns it into code, so the approval check -and the sandbox see the real name. +alias that lives only in the request; the decode boundary (``normalize_response``) +maps every alias in a reply back to the real name, so code, the approval check +and the sandbox only see real names. """ from __future__ import annotations @@ -35,7 +35,7 @@ PROVIDER_TOOL_NAME_RE, provider_safe_tool_name, provider_safe_tools, - resolve_tool_name, + resolve_tool_names, ) from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import ( _safe_bind, @@ -69,6 +69,14 @@ def _tool(name: str, **kwargs: Any) -> StructuredTool: ) +def _code(name: str) -> str: + return f"```python\nresult = await {name}(value=7)\nprint(result)\n```" + + +def _tool_call(name: str) -> AIMessage: + return AIMessage(content="", tool_calls=[{"name": name, "args": {"value": 7}, "id": "call_1"}]) + + def _adapter(tools_context: dict) -> AgentGraphAdapter: return AgentGraphAdapter( tracker=MagicMock(), @@ -153,7 +161,7 @@ def test_lone_surrogate_in_a_name_does_not_break_aliasing(): name = "tool_\ud800_name" alias = provider_safe_tool_name(name) assert PROVIDER_TOOL_NAME_RE.fullmatch(alias) - assert resolve_tool_name(alias, {name: _reconcile}) == name + assert resolve_tool_names(alias, {name: _reconcile}) == name # ── Binding ──────────────────────────────────────────────────────────────── @@ -167,8 +175,8 @@ def test_safe_bind_binds_an_aliased_copy_and_leaves_the_tool_alone(): assert bound.description == tool.description assert bound.args_schema is tool.args_schema assert bound.coroutine is tool.coroutine - assert bound.metadata == {"owner": "acme", "cuga_original_tool_name": LONG} - assert (tool.name, tool.metadata) == (LONG, {"owner": "acme"}) + assert bound.metadata == {"owner": "acme"} + assert tool.name == LONG def test_provider_request_carries_only_legal_names(): @@ -230,29 +238,36 @@ async def test_two_tools_bound_under_one_name_fail_loudly(): ) -# ── Mapping a returned alias back ────────────────────────────────────────── +# ── Mapping aliases back ─────────────────────────────────────────────────── -def test_resolve_tool_name_maps_an_alias_back(): +def test_resolve_tool_names_maps_every_alias_back_and_nothing_else(): ctx = {LONG: _reconcile, "get_weather": _reconcile} - assert resolve_tool_name(ALIAS, ctx) == LONG - assert resolve_tool_name("get_weather", ctx) == "get_weather" - assert resolve_tool_name("not_a_tool", ctx) == "not_a_tool" + assert resolve_tool_names(ALIAS, ctx) == LONG + assert resolve_tool_names("get_weather", ctx) == "get_weather" + assert resolve_tool_names("", ctx) == "" + # Only whole alias tokens change: not a longer token, the real name, or an unknown alias shape. + text = f"r = await {ALIAS}(value=7) # not {ALIAS}x, {LONG} or deadbeef_0123abcd" + assert resolve_tool_names(text, ctx) == text.replace(f"await {ALIAS}(", f"await {LONG}(") -def test_resolve_tool_name_refuses_ambiguous_names(monkeypatch): + +def test_resolve_tool_names_refuses_ambiguous_aliases(monkeypatch): with pytest.raises(RuntimeError, match="ambiguous"): - resolve_tool_name(ALIAS, {LONG: _reconcile, ALIAS: _reconcile}) + resolve_tool_names(ALIAS, {LONG: _reconcile, ALIAS: _reconcile}) monkeypatch.setattr(tool_names, "provider_safe_tool_name", lambda name: "shared_0123abcd") with pytest.raises(RuntimeError, match="ambiguous"): - resolve_tool_name("shared_0123abcd", {"tool_one": _reconcile, "tool_two": _reconcile}) + resolve_tool_names("shared_0123abcd", {"tool_one": _reconcile, "tool_two": _reconcile}) + + +# ── The decode boundary ──────────────────────────────────────────────────── @pytest.mark.parametrize( - "response", + "reply", [ - AIMessage(content="", tool_calls=[{"name": ALIAS, "args": {"value": 7}, "id": "call_1"}]), + _tool_call(ALIAS), SimpleNamespace( # raw OpenAI shape content="", tool_calls=None, @@ -260,17 +275,24 @@ def test_resolve_tool_name_refuses_ambiguous_names(monkeypatch): "tool_calls": [{"id": "call_1", "function": {"name": ALIAS, "arguments": '{"value": 7}'}}] }, ), + AIMessage(content=_code(ALIAS)), # the model wrote the alias into its own code ], - ids=["tool_calls", "additional_kwargs"], + ids=["tool_call", "raw_tool_call", "code_in_reply"], ) -def test_tool_call_under_the_alias_becomes_code_that_names_the_real_tool(response): - content, _ = _adapter({LONG: _reconcile}).normalize_response(response) +def test_reply_naming_the_alias_decodes_to_the_real_tool(reply): + content, _ = _adapter({LONG: _reconcile}).normalize_response(reply) assert f"await {LONG}(value=7)" in content assert ALIAS not in content +def test_reasoning_naming_the_alias_decodes_to_the_real_tool(): + reply = AIMessage(content="Done.", additional_kwargs={"reasoning_content": f"Call {ALIAS} next."}) + _, reasoning = _adapter({LONG: _reconcile}).normalize_response(reply) + assert reasoning == f"Call {LONG} next." + + @pytest.mark.asyncio -async def test_tool_use_failed_recovery_maps_the_alias_back(): +async def test_tool_use_failed_recovery_decodes_to_the_real_tool(): error = Exception( "Error code: 400 - {'error': {'message': 'Failed to call a function. tool_use_failed', " f"'failed_generation': '{{\"name\": \"{ALIAS}\", \"arguments\": {{\"value\": 7}}}}'}}}}" @@ -280,21 +302,23 @@ class _Rejecting: async def ainvoke(self, messages, config=None): raise error - response = await _adapter({LONG: _reconcile}).ainvoke_model(_Rejecting(), [], {}) - assert f"await {LONG}(value=7)" in response.content - assert ALIAS not in response.content + adapter = _adapter({LONG: _reconcile}) + content, _ = adapter.normalize_response(await adapter.ainvoke_model(_Rejecting(), [], {})) + assert f"await {LONG}(value=7)" in content + assert ALIAS not in content +@pytest.mark.parametrize( + "reply", [_tool_call(ALIAS), AIMessage(content=_code(ALIAS))], ids=["tool_call", "code"] +) @pytest.mark.asyncio -async def test_approval_policy_on_the_real_name_stops_a_call_made_under_the_alias(): +async def test_approval_policy_on_the_real_name_stops_a_call_made_under_the_alias(reply): """The call runs as the real tool, so a ToolApproval on the real name must still stop it. ToolApproval matches real names in the code text: code naming the alias would not match. """ adapter = _adapter({LONG: _reconcile}) - content, reasoning = adapter.normalize_response( - AIMessage(content="", tool_calls=[{"name": ALIAS, "args": {"value": 7}, "id": "call_1"}]) - ) + content, reasoning = adapter.normalize_response(reply) code = extract_code_from_model_response(content, reasoning) policy = ToolApproval( @@ -337,7 +361,8 @@ class _ProviderLikeModel: so an extra model call cannot put it out of step. """ - def __init__(self): + def __init__(self, reply_with: str): + self.reply_with = reply_with self.bound_names: list[list[str]] = [] def bind_tools(self, tools, **kwargs): @@ -349,8 +374,8 @@ async def ainvoke(self, messages, config=None, **kwargs): text = last.get("content") if isinstance(last, dict) else getattr(last, "content", "") if str(text).startswith(EXECUTION_OUTPUT_PREFIX): return AIMessage(content="Reconciled.") - call = {"name": self.bound_names[-1][0], "args": {"value": 7}, "id": "call_1"} - return AIMessage(content="", tool_calls=[call]) + name = self.bound_names[-1][0] + return AIMessage(content=_code(name)) if self.reply_with == "code" else _tool_call(name) @pytest.fixture @@ -361,8 +386,9 @@ def _reset_budgets(): tracker_module._block_tool_call_budget_context.set(None) +@pytest.mark.parametrize("reply_with", ["tool_call", "code"]) @pytest.mark.asyncio -async def test_native_call_under_the_alias_runs_the_real_tool_on_every_run(monkeypatch, _reset_budgets): +async def test_call_under_the_alias_runs_the_real_tool_on_every_run(monkeypatch, _reset_budgets, reply_with): from cuga.config import settings monkeypatch.setattr(settings.policy, "enabled", False, raising=False) @@ -385,7 +411,7 @@ async def reconcile(value: int) -> dict: tool = StructuredTool( name=LONG, description="Reconcile stock levels.", args_schema=_Args, coroutine=reconcile ) - model = _ProviderLikeModel() + model = _ProviderLikeModel(reply_with) graph = cuga_lite_graph.create_cuga_lite_graph( model=model, tool_provider=DirectLangChainToolsProvider(tools=[tool], app_name="acme"), From e653ea122add7fbdbb41801fbb7fbd98dc743669 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Wed, 7 Oct 2026 11:41:00 +0300 Subject: [PATCH 11/18] refactor(cuga-lite): trim provider-safe tool names to the two seams - Import from bind_tools.tool_names directly, so bind_tools/__init__.py is unchanged. - One-line comments at the two call sites, and shorter docstrings. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/graph_adapter.py | 5 +- .../nodes/cuga_lite/bind_tools/__init__.py | 11 +---- .../nodes/cuga_lite/bind_tools/tool_names.py | 46 +++++++------------ .../nodes/cuga_lite/helpers/bind_tools.py | 7 +-- tests/unit/test_provider_safe_tool_names.py | 1 - 5 files changed, 22 insertions(+), 48 deletions(-) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 3891b4c79..28b5324c1 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -29,7 +29,7 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import create_sandbox_node from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import resolve_model_with_bind_tools from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.find_tools import _first_user_message_text -from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools import resolve_tool_names +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import resolve_tool_names from cuga.backend.cuga_graph.nodes.cuga_lite.nl_auto_continue_classifier import ( BLOCKED_CLAIM_CORRECTION, BlockedClaimEvidence, @@ -173,8 +173,7 @@ async def resolve_bind_tools( def normalize_response(self, response: Any) -> Tuple[str, Optional[str]]: # Harmony framing is removed here, at the decode boundary, so every # downstream surface inherits clean text (see the base implementation). - # Provider-safe tool aliases (bind_tools/tool_names.py) are mapped back - # to real names here for the same reason. + # Provider-safe tool aliases are mapped back to real names for the same reason. content = strip_harmony_tokens(normalize_assistant_text(response.content)) if not content: tool_code = extract_code_from_response_tool_calls(response) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py index 39c8c1d6e..811051ad0 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/__init__.py @@ -1,8 +1,7 @@ """bind_tools cap + shortlist machinery for cuga_lite. Keeps cuga_lite_graph.py focused on orchestration. See :mod:`.cap` for the -provider-safe cap and shortlister flow, and :mod:`.tool_names` for -provider-safe tool names. +provider-safe cap and shortlister flow. """ from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.cap import ( @@ -11,18 +10,10 @@ bind_tools_max_count_from_settings, bind_tools_pad_to_cap_from_settings, ) -from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import ( - provider_safe_tool_name, - provider_safe_tools, - resolve_tool_names, -) __all__ = [ "BindToolsUnsupportedError", "apply_bind_tools_cap_and_merge", "bind_tools_max_count_from_settings", "bind_tools_pad_to_cap_from_settings", - "provider_safe_tool_name", - "provider_safe_tools", - "resolve_tool_names", ] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py index 65e355f5b..216270b1f 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/bind_tools/tool_names.py @@ -1,15 +1,11 @@ """Provider-safe tool names for native ``bind_tools``. -OpenAI-compatible endpoints (OpenAI, Azure OpenAI, LiteLLM proxies in front of -them) require every function name to match ``^[A-Za-z0-9_-]{1,64}$`` and reject -the whole request when one name does not, so a single long registry name -(``_``) makes every bound tool unusable. - -Such tools are bound under a legal alias that lives only in the request: -``_safe_bind`` binds :func:`provider_safe_tools`, and the decode boundary -(``AgentGraphAdapter.normalize_response``) maps every alias in a reply back -with :func:`resolve_tool_names`, so code, the approval check and the sandbox -only see real names. +OpenAI-compatible APIs reject a request when any tool name is outside +``^[A-Za-z0-9_-]{1,64}$``, so one long registry name (``_``) makes +every bound tool unusable. Such tools are bound under an alias that exists only +in the request: ``_safe_bind`` binds :func:`provider_safe_tools`, and +``AgentGraphAdapter.normalize_response`` maps aliases in a reply back with +:func:`resolve_tool_names`, so code, approval checks and the sandbox see real names. """ from __future__ import annotations @@ -40,13 +36,12 @@ def provider_safe_tool_name(name: str) -> str: - """Legal names come back unchanged; anything else gets a deterministic legal alias. + """Return ``name`` if providers accept it, else a deterministic legal alias. The alias is the name with illegal characters replaced by ``_``, cut to 55 - characters, then ``_`` and the first 8 hex digits of the SHA-1 of the - *original* name, so ``a.b`` and ``a_b`` stay distinct. Plain truncation is - not safe: long names of one family share their prefixes, and a merged name - makes the model call the wrong tool without any error. + characters, plus ``_`` and 8 hex digits of the SHA-1 of the original name. + Truncation alone would merge long names that share a prefix, and the model + would call the wrong tool. """ if PROVIDER_TOOL_NAME_RE.fullmatch(name): return name @@ -57,14 +52,10 @@ def provider_safe_tool_name(name: str) -> str: def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: - """Return ``tools`` with every illegal tool name replaced by its alias. + """Return ``tools`` with illegal names replaced by aliases, or ``tools`` itself if none are. - Aliased entries are copies (``model_copy``); the shared tool objects are - never mutated, so prompts, policies and tracking keep seeing real names. - When no name needs an alias, ``tools`` itself is returned. - - Raises ``RuntimeError`` when tools with different names would be bound - under the same name: the model would call one believing it is the other. + Aliased entries are copies; the shared tool objects are never mutated. + Raises ``RuntimeError`` when two tools would be bound under the same name. """ aliases: Dict[int, str] = {} for i, tool in enumerate(tools): @@ -97,14 +88,11 @@ def provider_safe_tools(tools: Sequence[Any]) -> Sequence[Any]: def resolve_tool_names(text: str, tool_names: Collection[str]) -> str: - """Replace every alias in ``text`` with the real tool name it stands for. - - ``text`` is a model reply, code or a single tool name; ``tool_names`` are - the real names that can run (the execution context's keys). Anything that - is not the alias of one of them, real names included, is left as it is. + """Replace each alias in ``text`` with the real name in ``tool_names`` it stands for. - Raises ``RuntimeError`` when an alias is ambiguous: it stands for more than - one tool, or for one tool while being the real name of another. + ``text`` can be a reply, code or a single tool name; everything else in it, + real names included, is left as it is. Raises ``RuntimeError`` when an alias + could mean more than one tool. """ if not text or not _ALIAS.search(text): return text diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py index 22bb74f79..c2dae5c25 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/helpers/bind_tools.py @@ -19,8 +19,8 @@ BindToolsUnsupportedError, apply_bind_tools_cap_and_merge, bind_tools_max_count_from_settings, - provider_safe_tools, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import provider_safe_tools def _record_bind_tools_degraded(reason: str) -> None: @@ -49,11 +49,8 @@ def _safe_bind(model: BaseChatModel, tools: List[StructuredTool]) -> Runnable: guard it would be caught by the cap's deliberate ``except RuntimeError: raise`` in ``resolve_model_with_bind_tools`` and crash ``call_model`` instead of degrading to the unbound (code-act) model. - - Names OpenAI-compatible providers reject are bound under an alias - (:func:`provider_safe_tools`); the response side maps it back. """ - tools = provider_safe_tools(tools) + tools = provider_safe_tools(tools) # alias names providers reject; normalize_response maps them back try: return model.bind_tools(tools) except NotImplementedError: diff --git a/tests/unit/test_provider_safe_tool_names.py b/tests/unit/test_provider_safe_tool_names.py index 667b55f9a..113e92b25 100644 --- a/tests/unit/test_provider_safe_tool_names.py +++ b/tests/unit/test_provider_safe_tool_names.py @@ -126,7 +126,6 @@ def test_long_name_gets_a_deterministic_64_char_alias(): name = "app_" + "b" * 61 # 65 characters expected = f"{name[:55]}_{hashlib.sha1(name.encode()).hexdigest()[:8]}" assert provider_safe_tool_name(name) == expected - assert provider_safe_tool_name(name) == expected assert len(expected) == 64 and PROVIDER_TOOL_NAME_RE.fullmatch(expected) From 9cfae1a7128d606f5fa6c610756b32e37d691d57 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:01:23 +0300 Subject: [PATCH 12/18] test(cuga-lite): record the feature-off CodeAct snapshot Golden for the function-calling rework: with no execution mode configured a CodeAct run must stay byte-identical (final state, every message, every outbound turn, variable bookkeeping). Recorded on this tree before any change; also pins that find_tools listings stay out of CodeAct variables. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../tests/snapshots/codeact_feature_off.json | 94 ++++++++++ .../tests/test_feature_off_snapshot.py | 165 ++++++++++++++++++ 2 files changed, 259 insertions(+) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/snapshots/codeact_feature_off.json create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_feature_off_snapshot.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/snapshots/codeact_feature_off.json b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/snapshots/codeact_feature_off.json new file mode 100644 index 000000000..01f5b1d89 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/snapshots/codeact_feature_off.json @@ -0,0 +1,94 @@ +{ + "calls": [ + [ + "echo", + 3 + ] + ], + "outbound": [ + [ + { + "content": "", + "role": "system" + }, + { + "content": "use the tools", + "role": "user" + } + ], + [ + { + "content": "", + "role": "system" + }, + { + "content": "use the tools", + "role": "user" + }, + { + "content": "```python\nr = await echo(value=3)\nprint(r)\n```", + "role": "assistant" + }, + { + "content": "Execution output:\n3\n\n## New Variables Created:\n# Variables Summary\n\n## r\n- Type: int\n- Items: 1\n- Description: Created during code execution\n- Value Preview: 3\n\n\n## Available Variables\n\n# Variables Summary\n\n## r\n- Type: int\n- Items: 1\n- Description: Created during code execution\n- Value Preview: 3\n\nYou can use these variables directly by their names.", + "role": "user" + } + ] + ], + "state": { + "chat_messages": [ + { + "content": "use the tools", + "type": "HumanMessage" + }, + { + "content": "```python\nr = await echo(value=3)\nprint(r)\n```", + "type": "AIMessage" + }, + { + "content": "Execution output:\n3\n\n## New Variables Created:\n# Variables Summary\n\n## r\n- Type: int\n- Items: 1\n- Description: Created during code execution\n- Value Preview: 3\n", + "type": "HumanMessage" + }, + { + "content": "It printed 3.", + "type": "AIMessage" + } + ], + "cuga_lite_metadata": {}, + "execution_complete": true, + "final_answer": "It printed 3.", + "mcp_few_shot_messages": [], + "metrics": {}, + "pi": "", + "reflection_apps": [], + "reflection_enable_find_tools": false, + "reflection_skills_enabled": false, + "reflection_skills_prompt_section": "", + "script": null, + "service_scope": { + "instance_id": "", + "tenant_id": "" + }, + "step_count": 3, + "task_todos": null, + "tool_budget_exhausted": false, + "tool_calls": [], + "tool_calls_used_run": 1, + "tool_calls_used_thread": 1, + "tools_prepared": true, + "variable_counter_state": 0, + "variable_creation_order": [ + "r" + ], + "variables_storage": { + "r": { + "count_items": 1, + "description": "Created during code execution", + "type": "int", + "value": 3 + } + }, + "verify_revise_streak": 0, + "verify_revise_total": 0 + } +} diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_feature_off_snapshot.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_feature_off_snapshot.py new file mode 100644 index 000000000..cbe033d11 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_feature_off_snapshot.py @@ -0,0 +1,165 @@ +"""Feature-off golden: a CodeAct run is byte-identical with function-calling code present. + +``snapshots/codeact_feature_off.json`` was recorded on the tree *before* the +function-calling execution path was routed through the sandbox. Every field of +the final graph state, every message, every outbound turn the model saw and the +sandbox's variable bookkeeping must still match. Timestamps are the only thing +normalised away. + +Re-record (only for an intentional CodeAct change, never to make this pass): +``CUGA_SNAPSHOT_WRITE=1 pytest ``. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.tools import StructuredTool +from langgraph.checkpoint.memory import MemorySaver + +from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState, create_cuga_lite_graph +from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import CodeExecutor +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module + +pytestmark = pytest.mark.unit + +SNAPSHOT = Path(__file__).parent / "snapshots" / "codeact_feature_off.json" +_VOLATILE_KEYS = {"created_at"} +CALLS: list = [] + + +class _Bound: + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + self.seen: list = [] + + def bind_tools(self, tools, **kwargs): + return _Bound(self) + + async def ainvoke(self, messages, config=None, **kwargs): + self.seen.append(list(messages)) + return self._responses.pop(0) + + +def _provider(*tools): + provider = MagicMock() + provider.get_all_tools = AsyncMock(return_value=list(tools)) + provider.get_apps = AsyncMock(return_value=[]) + provider.get_tools = AsyncMock(return_value=[]) + provider.app_name = "test_app" + return provider + + +def _echo_tool(): + async def echo(value: int) -> int: + """Echo a value.""" + CALLS.append(("echo", value)) + return value + + return StructuredTool.from_function(coroutine=echo, name="echo", description="Echo a value.") + + +@pytest.fixture(autouse=True) +def _quiet(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", False, raising=False) + CALLS.clear() + yield + CALLS.clear() + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + tracker_module._block_tool_call_cap_override_context.set(None) + + +def _normalize(value: Any) -> Any: + if isinstance(value, dict): + return {k: _normalize(v) for k, v in value.items() if k not in _VOLATILE_KEYS} + if isinstance(value, (list, tuple)): + return [_normalize(v) for v in value] + if hasattr(value, "type") and hasattr(value, "content"): # a message object + out = {"type": type(value).__name__, "content": value.content} + if getattr(value, "tool_calls", None): + out["tool_calls"] = value.tool_calls + if getattr(value, "tool_call_id", None): + out["tool_call_id"] = value.tool_call_id + return out + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def _outbound(messages: list) -> list: + """The dict-serialised CodeAct turns, system prompt replaced by a marker.""" + out = [] + for m in messages: + assert isinstance(m, dict), "CodeAct sends dict turns" + out.append( + {"role": m["role"], "content": "" if m["role"] == "system" else m["content"]} + ) + return out + + +@pytest.mark.asyncio +async def test_codeact_run_matches_the_recorded_snapshot(): + model = _ScriptedModel( + [ + AIMessage(content="```python\nr = await echo(value=3)\nprint(r)\n```"), + AIMessage(content="It printed 3."), + ] + ) + graph = create_cuga_lite_graph( + model=model, tool_provider=_provider(_echo_tool()), apps_list=[], thread_id="t" + ).compile(checkpointer=MemorySaver()) + result = await graph.ainvoke( + CugaLiteState(chat_messages=[HumanMessage(content="use the tools")]), + config={"configurable": {"thread_id": "codeact-snapshot", "enable_todos": False}}, + ) + + prompt = result["prepared_prompt"] + assert "```python" in prompt + assert all(turn[0]["content"] == prompt for turn in model.seen), "system prompt is the prepared prompt" + + state = _normalize({k: v for k, v in result.items() if k != "prepared_prompt"}) + snapshot = { + "calls": [list(c) for c in CALLS], + "state": state, + "outbound": [_outbound(turn) for turn in model.seen], + } + + if os.environ.get("CUGA_SNAPSHOT_WRITE"): + SNAPSHOT.write_text(json.dumps(snapshot, indent=2, sort_keys=True) + "\n") + pytest.skip(f"snapshot written to {SNAPSHOT}") + + recorded = json.loads(SNAPSHOT.read_text()) + assert snapshot == recorded, "CodeAct behaviour changed with the feature off" + + +@pytest.mark.asyncio +async def test_find_tools_listing_is_still_stripped_from_codeact_variables(): + """The executor drops find_tools listing markdown from a block's variables by default.""" + state = CugaLiteState(chat_messages=[]) + listing = "# Found 1 Matching Tool(s)\n**Query:** q\n## 1. `x`" + + async def lookup(): + return listing + + code = "tools = await lookup()\nkept = 1\n" + _, new_vars = await CodeExecutor.eval_with_tools_async(code=code, _locals={"lookup": lookup}, state=state) + + assert new_vars == {"kept": 1} From 1e9678bf7fb769058d5447ac4df1cb433915348b Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:31:45 +0300 Subject: [PATCH 13/18] feat(cuga-lite): run function-calling tool calls through the sandbox A native turn's tool_calls are translated into a short internal block (adapter/fc_actions.py) that the existing sandbox runs, and the block's variables are read back into one ToolMessage per call id. One execution pipeline instead of two: tool approval, the VERIFY gate, variables, output trimming, the block timeout and the budgets now apply to a native call as they apply to CodeAct code. The tool_exec node and its refuse-to-start approval guard are gone; the paused assistant turn keeps its tool_calls through an approval and the approved block resumes into the sandbox. The model side is untouched (bind_tools in, tool_calls in, ToolMessages out). Feature off is byte-identical: every new parameter defaults to the old behaviour and the recorded CodeAct snapshot passes unchanged. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../cuga_agent_core/graph/shared_graph.py | 19 +- .../policy/tool_approval_handler.py | 15 +- .../nodes/cuga_lite/adapter/fc_actions.py | 382 +++++++++++++ .../nodes/cuga_lite/adapter/graph_adapter.py | 138 ++--- .../nodes/cuga_lite/adapter/prepare_node.py | 6 + .../nodes/cuga_lite/adapter/sandbox_node.py | 94 +++- .../nodes/cuga_lite/adapter/tool_exec_node.py | 379 ------------- .../nodes/cuga_lite/cuga_lite_graph.py | 5 - .../cuga_lite/executors/code_executor.py | 9 +- .../tests/test_execution_mode_config.py | 42 +- .../tests/test_fc_call_model_hook.py | 162 +++--- .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 125 ++++- .../cuga_lite/tests/test_tool_exec_node.py | 516 ------------------ src/cuga/backend/cuga_graph/policy/storage.py | 10 - .../policy/tests/test_storage_strict_list.py | 34 -- .../backend/cuga_graph/utils/agent_loop.py | 20 +- src/cuga/settings.toml | 2 +- tests/unit/test_agent_loop_fc_events.py | 30 +- .../unit/test_fc_sdk_fail_closed_approval.py | 111 ---- tests/unit/test_fc_sdk_two_turn.py | 9 +- 20 files changed, 832 insertions(+), 1276 deletions(-) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py delete mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py delete mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py delete mode 100644 src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py delete mode 100644 tests/unit/test_fc_sdk_fail_closed_approval.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py index 0deaf09cd..d238d9ca1 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/graph/shared_graph.py @@ -4,8 +4,8 @@ START → prepare --Command--> call_model ↔ execute (loop) → END -CugaLite may add a fourth node, ``tool_exec`` (native function-calling -mode), which routes back into ``call_model`` with a Command. +CugaLite's function-calling mode uses the same three nodes: a native tool-call +turn is translated into a block the execute node runs. Both CugaLite and CugaSupervisor share this structure. The nodes themselves are provided by the caller (produced by adapter factories), so the graph @@ -18,7 +18,7 @@ from __future__ import annotations -from typing import Callable, Optional, Type +from typing import Callable, Type from langgraph.graph import START, StateGraph @@ -32,7 +32,6 @@ def build_agent_graph( prepare_node: Callable, call_model_node: Callable, execute_node: Callable, - tool_exec_node: Optional[Callable] = None, ) -> StateGraph: """Wire and return an UNCOMPILED 3-node agent StateGraph. @@ -46,12 +45,6 @@ def build_agent_graph( call_model_node: Async node function for the call_model step (use ``create_call_model_node`` from ``shared_nodes.py``). execute_node: Async node function for the execute/sandbox step. - tool_exec_node: Optional node that executes native ``tool_calls`` and - replies with ``ToolMessage``s (CugaLite function-calling mode). When - supplied it is added as ``"tool_exec"``; it routes back to - ``call_model`` (or to END on error) with a Command, and ``call_model`` - only routes to it in that mode, so it is dormant on every CodeAct run. - Omitted by the Supervisor graph. Returns: An uncompiled ``StateGraph``. Call ``.compile(checkpointer=...)`` @@ -68,10 +61,4 @@ def build_agent_graph( # Execute node returns a state update (not Command); loop back for the NL answer. graph.add_edge(adapter.execute_node_name, "call_model") - if tool_exec_node is not None: - # No static edge back: the node routes with Command (call_model on a normal - # exit, END on an error exit), so a terminal error never schedules one more - # model turn the way a static edge would. - graph.add_node("tool_exec", tool_exec_node) - return graph diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/policy/tool_approval_handler.py b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/policy/tool_approval_handler.py index 59091f1af..a9c5620f3 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/policy/tool_approval_handler.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_agent_core/policy/tool_approval_handler.py @@ -119,8 +119,15 @@ async def check_and_create_approval_interrupt( code: str, content: str, config: dict = None, + *, + ai_message: Optional[AIMessage] = None, ) -> Optional[Command]: - """Check if code requires approval and create an interrupt if needed.""" + """Check if code requires approval and create an interrupt if needed. + + ``ai_message`` is the assistant turn to persist with the interrupt; the + default is a text message carrying ``content``. A function-calling turn + passes its own ``AIMessage`` so the ``tool_calls`` survive the pause. + """ from cuga.backend.cuga_graph.policy.configurable import PolicyConfigurable try: @@ -173,7 +180,7 @@ async def check_and_create_approval_interrupt( adapter.set_metadata(state, approval_metadata) return ToolApprovalHandler._create_approval_interrupt( - adapter, state, code, content, preview_lines + adapter, state, code, content, preview_lines, ai_message=ai_message ) except Exception as e: @@ -187,6 +194,8 @@ def _create_approval_interrupt( code: str, content: str, preview_lines: List[str], + *, + ai_message: Optional[AIMessage] = None, ) -> Command: """Create an interrupt Command for tool approval.""" from cuga.backend.cuga_graph.nodes.human_in_the_loop.followup_model import create_tool_approval_action @@ -223,7 +232,7 @@ def _create_approval_interrupt( ) updated_messages, error_message = _core_append_with_step_limit( - adapter, state, [AIMessage(content=content)] + adapter, state, [ai_message if ai_message is not None else AIMessage(content=content)] ) if error_message: return _core_create_error_command(adapter, updated_messages, error_message, state.step_count) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py new file mode 100644 index 000000000..60b342fe3 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py @@ -0,0 +1,382 @@ +"""Function-calling actions: native ``tool_calls`` in, a sandbox block out, ``ToolMessage``s back. + +``cuga_lite_execution_mode = "function_calling"`` has no executor of its own. The +calls the model attached to its ``AIMessage`` are translated into a short Python +block that the existing sandbox runs, so tool approval, the pre-execute VERIFY +gate, variables, output trimming, the block timeout and the tool-call budgets +apply to a native call exactly as they apply to a CodeAct block. The block's +variables are then read back into one ``ToolMessage`` per call id, so the +transcript the model sees stays a native one. + +The block is internal: never shown to the model, never persisted in chat +history, never sent to a provider. + +What keeps the tool side identical to a direct call: + +- a tool whose name is a Python identifier is called by that name, so the + approval text scan and VERIFY see the real tool name; any other name goes + through the injected ``_fc_tools`` lookup, budget-counted like the rest; +- arguments are always passed as ``**{...}`` so keys like ``from`` work, and + only JSON values are encoded — providers deliver JSON, so nothing is lost; +- each call sits in its own ``try``: one failing call never aborts its siblings, + and the exception is kept as a marker the reply builder formats; +- the result is kept through ``_fc_keep`` so a non-JSON value survives as the + very text the reply carries instead of being dropped from the variables; +- calls that cannot run (no name, unknown tool, unparsable arguments, + provider-rejected, deferred under step discipline) get their reply decided + at planning time and emit no code. +""" + +from __future__ import annotations + +import json +import keyword +import math +import re +from typing import Any, Dict, List, Optional, Tuple + +from langchain_core.messages import ToolMessage + +FC_PENDING_KEY = "fc_pending" +"""Metadata key: the turn's call plan while the sandbox runs the block; cleared on every sandbox exit.""" + +FC_KEEP_KEY = "_fc_keep" +FC_TOOLS_KEY = "_fc_tools" +FC_BUDGET_EXC_KEY = "_fc_budget_exc" +RESULT_VAR_PREFIX = "tool_result_" +_ERROR_KEY = "__cuga_fc_error__" +_TIMEOUT_PREFIX = "Error during execution: Execution timed out" + +DEFERRED_CALL_MESSAGE = ( + "Deferred: step discipline is on, so only the first tool call of a turn is attempted. " + "Read that result, then re-issue this call in your next turn if you still need it." +) + +TRUNCATION_MARKER = "\n... [result truncated: {limit} characters of output remain for this batch]" + +NOT_EXECUTED_REPLY = "Not executed: tool execution failed before this call ran ({reason})." + + +# ── Pieces shared with the direct-call contract ────────────────────────────── + + +def _tool_call_parts(call: Any) -> tuple[str, str, Any]: + """``(id, name, args)`` from a LangChain tool_call dict or a raw OpenAI-shaped one.""" + if not isinstance(call, dict): + return "", "", None + fn = call.get("function") if isinstance(call.get("function"), dict) else {} + call_id = str(call.get("id") or "") + name = str(call.get("name") or fn.get("name") or "") + args = call.get("args") if "args" in call else fn.get("arguments") + return call_id, name, args + + +def _coerce_args(args: Any) -> tuple[Optional[Dict[str, Any]], Optional[str]]: + """Return ``(kwargs, error)``. Strings are JSON-decoded; anything else is rejected.""" + if args is None: + return {}, None + if isinstance(args, dict): + return args, None + if isinstance(args, str): + text = args.strip() + if not text: + return {}, None + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + return ( + None, + f"Could not parse tool arguments as JSON ({exc}). Re-issue the call with valid JSON arguments.", + ) + if isinstance(parsed, dict): + return parsed, None + return None, f"Tool arguments must be a JSON object, got {type(parsed).__name__}." + return None, f"Tool arguments must be an object, got {type(args).__name__}." + + +def _stringify(result: Any) -> str: + if isinstance(result, str): + return result + try: + return json.dumps(result, ensure_ascii=False, default=str) + except Exception: + return str(result) + + +class _BatchOutputBudget: + """One turn's batch of results shares ``execution_output_max_length``, like one block. + + Per-result truncation alone would let N calls return N x the limit; this keeps + the whole batch under the ceiling the sandbox applies to a block's output. + """ + + def __init__(self, limit: Any): + self.limit = int(limit or 0) + self.remaining = self.limit + + def take(self, text: str) -> str: + if not self.limit: + return text + if len(text) > self.remaining: + cap = max(self.remaining, 0) + text = text[:cap] + TRUNCATION_MARKER.format(limit=cap) + self.remaining = max(self.remaining - len(text), 0) + return text + + +def _error_message(text: str, *, call_id: str, name: str) -> ToolMessage: + return ToolMessage(content=text, tool_call_id=call_id, name=name or "unknown", status="error") + + +def fc_keep(value: Any) -> Any: + """Keep a tool result as a block variable without losing it. + + JSON-shaped values are stored as they are. Anything else is stored as the + text the reply carries (``_stringify``), so the variables filter never drops + a result and the reply is the same text a direct call would have produced. + """ + try: + json.dumps(value, ensure_ascii=False) + return value + except Exception: + return _stringify(value) + + +# ── Planning: tool_calls -> block ──────────────────────────────────────────── + + +def _literal(value: Any) -> str: + """Python source for a JSON value. Raises ``ValueError`` for anything that is not JSON.""" + if value is None or isinstance(value, bool): + return repr(value) + if isinstance(value, int): + return repr(value) + if isinstance(value, float): + if math.isnan(value) or math.isinf(value): + raise ValueError("non-JSON number") + return repr(value) + if isinstance(value, str): + return json.dumps(value, ensure_ascii=False) + if isinstance(value, (list, tuple)): + return "[" + ", ".join(_literal(v) for v in value) + "]" + if isinstance(value, dict): + items = ", ".join( + f"{json.dumps(str(k), ensure_ascii=False)}: {_literal(v)}" for k, v in value.items() + ) + return "{" + items + "}" + raise ValueError(f"unsupported argument type {type(value).__name__}") + + +def _result_var(call_id: str, taken: set) -> str: + base = RESULT_VAR_PREFIX + re.sub(r"\W", "_", call_id) + name, n = base, 1 + while name in taken: + n += 1 + name = f"{base}_{n}" + taken.add(name) + return name + + +def _pre_decided(content: str) -> Dict[str, str]: + return {"content": content, "status": "error"} + + +_ERROR_MARKER_SRC = ( + "{" + json.dumps(_ERROR_KEY) + ': str(_fc_e), "kind": "TypeError" if isinstance(_fc_e, TypeError) ' + 'else ("budget" if isinstance(_fc_e, ' + FC_BUDGET_EXC_KEY + ') else "error")}' +) + + +def plan_tool_calls( + calls: List[Any], + invalid: List[Any], + tools_context: Dict[str, Any], + *, + one_per_step: bool, +) -> Tuple[Optional[str], List[Dict[str, Any]]]: + """Translate one turn's calls into ``(block, plan)``. + + ``plan`` has one JSON-serialisable entry per issued id, in order: + ``{"id", "name", "var", "lookup", "reply"}``. ``var`` names the block + variable that will hold the call's result; ``reply`` is set instead when + the call was answered here and emits no code. ``block`` is ``None`` when + nothing is left to execute. + """ + plan: List[Dict[str, Any]] = [] + lines: List[str] = [] + taken: set = set() + + for index, call in enumerate(calls): + call_id, name, raw_args = _tool_call_parts(call) + call_id = call_id or f"call_{index}" + entry: Dict[str, Any] = {"id": call_id, "name": name, "var": None, "lookup": False, "reply": None} + plan.append(entry) + + if not name: + entry["reply"] = _pre_decided("Tool call carried no tool name.") + continue + if one_per_step and index > 0: + # By position, not by success: an error on the first call is the result + # the model must read before it decides the next one. + entry["reply"] = _pre_decided(DEFERRED_CALL_MESSAGE) + continue + fn = tools_context.get(name) + if fn is None or not callable(fn): + known = ", ".join(sorted(k for k in tools_context if not k.startswith("_"))) + entry["reply"] = _pre_decided( + f"Unknown tool '{name}'. Choose one of the provided tools: {known or '(none)'}." + ) + continue + kwargs, arg_error = _coerce_args(raw_args) + if arg_error: + entry["reply"] = _pre_decided(arg_error) + continue + try: + args_src = _literal(kwargs) + except ValueError as exc: + entry["reply"] = _pre_decided( + f"Tool arguments could not be encoded ({exc}). Re-issue the call with plain JSON values." + ) + continue + + var = _result_var(call_id, taken) + entry["var"] = var + if name.isidentifier() and not keyword.iskeyword(name): + target = name + else: + entry["lookup"] = True + target = f"{FC_TOOLS_KEY}[{json.dumps(name, ensure_ascii=False)}]" + lines += [ + "try:", + f" {var} = {FC_KEEP_KEY}(await {target}(**{args_src}))", + "except Exception as _fc_e:", + f" {var} = {_ERROR_MARKER_SRC}", + ] + + for index, bad in enumerate(invalid): + call_id, name, _ = _tool_call_parts(bad) + reason = (bad.get("error") if isinstance(bad, dict) else None) or "malformed tool call" + plan.append( + { + "id": call_id or f"invalid_{index}", + "name": name, + "var": None, + "lookup": False, + "reply": _pre_decided( + f"The provider could not parse this tool call ({reason}). Re-issue it with valid arguments." + ), + } + ) + + block = "\n".join(lines) + "\n" if lines else None + return block, plan + + +# ── Execution side: what the sandbox injects and reads back ────────────────── + + +def fc_pending_plan(adapter: Any, state: Any) -> Optional[List[Dict[str, Any]]]: + """The plan a function-calling turn left for this sandbox step, else ``None``.""" + plan = (adapter.get_metadata(state) or {}).get(FC_PENDING_KEY) + return [dict(e) for e in plan] if plan else None + + +def metadata_without_plan(adapter: Any, state: Any) -> Dict[str, Any]: + """Turn metadata with the pending plan cleared — every sandbox exit writes this back.""" + return {k: v for k, v in (adapter.get_metadata(state) or {}).items() if k != FC_PENDING_KEY} + + +def fc_context_overlay(tools_context: Dict[str, Any], plan: List[Dict[str, Any]]) -> Dict[str, Any]: + """Names the block needs beyond the tools: the keeper, the budget exception and the lookup table. + + All ``_``-prefixed, so the executor neither counts them as tools nor keeps + them as variables. Lookup entries are budget-counted here because the + executor only wraps top-level callables. + """ + from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( + ToolCallBudgetExceeded, + counted_tool_call, + ) + + overlay: Dict[str, Any] = {FC_KEEP_KEY: fc_keep, FC_BUDGET_EXC_KEY: ToolCallBudgetExceeded} + lookup = { + e["name"]: counted_tool_call(tools_context[e["name"]]) + for e in plan + if e.get("lookup") and e.get("name") in tools_context + } + if lookup: + overlay[FC_TOOLS_KEY] = lookup + return overlay + + +def _format_error(name: str, marker: Dict[str, Any]) -> str: + text = str(marker.get(_ERROR_KEY, "")) + kind = marker.get("kind") + if kind == "TypeError": + return f"Tool '{name}' rejected these arguments: {text}. Check the parameter names and types." + if kind == "budget": + return text + return f"Tool '{name}' failed: {text}" + + +def _unreached_reason(output: str, timeout: Any) -> str: + """Why a planned call left no result: the block timed out, or something outside the per-call guard.""" + text = (output or "").strip() + if text.startswith(_TIMEOUT_PREFIX): + return f"Tool '{{name}}' timed out after {timeout}s. Try a narrower call or a different tool." + first_line = text.splitlines()[0] if text else "no result was recorded" + return NOT_EXECUTED_REPLY.format(reason=first_line.replace("{", "{{").replace("}", "}}")) + + +def replies_from_execution( + plan: List[Dict[str, Any]], + new_vars: Dict[str, Any], + output: str, + *, + output_limit: Any, + timeout: Any, +) -> Tuple[List[ToolMessage], Dict[str, Any], List[str]]: + """One ``ToolMessage`` per issued id, read back from the block's variables. + + Returns ``(replies, variables_to_keep, variables_to_drop)``: error markers + are formatted into their reply and dropped from the variables. + """ + budget = _BatchOutputBudget(output_limit) + keep = dict(new_vars) + drop: List[str] = [] + unreached = _unreached_reason(output, timeout) + replies: List[ToolMessage] = [] + for entry in plan: + call_id, name = entry["id"], entry.get("name") or "" + pre = entry.get("reply") + if pre: + replies.append(_error_message(pre["content"], call_id=call_id, name=name)) + continue + var = entry.get("var") + if var is not None and var in new_vars: + value = new_vars[var] + if isinstance(value, dict) and _ERROR_KEY in value: + keep.pop(var, None) + drop.append(var) + replies.append(_error_message(_format_error(name, value), call_id=call_id, name=name)) + else: + replies.append( + ToolMessage(content=budget.take(_stringify(value)), tool_call_id=call_id, name=name) + ) + continue + replies.append(_error_message(unreached.format(name=name), call_id=call_id, name=name)) + return replies, keep, drop + + +def replies_without_execution(plan: List[Dict[str, Any]], reason: str) -> List[ToolMessage]: + """Every id answered when the block did not run (VERIFY revise, a failure outside the block). + + Pre-decided replies are delivered as they are; every call that would have + run gets ``reason``. + """ + out: List[ToolMessage] = [] + for entry in plan: + pre = entry.get("reply") + text = pre["content"] if pre else reason + out.append(_error_message(text, call_id=entry["id"], name=entry.get("name") or "")) + return out diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 6b4034220..2038aa3d2 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -20,6 +20,7 @@ format_current_plan_section, format_task_todos_system_block, ) +from cuga.backend.cuga_graph.nodes.cuga_agent_core.policy.tool_approval_handler import ToolApprovalHandler from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( EMPTY_RESPONSE_CORRECTION, EMPTY_RESPONSE_CORRECTION_KEY, @@ -29,20 +30,26 @@ enforce_step_limit, ) from cuga.backend.cuga_graph.utils.harmony import contains_harmony_tokens, strip_harmony_tokens +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import ( + FC_PENDING_KEY, + plan_tool_calls, + replies_without_execution, +) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.prepare_node import create_prepare_tools_and_apps_node from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.response_utils import ( clean_empty_response_retry_meta, extract_code_from_response_tool_calls, ) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import create_sandbox_node -from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import create_tool_exec_node from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.bind_tools import ( _bind_tools_mode_from_settings, resolve_model_with_bind_tools, ) from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( EXECUTION_MODE_FUNCTION_CALLING, + STEP_DISCIPLINE_ONE_TOOL_PER_STEP, resolve_execution_mode, + resolve_step_discipline, resolved_runtime_model_name, runtime_defaults_for_model, ) @@ -65,23 +72,18 @@ "Code is never run here. Issue the tool call natively instead, or give the final answer as plain text." ) -FC_TOOL_APPROVAL_UNSUPPORTED = ( - "Function-calling mode does not support tool-approval policies yet. An enabled tool-approval " - "policy exists, so this run was stopped before any tool ran. Use cuga_lite_execution_mode = " - '"codeact" for this agent, or disable the policy.' -) - -FC_TOOL_APPROVAL_UNVERIFIED = ( - "Function-calling mode could not verify whether a tool-approval policy exists ({error}). " - "It fails closed: no tool ran. Fix policy storage, or use cuga_lite_execution_mode = \"codeact\"." -) - FC_BIND_FAILED = ( "Function-calling mode could not advertise any tool natively (bind_tools failed or is unsupported " "by this model, or there is no executable tool to bind), so no native tool call is possible. " 'It fails closed: no model call was made. Use cuga_lite_execution_mode = "codeact" for this model.' ) +FC_VARIABLES_NOTE = ( + "These are results of earlier tool calls in this conversation. Call a tool again only if you " + "need more than the preview shows." +) +_CODEACT_VARIABLES_NOTE = "You can use these variables directly by their names." + FC_STEP_LIMIT_CALL_REPLY = "Not executed: the step limit was reached before this call could run." FC_BUDGET_CALL_REPLY = ( "Not executed: the tool budget for this turn is spent. Answer from the data already retrieved." @@ -380,7 +382,7 @@ async def resolve_bind_tools( ) else: # Nothing executable was recorded: advertising the full catalogue - # would offer tools tool_exec cannot run. Bind nothing; the turn + # would offer tools the sandbox cannot run. Bind nothing; the turn # then fails closed in execute_call_model_fc. logger.error("[fc] no executable tools recorded by prepare; binding nothing") return None @@ -489,25 +491,6 @@ def _resolved_bind_mode(self, configurable: dict) -> str: return str(candidate).strip().lower() return _bind_tools_mode_from_settings() - async def _tool_approval_policies_exist(self, config: Any) -> bool: - """True when an enabled tool-approval policy is stored. - - Function-calling has no approval interrupt yet — CodeAct's runs on the - generated code string, after the seam — so the mode refuses to start - rather than run a guarded tool unprompted. Raises when the answer - cannot be established (policy system not initialised, storage backend - down): ``strict=True`` keeps the storage layer from turning a backend - failure into an empty list, and the caller fails closed on any error. - """ - from cuga.backend.cuga_graph.policy.configurable import PolicyConfigurable - from cuga.backend.cuga_graph.policy.models import PolicyType - - policy_system = PolicyConfigurable.from_config(config or {}) - policies = await policy_system.agent.storage.list_policies( - policy_type=PolicyType.TOOL_APPROVAL, enabled_only=True, limit=1, strict=True - ) - return bool(policies) - async def execute_call_model_fc( self, *, @@ -531,15 +514,17 @@ async def execute_call_model_fc( history normalised into a provider-valid transcript — hands them straight to ``bound.ainvoke`` so the shared dict serializer is never involved — then: - - ``tool_calls`` present -> ``Command(goto="tool_exec")``, assistant turn kept verbatim + - ``tool_calls`` present -> translated into a block (``fc_actions``) and routed + to the sandbox like CodeAct code, assistant turn kept verbatim; the sandbox + answers every id with a ``ToolMessage`` - otherwise -> final answer, ``END`` - an empty reply gets one corrective turn, like the CodeAct path - a fenced code block with no ``tool_calls`` is a mode violation: it is never executed; the model gets one corrective turn instead. - Refuses to start when an enabled tool-approval policy exists — or when - that cannot be verified: there is no approval interrupt on this path - yet, and silently running a guarded tool is worse than not running. + Tool approval is the CodeAct check run on the translated block: the real + tool names are in its text, so a matching policy interrupts before the + sandbox, and the approved block resumes into it. """ if self._execution_mode(configurable) != EXECUTION_MODE_FUNCTION_CALLING: return None @@ -551,22 +536,6 @@ async def execute_call_model_fc( cfg = configurable or {} history: list = list(modified_messages) - if settings.policy.enabled: - # Fail closed, stricter than CodeAct's own check: a refusal here is - # cheap and explicit, a guarded tool running unprompted is not. - try: - blocked_reason = ( - FC_TOOL_APPROVAL_UNSUPPORTED if await self._tool_approval_policies_exist(config) else None - ) - except Exception as exc: - logger.error("[fc] could not verify tool-approval policies ({}); refusing to start", exc) - blocked_reason = FC_TOOL_APPROVAL_UNVERIFIED.format(error=exc) - if blocked_reason: - logger.error("[fc] refusing to start: {}", blocked_reason) - return create_error_command( - self, history, AIMessage(content=blocked_reason), state.step_count - ) - if not budget_exhausted and bound is active_model: # bind_tools failed, is unsupported by the model, or had nothing to bind — # call_model fell back to the unbound model. That run cannot make a native @@ -578,12 +547,13 @@ async def execute_call_model_fc( msgs.extend(_few_shot_to_messages(self.get_few_shot_messages(state))) msgs.extend(_normalize_history_for_replay(history)) if variables_addendum: - # Outbound only, like call_model's CodeAct path (#600): never persisted. + # Outbound only, like call_model's CodeAct path (#600): never persisted. The + # variables are earlier tool results; the CodeAct closing line would invite + # the very code block this mode treats as a violation. + fc_addendum = variables_addendum.replace(_CODEACT_VARIABLES_NOTE, FC_VARIABLES_NOTE) for i in range(len(msgs) - 1, -1, -1): if isinstance(msgs[i], HumanMessage): - msgs[i] = msgs[i].model_copy( - update={"content": _message_text(msgs[i]) + variables_addendum} - ) + msgs[i] = msgs[i].model_copy(update={"content": _message_text(msgs[i]) + fc_addendum}) break if budget_exhausted: # Outbound only as well. @@ -649,14 +619,59 @@ async def execute_call_model_fc( base_meta = dict(self.build_metadata_update(state, playbook_fired=playbook_fired) or {}) if has_calls: - logger.info("[fc] {} native tool_call(s) -> tool_exec", len(tool_calls) + len(invalid_tool_calls)) + one_per_step = ( + resolve_step_discipline(cfg, self._runtime_model_name(cfg)) + == STEP_DISCIPLINE_ONE_TOOL_PER_STEP + ) + block, plan = plan_tool_calls( + tool_calls, invalid_tool_calls, self._tools_context, one_per_step=one_per_step + ) + if block is None: + # Every call was answered at planning time (no name, unknown tool, bad + # arguments, provider-rejected, deferred): nothing to execute. Charge the + # execute step the sandbox would have taken and hand the replies back. + replies = replies_without_execution(plan, "") + exec_step = new_step_count + 1 + limit_cmd = enforce_step_limit( + self, + state=state, + messages=final_messages + replies, + new_step_count=exec_step, + limit=max_steps, + ) + if limit_cmd is not None: + return limit_cmd + logger.info("[fc] {} native tool_call(s), none executable -> call_model", len(plan)) + return Command( + goto="call_model", + update={ + self.messages_key: final_messages + replies, + "script": None, + "step_count": exec_step, + self.metadata_key: base_meta, + }, + ) + if settings.policy.enabled: + # The approval check CodeAct runs on its block. The tool names are in + # the block text, so a policy on them interrupts here; the approved + # block resumes into the sandbox with the plan still in metadata. + approval_cmd = await ToolApprovalHandler.check_and_create_approval_interrupt( + self, state, block, content, config, ai_message=response + ) + if approval_cmd is not None: + approval_cmd.update[self.metadata_key] = { + **(approval_cmd.update.get(self.metadata_key) or {}), + FC_PENDING_KEY: plan, + } + return approval_cmd + logger.info("[fc] {} native tool_call(s) -> sandbox", len(plan)) return Command( - goto="tool_exec", + goto=self.execute_node_name, update={ self.messages_key: final_messages, - "script": None, + "script": block, "step_count": new_step_count, - self.metadata_key: base_meta, + self.metadata_key: {**base_meta, FC_PENDING_KEY: plan}, }, ) @@ -730,9 +745,6 @@ async def execute_call_model_fc( }, ) - def build_tool_exec_node(self): - return create_tool_exec_node(self) - def build_prepare_node(self, lc_bind_tools_meta: dict): return create_prepare_tools_and_apps_node(self, lc_bind_tools_meta) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index b2168db7a..823719e01 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -55,6 +55,7 @@ render_fc_prompt, resolve_cuga_lite_few_shots_enabled, ) +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import FC_PENDING_KEY from cuga.backend.cuga_graph.nodes.cuga_lite.helpers.app_auth import call_authenticate_apps from cuga.backend.cuga_graph.policy.enactment import PolicyEnactment from cuga.backend.skills import ( @@ -843,6 +844,11 @@ async def _wrapped(*args, **kwargs): # drop the counter here — prepare runs exactly once per turn on both paths. _turn_metadata = dict(state.cuga_lite_metadata or {}) _turn_metadata.pop("fc_mode_violations", None) + if not _turn_metadata.get("user_approved"): + # A plan left by a turn that never reached the sandbox (a denied approval, a + # crash). On an approval resume the subgraph is re-entered through prepare and + # the plan is what the approved block answers, so it stays. + _turn_metadata.pop(FC_PENDING_KEY, None) update_payload: dict[str, Any] = { "tools_prepared": True, "prepared_prompt": dynamic_prompt, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py index 87f124e45..ada340a46 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py @@ -18,6 +18,14 @@ ) from cuga.backend.cuga_graph.nodes.cuga_agent_core.policy.execution_policy import ExecutionRouter from cuga.backend.cuga_graph.nodes.cuga_agent_core.policy.tool_approval_handler import ToolApprovalHandler +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import ( + NOT_EXECUTED_REPLY, + fc_context_overlay, + fc_pending_plan, + metadata_without_plan, + replies_from_execution, + replies_without_execution, +) from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.response_utils import reflection_current_task from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import ( CodeExecutor, @@ -148,6 +156,15 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): # Add tools to context context = {**existing_vars, **adapter._tools_context} + # A function-calling turn left its call plan: the block is the translated + # native calls, and every exit below answers each id with a ToolMessage + # instead of an "Execution output:" message (see fc_actions). Empty for a + # CodeAct block, where nothing below changes. + fc_plan = fc_pending_plan(adapter, state) + if fc_plan: + context.update(fc_context_overlay(adapter._tools_context, fc_plan)) + fc_meta = {adapter.metadata_key: metadata_without_plan(adapter, state)} if fc_plan else {} + # Start tool call tracking (enabled via invoke parameter, or internally # whenever a weak-schema tool's output shape hasn't been observed yet). # "timings_only" (set when tracking is forced for the run receipt) records @@ -211,11 +228,14 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): if decision.gate == "revise": ToolCallTracker.stop_tracking() msg = verify_blocked_message(decision.alert) - new_message = HumanMessage(content=msg) + new_messages = ( + replies_without_execution(fc_plan, msg) if fc_plan else [HumanMessage(content=msg)] + ) updated_messages, error_message = core_append_with_step_limit( - adapter, state, [new_message], max_steps + adapter, state, new_messages, max_steps ) skip_updates = { + **fc_meta, "variables_storage": state.variables_storage, "variable_counter_state": state.variable_counter_state, "variable_creation_order": state.variable_creation_order, @@ -271,6 +291,11 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): thread_id=current_thread_id, apps_list=current_apps_list, plan=_exec_plan, + # A function-calling block calls in-process callables, so it runs on + # the local executor whatever the Python backend; a find_tools result + # must stay among its variables to be read back as the reply. + mode="local" if fc_plan else None, + keep_listing_vars=bool(fc_plan), ) finally: if cap_token is not None: @@ -290,6 +315,20 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): # Output is already formatted and trimmed by code_executor logger.debug(f"\n\n------\n\n📝 Execution output:\n\n{output}\n\n------\n\n") + fc_replies: list = [] + if fc_plan: + fc_replies, new_vars, dropped = replies_from_execution( + fc_plan, + new_vars, + output, + output_limit=settings.advanced_features.execution_output_max_length, + timeout=settings.advanced_features.sandbox_execution_timeout, + ) + # Error markers: the reply carries them, the variables do not. + for name in dropped: + if name in state.variables_manager.get_variable_names(): + state.variables_manager.remove_variable(name) + # Update variables using CugaLiteState's variables_manager # This automatically updates state.variables_storage for name, value in new_vars.items(): @@ -300,7 +339,7 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): ) reflection_output = "" - if reflection_enabled: + if reflection_enabled and not fc_plan: try: active_model = configurable.get("llm") or _llm_manager.get_model( settings.agent.planner.model @@ -362,16 +401,36 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): f"{execution_message_content}\n\n---\n\nSummary:\n{reflection_output}" ) - adapter._tracker.collect_step( - step=Step( - name="User_return", - data=execution_message_content, + if fc_plan: + adapter._tracker.collect_step( + step=Step( + name="Tool_results", + data=json.dumps( + [ + { + "tool_call_id": m.tool_call_id, + "name": m.name, + "status": m.status, + "content": m.content, + } + for m in fc_replies + ], + ensure_ascii=False, + default=str, + ), + ) ) - ) - - new_message = HumanMessage(content=execution_message_content) + new_messages = fc_replies + else: + adapter._tracker.collect_step( + step=Step( + name="User_return", + data=execution_message_content, + ) + ) + new_messages = [HumanMessage(content=execution_message_content)] updated_messages, error_message = core_append_with_step_limit( - adapter, state, [new_message], max_steps + adapter, state, new_messages, max_steps ) # Collect tool calls from this execution @@ -388,6 +447,7 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): error_message, state.step_count, additional_updates={ + **fc_meta, "variables_storage": state.variables_storage, "variable_counter_state": state.variable_counter_state, "variable_creation_order": state.variable_creation_order, @@ -404,6 +464,7 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): todo_state_update = extract_task_todos_from_new_vars(new_vars) base_update = { + **fc_meta, "chat_messages": updated_messages, "variables_storage": state.variables_storage, "variable_counter_state": state.variable_counter_state, @@ -428,9 +489,13 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): error_msg = f"Error during execution: {str(e)}" logger.error(error_msg) - new_message = HumanMessage(content=error_msg) + new_messages = ( + replies_without_execution(fc_plan, NOT_EXECUTED_REPLY.format(reason=e)) + if fc_plan + else [HumanMessage(content=error_msg)] + ) updated_messages, limit_error_message = core_append_with_step_limit( - adapter, state, [new_message], max_steps + adapter, state, new_messages, max_steps ) if limit_error_message: @@ -439,10 +504,11 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): updated_messages, limit_error_message, state.step_count, - additional_updates=_budget_updates(), + additional_updates={**fc_meta, **_budget_updates()}, ) return { + **fc_meta, "chat_messages": updated_messages, "error": error_msg, "final_answer": error_msg, diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py deleted file mode 100644 index 43c2cecd3..000000000 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/tool_exec_node.py +++ /dev/null @@ -1,379 +0,0 @@ -"""``tool_exec`` node — executes native ``tool_calls`` in function-calling mode. - -The counterpart of ``sandbox_node`` for ``cuga_lite_execution_mode = -"function_calling"``: where the sandbox runs a generated Python block, this -node runs the ``tool_calls`` the model attached to its last ``AIMessage`` and -replies to each with a ``ToolMessage`` carrying the call's exact -``tool_call_id``, so the conversation stays replayable to the provider. - -Guard envelope --------------- -The budgets and the tracker are *not* on the callables in ``_tools_context`` — -they are applied by ``eval_with_tools_async`` and the sandbox node. A node that -merely called the callable would bypass the per-run / per-thread / per-block -tool-call caps and the run receipt, exactly the bug #560 shipped with on the -supervisor graph. So this node mirrors the sandbox node's envelope in the same -order: ``start_tracking`` -> ``seed_call_budget`` -> ``seed_block_budget`` -> -each call under ``counted_tool_call`` and a timeout -> ``stop_tracking`` -> -``_budget_updates()`` on every exit. - -Every outcome is a ``ToolMessage``, never a raised exception: unknown tool, -malformed arguments, a tool exception, a timeout, or a spent budget. One -failing call does not abort its siblings, and every id the model issued — -including provider-side ``invalid_tool_calls`` — receives a reply. -""" - -from __future__ import annotations - -import asyncio -import inspect -import json -from typing import Any, Callable, Dict, List, Optional - -from langchain_core.messages import AIMessage, ToolMessage -from langchain_core.runnables import RunnableConfig -from langgraph.types import Command -from loguru import logger - -from cuga.backend.activity_tracker.tracker import Step -from cuga.backend.cuga_graph.nodes.cuga_agent_core.execution.code_extraction import make_tool_awaitable -from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.graph_nodes import ( - append_chat_messages_with_step_limit as core_append_with_step_limit, - create_error_command as core_create_error_command, -) -from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.sandbox_node import _budget_updates -from cuga.backend.cuga_graph.nodes.cuga_lite.model_runtime_profile import ( - STEP_DISCIPLINE_ONE_TOOL_PER_STEP, - resolve_step_discipline, - resolved_runtime_model_name, -) -from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( - ToolCallBudgetExceeded, - ToolCallTracker, - counted_tool_call, -) -from cuga.config import settings - -DEFERRED_CALL_MESSAGE = ( - "Deferred: step discipline is on, so only the first tool call of a turn is attempted. " - "Read that result, then re-issue this call in your next turn if you still need it." -) - -TRUNCATION_MARKER = "\n... [result truncated: {limit} characters of output remain for this batch]" - - -def _tool_call_parts(call: Any) -> tuple[str, str, Any]: - """``(id, name, args)`` from a LangChain tool_call dict or a raw OpenAI-shaped one.""" - if not isinstance(call, dict): - return "", "", None - fn = call.get("function") if isinstance(call.get("function"), dict) else {} - call_id = str(call.get("id") or "") - name = str(call.get("name") or fn.get("name") or "") - args = call.get("args") if "args" in call else fn.get("arguments") - return call_id, name, args - - -def _coerce_args(args: Any) -> tuple[Optional[Dict[str, Any]], Optional[str]]: - """Return ``(kwargs, error)``. Strings are JSON-decoded; anything else is rejected.""" - if args is None: - return {}, None - if isinstance(args, dict): - return args, None - if isinstance(args, str): - text = args.strip() - if not text: - return {}, None - try: - parsed = json.loads(text) - except json.JSONDecodeError as exc: - return ( - None, - f"Could not parse tool arguments as JSON ({exc}). Re-issue the call with valid JSON arguments.", - ) - if isinstance(parsed, dict): - return parsed, None - return None, f"Tool arguments must be a JSON object, got {type(parsed).__name__}." - return None, f"Tool arguments must be an object, got {type(args).__name__}." - - -def _stringify(result: Any) -> str: - if isinstance(result, str): - return result - try: - return json.dumps(result, ensure_ascii=False, default=str) - except Exception: - return str(result) - - -class _BatchOutputBudget: - """One turn's batch of results shares ``execution_output_max_length``, like one block. - - Per-result truncation alone would let N calls return N x the limit; this keeps - the whole batch under the ceiling the sandbox applies to a block's output. - """ - - def __init__(self, limit: Any): - self.limit = int(limit or 0) - self.remaining = self.limit - - def take(self, text: str) -> str: - if not self.limit: - return text - if len(text) > self.remaining: - cap = max(self.remaining, 0) - text = text[:cap] + TRUNCATION_MARKER.format(limit=cap) - self.remaining = max(self.remaining - len(text), 0) - return text - - -def _error_message(text: str, *, call_id: str, name: str) -> ToolMessage: - return ToolMessage(content=text, tool_call_id=call_id, name=name or "unknown", status="error") - - -def create_tool_exec_node(adapter: Any) -> Callable: - async def tool_exec(state: Any, config: Optional[RunnableConfig] = None): - configurable = config.get("configurable", {}) if config else {} - from cuga.backend.cuga_graph.utils.langfuse_tracing import sync_langfuse_callbacks_from_config - - sync_langfuse_callbacks_from_config(config) - track_tool_calls = configurable.get("track_tool_calls", False) - max_steps = configurable.get("cuga_lite_max_steps") if "cuga_lite_max_steps" in configurable else None - model_name = resolved_runtime_model_name( - configurable_llm=configurable.get("llm"), graph_default_model=getattr(adapter, "_model", None) - ) - one_per_step = resolve_step_discipline(configurable, model_name) == STEP_DISCIPLINE_ONE_TOOL_PER_STEP - # One batch is the function-calling analogue of one code block, so the - # outer bound per call is the block wall clock. Registry tools keep their - # own inner ``tool_call_timeout``. - timeout = getattr(settings.advanced_features, "sandbox_execution_timeout", 30) or None - output_budget = _BatchOutputBudget( - getattr(settings.advanced_features, "execution_output_max_length", 0) - ) - - messages = adapter.get_messages(state) - last = messages[-1] if messages else None - calls: List[Any] = ( - list(getattr(last, "tool_calls", None) or []) if isinstance(last, AIMessage) else [] - ) - invalid: List[Any] = ( - list(getattr(last, "invalid_tool_calls", None) or []) if isinstance(last, AIMessage) else [] - ) - - # ── Guard envelope (mirrors sandbox_node, same order) ────────────── - ToolCallTracker.start_tracking( - enabled=bool(track_tool_calls), - timings_only=track_tool_calls == "timings_only", - ) - ToolCallTracker.seed_call_budget( - getattr(state, "tool_calls_used_run", 0), - getattr(state, "tool_calls_used_thread", 0), - ) - # One turn's batch of calls is the function-calling analogue of one code block. - ToolCallTracker.seed_block_budget() - - results: List[ToolMessage] = [] - executed = 0 - execution_tool_calls: List[Dict[str, Any]] = [] - try: - try: - executed = await _run_batch( - adapter, - calls, - invalid, - results, - one_per_step=one_per_step, - timeout=timeout, - output_budget=output_budget, - ) - finally: - execution_tool_calls = ToolCallTracker.stop_tracking() - - if not calls and not invalid: - # Nothing executable: advance the step and hand back to the model so a - # mis-route can never become an infinite no-op loop. - logger.warning("tool_exec entered with no tool_calls on the last message") - return Command( - goto="call_model", - update={"step_count": state.step_count + 1, "script": None, **_budget_updates()}, - ) - - try: - adapter._tracker.collect_step( - step=Step( - name="Tool_results", - data=json.dumps( - [ - { - "tool_call_id": m.tool_call_id, - "name": m.name, - "status": m.status, - "content": m.content, - } - for m in results - ], - ensure_ascii=False, - default=str, - ), - ) - ) - except Exception as exc: - logger.debug(f"tool_exec tracker error: {exc}") - - updated_messages, error_message = core_append_with_step_limit(adapter, state, results, max_steps) - accumulated_tool_calls = (state.tool_calls or []) + ( - execution_tool_calls if track_tool_calls else [] - ) - - if error_message: - return core_create_error_command( - adapter, - updated_messages, - error_message, - state.step_count, - additional_updates={"tool_calls": accumulated_tool_calls, **_budget_updates()}, - ) - - logger.info( - "[fc] tool_exec: {} call(s) -> {} ToolMessage(s), {} executed", - len(calls), - len(results), - executed, - ) - return Command( - goto="call_model", - update={ - adapter.messages_key: updated_messages, - "script": None, - "step_count": state.step_count + 1, - "tool_calls": accumulated_tool_calls, - **_budget_updates(), - }, - ) - except Exception as exc: - # Mirrors the sandbox's outer guard: a failure outside the per-call - # handling must still report the budget the batch spent, or the - # checkpoint keeps the pre-batch counts (see _budget_updates). - logger.error(f"tool_exec failed: {exc}") - # Tools may already have run: keep their replies, and answer every id - # the batch never reached, so the persisted turn has no dangling call. - answered = {m.tool_call_id for m in results} - replies = list(results) - for index, call in enumerate(list(calls) + list(invalid)): - call_id, name, _ = _tool_call_parts(call) - call_id = call_id or f"call_{index}" - if call_id not in answered: - answered.add(call_id) - replies.append( - _error_message( - f"Not executed: tool execution failed before this call ran ({exc}).", - call_id=call_id, - name=name, - ) - ) - return core_create_error_command( - adapter, - adapter.get_messages(state) + replies, - AIMessage(content=f"Error during tool execution: {exc}"), - state.step_count, - additional_updates={ - "tool_calls": list(state.tool_calls or []) - + (execution_tool_calls if track_tool_calls else []), - **_budget_updates(), - }, - ) - - return tool_exec - - -async def _run_batch( - adapter: Any, - calls: List[Any], - invalid: List[Any], - results: List[ToolMessage], - *, - one_per_step: bool, - timeout: Optional[float], - output_budget: _BatchOutputBudget, -) -> int: - """Run one turn's calls in order, appending one ``ToolMessage`` per id to ``results``. - - Returns how many calls actually executed. Every outcome is a message, never - an exception, so one failing call cannot take its siblings down. - """ - executed = 0 - budget_stop: Optional[str] = None - for index, call in enumerate(calls): - call_id, name, raw_args = _tool_call_parts(call) - call_id = call_id or f"call_{index}" - if not name: - results.append(_error_message("Tool call carried no tool name.", call_id=call_id, name="")) - continue - if budget_stop: - results.append(_error_message(budget_stop, call_id=call_id, name=name)) - continue - if one_per_step and index > 0: - # By position, not by success: an error on the first call is the result - # the model must read before it decides the next one. - results.append(_error_message(DEFERRED_CALL_MESSAGE, call_id=call_id, name=name)) - continue - - fn = adapter._tools_context.get(name) - if fn is None or not callable(fn): - known = ", ".join(sorted(k for k in adapter._tools_context if not k.startswith("_"))) - results.append( - _error_message( - f"Unknown tool '{name}'. Choose one of the provided tools: {known or '(none)'}.", - call_id=call_id, - name=name, - ) - ) - continue - - kwargs, arg_error = _coerce_args(raw_args) - if arg_error: - results.append(_error_message(arg_error, call_id=call_id, name=name)) - continue - - awaitable = fn if inspect.iscoroutinefunction(fn) else make_tool_awaitable(fn) - counted = counted_tool_call(awaitable) - try: - result = await asyncio.wait_for(counted(**kwargs), timeout=timeout) - executed += 1 - results.append( - ToolMessage(content=output_budget.take(_stringify(result)), tool_call_id=call_id, name=name) - ) - except asyncio.TimeoutError: - results.append( - _error_message( - f"Tool '{name}' timed out after {timeout}s. Try a narrower call or a different tool.", - call_id=call_id, - name=name, - ) - ) - except ToolCallBudgetExceeded as exc: - # Every remaining id still gets a reply — without invoking anything. - if exc.scope != "block": - budget_stop = str(exc) - results.append(_error_message(str(exc), call_id=call_id, name=name)) - except TypeError as exc: - results.append( - _error_message( - f"Tool '{name}' rejected these arguments: {exc}. Check the parameter names and types.", - call_id=call_id, - name=name, - ) - ) - except Exception as exc: # one failing call must not abort its siblings - results.append(_error_message(f"Tool '{name}' failed: {exc}", call_id=call_id, name=name)) - - for index, bad in enumerate(invalid): - call_id, name, _ = _tool_call_parts(bad) - reason = (bad.get("error") if isinstance(bad, dict) else None) or "malformed tool call" - results.append( - _error_message( - f"The provider could not parse this tool call ({reason}). Re-issue it with valid arguments.", - call_id=call_id or f"invalid_{index}", - name=name, - ) - ) - return executed diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py index b3adcbe56..b855d6362 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/cuga_lite_graph.py @@ -229,10 +229,6 @@ def create_cuga_lite_graph( prepare_node = adapter.build_prepare_node(lc_bind_tools_meta) sandbox_node = adapter.build_sandbox_node(thread_id, apps_list) call_model_node = _create_shared_call_model_node(adapter, model, settings) - # Native function-calling node. Always wired so the mode can be switched per - # invoke without recompiling; call_model only routes to it in that mode, so - # it is dormant on every CodeAct run. - tool_exec_node = adapter.build_tool_exec_node() return build_agent_graph( adapter=adapter, @@ -240,5 +236,4 @@ def create_cuga_lite_graph( prepare_node=prepare_node, call_model_node=call_model_node, execute_node=sandbox_node, - tool_exec_node=tool_exec_node, ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/code_executor.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/code_executor.py index 30947b78f..d4e395a25 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/code_executor.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/executors/code_executor.py @@ -118,6 +118,7 @@ async def eval_with_tools_async( mode: Optional[Literal['local', 'e2b', 'opensandbox']] = None, plan: Optional[ExecutionPlan] = None, variable_manager: Optional[Any] = None, + keep_listing_vars: bool = False, ) -> tuple[str, dict[str, Any]]: """Execute code with async tools available in the local namespace. @@ -133,6 +134,9 @@ async def eval_with_tools_async( Python execution path unless ``mode`` is given explicitly. variable_manager: Variable manager to record new variables into. Defaults to ``state.variables_manager`` (preserves prior behavior). + keep_listing_vars: Keep find_tools listing markdown among the block's + variables (function-calling reads the reply from there). Default + strips it, as before. Returns: Tuple of (execution result, new variables dictionary) @@ -177,6 +181,7 @@ async def eval_with_tools_async( plan=plan, variable_manager=variable_manager, skills_on=skills_on, + keep_listing_vars=keep_listing_vars, ) finally: reset_skills_relaxed_execution(skills_token) @@ -195,6 +200,7 @@ async def _eval_with_tools_async_impl( plan: Optional[ExecutionPlan] = None, variable_manager: Optional[Any] = None, skills_on: bool, + keep_listing_vars: bool = False, ) -> tuple[str, dict[str, Any]]: result = "" @@ -304,7 +310,8 @@ async def _eval_with_tools_async_impl( {k: v for k, v in new_vars.items() if k not in changed_keys}, keep_last_n ) new_vars.update(changed_vars) - new_vars = _omit_find_tools_listing_vars(new_vars) + if not keep_listing_vars: + new_vars = _omit_find_tools_listing_vars(new_vars) # Format/trim the output before adding variables result = format_execution_output(result) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py index f74bd7f04..b63fc1fe2 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_execution_mode_config.py @@ -3,7 +3,7 @@ One assertion per control plane — settings.toml, validators, ``configurable``, the per-model runtime profile, the SDK constructor / ``invoke`` / ``stream`` — so a knob wired into only some surfaces cannot ship (the usual "my config does -nothing" bug). Plus the graph wiring: the ``tool_exec`` node is present on the +nothing" bug). Plus the graph wiring: no mode-specific node exists on the CugaLite graph regardless of mode, and absent when a graph does not opt in. """ @@ -11,7 +11,7 @@ import inspect from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -242,48 +242,30 @@ async def call_model(state, config=None): async def execute(state, config=None): return {} - async def tool_exec(state, config=None): - return {} - - return prepare, call_model, execute, tool_exec + return prepare, call_model, execute -def test_build_agent_graph_adds_tool_exec_only_when_given(): +def test_build_agent_graph_has_exactly_the_three_shared_nodes(): + """Function-calling adds no node: a native turn is a block the execute node runs.""" from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_graph import build_agent_graph from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState - prepare, call_model, execute, tool_exec = _dummy_nodes() + prepare, call_model, execute = _dummy_nodes() adapter = SimpleNamespace(execute_node_name="sandbox") - - without = build_agent_graph( + graph = build_agent_graph( adapter=adapter, state_class=CugaLiteState, prepare_node=prepare, call_model_node=call_model, execute_node=execute, ) - assert "tool_exec" not in without.nodes, "a graph that does not opt in must be unchanged" + assert set(graph.nodes) - {"__start__"} == {"prepare", "call_model", "sandbox"} + assert ("sandbox", "call_model") in graph.edges - with_node = build_agent_graph( - adapter=adapter, - state_class=CugaLiteState, - prepare_node=prepare, - call_model_node=call_model, - execute_node=execute, - tool_exec_node=tool_exec, - ) - assert "tool_exec" in with_node.nodes - assert ("tool_exec", "call_model") not in with_node.edges, "it routes with a Command, no static edge" +def test_cuga_lite_graph_has_no_mode_specific_node(): -def test_cuga_lite_graph_always_wires_tool_exec_so_mode_can_switch_per_invoke(): from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import create_cuga_lite_graph - provider = MagicMock() - provider.get_all_tools = AsyncMock(return_value=[]) - provider.get_apps = AsyncMock(return_value=[]) - provider.get_tools = AsyncMock(return_value=[]) - graph = create_cuga_lite_graph(model=MagicMock(), tool_provider=provider, apps_list=[]) - - assert {"prepare", "call_model", "sandbox", "tool_exec"} <= set(graph.nodes) - assert ("tool_exec", "call_model") not in graph.edges + graph = create_cuga_lite_graph(model=MagicMock(), tool_provider=None, apps_list=[], thread_id="t") + assert set(graph.nodes) - {"__start__"} == {"prepare", "call_model", "sandbox"} diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index d7d4c8f69..1adcd83f3 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -2,11 +2,12 @@ Codeact: ``None``, nothing invoked. Function-calling: the bound model gets real message objects (system prompt, few-shot demos, history normalised into a -provider-valid transcript), and the response routes on ``tool_calls`` — to -``tool_exec`` with the assistant turn kept verbatim, or to END as the final -answer. Every id the model issued is answered before the run can end, an empty -reply gets one retry, a fenced code block is a mode violation (never executed), -and the mode refuses to start while a tool-approval policy is configured. +provider-valid transcript), and the response routes on ``tool_calls`` — to the +sandbox as a translated block with the assistant turn kept verbatim, or to END +as the final answer. Every id the model issued is answered before the run can +end, an empty reply gets one retry, a fenced code block is a mode violation +(never executed), and a matching tool-approval policy interrupts on the block +exactly as it does for CodeAct code. """ from __future__ import annotations @@ -25,13 +26,12 @@ ) from cuga.backend.cuga_graph.nodes.cuga_agent_core.graph.shared_nodes import TOOL_BUDGET_EXHAUSTED_INSTRUCTION from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import graph_adapter as ga +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import FC_PENDING_KEY from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( FC_BIND_FAILED, FC_BUDGET_CALL_REPLY, FC_MODE_VIOLATION_CORRECTION, FC_STEP_LIMIT_CALL_REPLY, - FC_TOOL_APPROVAL_UNSUPPORTED, - FC_TOOL_APPROVAL_UNVERIFIED, FC_UNANSWERED_CALL_REPLY, AgentGraphAdapter, _looks_like_python_block, @@ -57,14 +57,20 @@ async def ainvoke(self, messages, config=None, **kwargs): return self.response -def _adapter(tools_context_ref=None) -> AgentGraphAdapter: - return AgentGraphAdapter( +async def _add(a: int, b: int) -> int: + return a + b + + +def _adapter(tools_context_ref=None, tools=None) -> AgentGraphAdapter: + adapter = AgentGraphAdapter( tracker=MagicMock(), base_callbacks=[], task_todos_ref=[], tools_context_ref=tools_context_ref if tools_context_ref is not None else {}, base_tool_provider=None, ) + adapter._tools_context.update(tools if tools is not None else {"add": _add}) + return adapter def _state(messages=None, few_shots=None, step_count=0, max_steps=None, metadata=None): @@ -114,24 +120,40 @@ async def test_codeact_returns_none_and_invokes_nothing(): @pytest.mark.asyncio -async def test_tool_calls_route_to_tool_exec_with_the_assistant_turn_verbatim(): +async def test_tool_calls_route_to_the_sandbox_as_a_block_with_the_assistant_turn_verbatim(): response = AIMessage(content="", tool_calls=[_CALL]) model = _Model(response) state = _state() cmd = await _turn(_adapter(), model, state, FC) - assert cmd.goto == "tool_exec" + assert cmd.goto == "sandbox" persisted = cmd.update["chat_messages"] assert persisted[:-1] == state.chat_messages - assert persisted[-1] is response, "tool_calls must reach tool_exec untouched, ids included" - assert cmd.update["script"] is None and cmd.update["step_count"] == 1 + assert persisted[-1] is response, "the assistant turn is persisted untouched, ids included" + assert 'await add(**{"a": 1, "b": 2})' in cmd.update["script"], cmd.update["script"] + assert cmd.update["step_count"] == 1 + (entry,) = cmd.update["cuga_lite_metadata"][FC_PENDING_KEY] + assert entry["id"] == "c1" and entry["var"] == "tool_result_c1" and entry["reply"] is None (outbound,) = model.seen assert isinstance(outbound[0], SystemMessage) and outbound[0].content == "SYS" assert outbound[1:] == state.chat_messages, "history goes out as message objects, not the CodeAct dicts" +@pytest.mark.asyncio +async def test_calls_nothing_can_execute_are_answered_without_the_sandbox(): + """Unknown tool, bad arguments, provider-rejected: every id is answered here, and the + execute step the sandbox would have taken is still charged.""" + unknown = {"name": "nope", "args": {}, "id": "c9", "type": "tool_call"} + cmd = await _turn(_adapter(), _Model(AIMessage(content="", tool_calls=[unknown])), _state(), FC) + + assert cmd.goto == "call_model" and cmd.update["script"] is None and cmd.update["step_count"] == 2 + reply = cmd.update["chat_messages"][-1] + assert isinstance(reply, ToolMessage) and reply.tool_call_id == "c9" and reply.status == "error" + assert reply.content == "Unknown tool 'nope'. Choose one of the provided tools: add." + + @pytest.mark.asyncio async def test_plain_text_is_the_final_answer(): cmd = await _turn(_adapter(), _Model(AIMessage(content="It is 3.")), _state(), FC) @@ -302,22 +324,76 @@ def test_replay_rebuilds_bare_shells_and_closes_dangling_calls(): # ── guards and binding ─────────────────────────────────────────────────────── +def _policy_system(match): + return SimpleNamespace(agent=SimpleNamespace(check_tool_approval_for_code=AsyncMock(return_value=match))) + + +def _approval_match(): + policy = SimpleNamespace( + id="p1", + name="Approve add", + required_tools=["add"], + required_apps=[], + approval_message="Adding needs approval.", + show_code_preview=True, + ) + return SimpleNamespace( + matched=True, policy=policy, confidence=1.0, reasoning="r", trigger_details={"matched_tools": ["add"]} + ) + + @pytest.mark.asyncio -async def test_refuses_to_start_when_a_tool_approval_policy_exists(monkeypatch): - """No approval interrupt exists on this path yet: fail closed, before any model call.""" +async def test_a_matching_tool_approval_policy_interrupts_before_the_sandbox(monkeypatch): + """The CodeAct approval check runs on the translated block, so a policy on the tool + name interrupts; the assistant turn is persisted with its tool_calls and the plan + rides the metadata, so the approved block resumes into the sandbox.""" from cuga.config import settings monkeypatch.setattr(settings.policy, "enabled", True, raising=False) - model = _Model(AIMessage(content="", tool_calls=[_CALL])) - - with patch.object(AgentGraphAdapter, "_tool_approval_policies_exist", new=AsyncMock(return_value=True)): + response = AIMessage(content="", tool_calls=[_CALL]) + model = _Model(response) + system = _policy_system(_approval_match()) + + with ( + patch( + "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.from_config", return_value=system + ), + patch( + "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.create_context_from_state", + return_value=SimpleNamespace(user_input="q"), + ), + ): cmd = await _turn(_adapter(), model, _state(), FC) - assert cmd.goto == END and cmd.update["error"] == FC_TOOL_APPROVAL_UNSUPPORTED - assert model.seen == [], "refused before the model was invoked" - with patch.object(AgentGraphAdapter, "_tool_approval_policies_exist", new=AsyncMock(return_value=False)): - cmd = await _turn(_adapter(), model, _state(), FC) - assert cmd.goto == "tool_exec" and len(model.seen) == 1, "policies without approval rules do not block" + assert cmd.goto == END and cmd.update["hitl_action"] is not None + assert "Adding needs approval." in cmd.update["final_answer"] + assert 'await add(**{"a": 1, "b": 2})' in cmd.update["script"], "the approved block is what resumes" + assert cmd.update["chat_messages"][-1] is response, "tool_calls survive the pause" + meta = cmd.update["cuga_lite_metadata"] + assert meta["approval_required"] is True and meta[FC_PENDING_KEY][0]["id"] == "c1" + (code, _context), _ = system.agent.check_tool_approval_for_code.call_args + assert "add(" in code + + +@pytest.mark.asyncio +async def test_no_matching_policy_routes_to_the_sandbox(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", True, raising=False) + system = _policy_system(SimpleNamespace(matched=False)) + + with ( + patch( + "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.from_config", return_value=system + ), + patch( + "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.create_context_from_state", + return_value=SimpleNamespace(user_input="q"), + ), + ): + cmd = await _turn(_adapter(), _Model(AIMessage(content="", tool_calls=[_CALL])), _state(), FC) + + assert cmd.goto == "sandbox" and FC_PENDING_KEY in cmd.update["cuga_lite_metadata"] @pytest.mark.asyncio @@ -331,7 +407,7 @@ async def fake_resolve(active_model, **kwargs): with patch.object(ga, "resolve_model_with_bind_tools", new=AsyncMock(side_effect=fake_resolve)): adapter = _adapter() # No executable set recorded: bind nothing (the turn then fails closed) rather - # than open the registry-wide catalogue to tools tool_exec cannot run. + # than open the registry-wide catalogue to tools the sandbox cannot run. assert await adapter.resolve_bind_tools(_state(), object(), dict(FC), None) is None await adapter.resolve_bind_tools( _state(), object(), {**FC, "cuga_lite_bind_tools_mode": "find_tools"}, None @@ -362,44 +438,6 @@ async def fake_resolve(active_model, **kwargs): ) -@pytest.mark.asyncio -async def test_refuses_when_the_policy_query_fails(monkeypatch): - """Fail closed on infrastructure failure too: 'could not verify' is not 'no policies'.""" - from cuga.config import settings - - monkeypatch.setattr(settings.policy, "enabled", True, raising=False) - model = _Model(AIMessage(content="", tool_calls=[_CALL])) - - with patch.object( - AgentGraphAdapter, - "_tool_approval_policies_exist", - new=AsyncMock(side_effect=RuntimeError("backend down")), - ): - cmd = await _turn(_adapter(), model, _state(), FC) - - assert cmd.goto == END and model.seen == [] - assert cmd.update["error"] == FC_TOOL_APPROVAL_UNVERIFIED.format(error="backend down") - - -@pytest.mark.asyncio -async def test_guard_asks_storage_strictly_so_a_backend_failure_is_not_an_empty_list(): - """PolicyStorage.list_policies swallows backend errors by default; the guard must not.""" - from types import SimpleNamespace - - calls = [] - - async def list_policies(**kwargs): - calls.append(kwargs) - return [] - - fake_system = SimpleNamespace(agent=SimpleNamespace(storage=SimpleNamespace(list_policies=list_policies))) - with patch( - "cuga.backend.cuga_graph.policy.configurable.PolicyConfigurable.from_config", return_value=fake_system - ): - assert await _adapter()._tool_approval_policies_exist({}) is False - assert calls and calls[0]["strict"] is True - - @pytest.mark.asyncio async def test_refuses_when_no_tools_are_bound(): """bind_tools failed / unsupported / nothing to bind: call_model hands the unbound model diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index a86a6d8a8..68a4b9fb5 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -2,13 +2,12 @@ Node tests pin what each node returns; only a real run proves the loop — prepare selects the FC prompt, call_model routes native ``tool_calls`` to -``tool_exec``, the ``ToolMessage`` replies are replayed to the model with their +the sandbox as a translated block, the ``ToolMessage`` replies are replayed to the model with their ids, and a text reply ends the run. The scripted model asserts on what it is sent, so a broken replay fails loudly instead of the test passing by accident. The last test is the feature-off golden: with no mode configured the graph -behaves exactly as before — dict-serialised CodeAct turns, the sandbox, and a -``tool_exec`` node that is present but never entered. +behaves exactly as before — dict-serialised CodeAct turns and the sandbox. """ from __future__ import annotations @@ -22,9 +21,8 @@ from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import ( FC_MODE_VIOLATION_CORRECTION, - AgentGraphAdapter, ) -from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import DEFERRED_CALL_MESSAGE +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import DEFERRED_CALL_MESSAGE from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState, create_cuga_lite_graph from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module @@ -244,11 +242,7 @@ async def test_code_block_in_fc_mode_is_corrected_not_executed(): @pytest.mark.asyncio -async def test_feature_off_is_the_old_codeact_run_and_tool_exec_is_never_entered(monkeypatch): - async def forbidden(state, config=None): - raise AssertionError("tool_exec entered on a CodeAct run") - - monkeypatch.setattr(AgentGraphAdapter, "build_tool_exec_node", lambda self: forbidden) +async def test_feature_off_is_the_old_codeact_run(monkeypatch): model = _ScriptedModel( [ AIMessage(content="```python\nr = await echo(value=3)\nprint(r)\n```"), @@ -259,7 +253,9 @@ async def forbidden(state, config=None): graph = create_cuga_lite_graph( model=model, tool_provider=_provider(_echo_tool()), apps_list=[], thread_id="t" ) - assert "tool_exec" in graph.nodes, "the node is wired (dormant) so the mode can switch per invoke" + assert set(graph.nodes) - {"__start__"} == {"prepare", "call_model", "sandbox"}, ( + "no extra node for the mode" + ) result = await graph.compile(checkpointer=MemorySaver()).ainvoke( CugaLiteState(chat_messages=[HumanMessage(content="use the tools")]), config=_config("codeact-golden") ) @@ -320,3 +316,110 @@ async def test_static_prompt_is_ignored_in_fc_mode(): ) assert "native function-calling" in result["prepared_prompt"] assert "Always answer with a ```python block" not in result["prepared_prompt"] + + +# ── what routing through the sandbox gives a native call ───────────────────── + + +@pytest.mark.asyncio +async def test_verify_revise_answers_every_call_without_running_it(monkeypatch): + """The pre-execute VERIFY gate decides on the translated block like on CodeAct code; + a revise verdict reaches the model as an error reply per call id, nothing runs.""" + from unittest.mock import AsyncMock + + from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import sandbox_node + from cuga.backend.cuga_graph.nodes.cuga_lite.reflection.pre_execute import verify_blocked_message + from cuga.backend.cuga_graph.nodes.cuga_lite.reflection.verify_result import VerifyDecision + + monkeypatch.setattr( + sandbox_node, + "decide_pre_execute_verify", + AsyncMock(return_value=VerifyDecision(gate="revise", alert="echoing 9 looks wrong")), + ) + + def after_revise(messages): + reply = _last_tool_message(messages) + assert reply.tool_call_id == "c1" and reply.status == "error" + assert reply.content == verify_blocked_message("echoing 9 looks wrong") + return AIMessage(content="Understood, not echoing.") + + model = _ScriptedModel( + [AIMessage(content="", tool_calls=[_tc("echo", {"value": 9}, "c1")]), after_revise] + ) + config = _config( + "fc-verify", cuga_lite_execution_mode="function_calling", pre_execute_verify_enabled=True + ) + + result = await _run(model, config, _echo_tool()) + + assert CALLS == [] and result["final_answer"] == "Understood, not echoing." + assert result["verify_revise_streak"] == 1 + + +@pytest.mark.asyncio +async def test_e2b_python_backend_still_runs_a_native_call_locally(monkeypatch): + """The translated block calls in-process callables, so it never goes to a remote sandbox.""" + from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import CodeExecutor + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "e2b_sandbox", True, raising=False) + + def forbidden(): + raise AssertionError("a function-calling block reached the E2B executor") + + monkeypatch.setattr(CodeExecutor, "_get_e2b_executor", classmethod(lambda cls: forbidden())) + model = _ScriptedModel( + [AIMessage(content="", tool_calls=[_tc("echo", {"value": 7}, "c1")]), AIMessage(content="7.")] + ) + + result = await _run(model, _config("fc-e2b", cuga_lite_execution_mode="function_calling"), _echo_tool()) + + assert CALLS == [("echo", 7)] and result["final_answer"] == "7." + + +@pytest.mark.asyncio +async def test_a_large_result_is_truncated_in_the_reply_but_kept_whole_as_a_variable(): + from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import TRUNCATION_MARKER + from cuga.config import settings + + limit = int(settings.advanced_features.execution_output_max_length) + big = "x" * (limit + 100) + + async def dump() -> str: + """A large payload.""" + return big + + tool = StructuredTool.from_function(coroutine=dump, name="dump", description="A large payload.") + + def after(messages): + reply = _last_tool_message(messages) + assert reply.content == "x" * limit + TRUNCATION_MARKER.format(limit=limit) + return AIMessage(content="Got it.") + + model = _ScriptedModel([AIMessage(content="", tool_calls=[_tc("dump", {}, "c1")]), after]) + + result = await _run(model, _config("fc-big", cuga_lite_execution_mode="function_calling"), tool) + + assert result["final_answer"] == "Got it." + assert len(result["variables_storage"]["tool_result_c1"]["value"]) == limit + 100, ( + "the whole result survives as a variable for later turns" + ) + + +@pytest.mark.asyncio +async def test_a_tool_whose_name_is_not_an_identifier_round_trips(): + async def dashed(q: str) -> str: + """Dashed.""" + CALLS.append(("my-tool", q)) + return "ok:" + q + + tool = StructuredTool.from_function(coroutine=dashed, name="my-tool", description="Dashed.") + model = _ScriptedModel( + [AIMessage(content="", tool_calls=[_tc("my-tool", {"q": "hi"}, "c1")]), AIMessage(content="done")] + ) + + result = await _run(model, _config("fc-dash", cuga_lite_execution_mode="function_calling"), tool) + + assert CALLS == [("my-tool", "hi")] + reply = [m for m in result["chat_messages"] if isinstance(m, ToolMessage)][0] + assert (reply.tool_call_id, reply.name, reply.content) == ("c1", "my-tool", "ok:hi") diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py deleted file mode 100644 index 003a766e5..000000000 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_tool_exec_node.py +++ /dev/null @@ -1,516 +0,0 @@ -"""``tool_exec`` node: every ``tool_call`` gets a ``ToolMessage``, inside the guard envelope. - -The node is the function-calling counterpart of the sandbox, so the assertions -mirror what #560 pinned for the sandbox: budgets and the tracker apply to every -call even though the callables in ``_tools_context`` are bare, and every exit -carries the budget fields. On top of that, the FC-specific contract: ids are -preserved, one failing call never aborts its siblings, provider-side invalid -calls are answered, and step discipline defers everything but the first call. -""" - -from __future__ import annotations - -import asyncio -from types import SimpleNamespace -from typing import Any, List, Optional -from unittest.mock import MagicMock - -import pytest -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage -from langgraph.types import Command - -from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.tool_exec_node import ( - DEFERRED_CALL_MESSAGE, - create_tool_exec_node, -) -from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module -from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( - ToolCallTracker, - make_recording_awaitable, -) - -pytestmark = pytest.mark.unit - - -class _Adapter: - messages_key = "chat_messages" - metadata_key = "cuga_lite_metadata" - execute_node_name = "sandbox" - sender_name = "CugaLite" - - def __init__(self, tools: dict, max_steps: int = 50): - self._tools_context = tools - self._tracker = MagicMock() - self._model = None - self._max_steps = max_steps - - def get_messages(self, state: Any) -> List[BaseMessage]: - return list(state.chat_messages or []) - - def resolve_max_steps(self, state: Any, override: Optional[int]) -> int: - return override if override is not None else self._max_steps - - -def _state(last: AIMessage, step_count: int = 0): - return SimpleNamespace( - chat_messages=[HumanMessage(content="add 1 and 2"), last], - step_count=step_count, - tool_calls=[], - tool_calls_used_run=0, - tool_calls_used_thread=0, - ) - - -def _call(name: str, args: dict, call_id: str) -> dict: - return {"name": name, "args": args, "id": call_id, "type": "tool_call"} - - -def _update(result) -> dict: - """Normal exits are ``Command(goto="call_model")``; error exits ``Command(goto=END)``.""" - assert isinstance(result, Command), type(result) - return result.update - - -def _tool_messages(result) -> List[ToolMessage]: - return [m for m in _update(result)["chat_messages"] if isinstance(m, ToolMessage)] - - -@pytest.fixture(autouse=True) -def _reset_budget_contexts(): - yield - tracker_module._tool_call_budget_context.set(None) - tracker_module._thread_tool_call_budget_context.set(None) - tracker_module._block_tool_call_budget_context.set(None) - tracker_module._block_tool_call_cap_override_context.set(None) - - -def _caps(monkeypatch, *, block=0, run=0, thread=0): - from cuga.config import settings - - monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_block", block, raising=False) - monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_run", run, raising=False) - monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_thread", thread, raising=False) - - -# ── happy path ─────────────────────────────────────────────────────────────── - - -@pytest.mark.asyncio -async def test_every_call_runs_and_ids_are_preserved(monkeypatch): - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - def greet(name: str) -> str: # sync tools are made awaitable, like the sandbox does - return f"hi {name}" - - node = create_tool_exec_node(_Adapter({"add": add, "greet": greet})) - last = AIMessage( - content="", tool_calls=[_call("add", {"a": 1, "b": 2}, "c1"), _call("greet", {"name": "bob"}, "c2")] - ) - - result = await node(_state(last), config={"configurable": {}}) - - msgs = _tool_messages(result) - assert [(m.tool_call_id, m.name, m.content, m.status) for m in msgs] == [ - ("c1", "add", "3", "success"), - ("c2", "greet", "hi bob", "success"), - ] - assert _update(result)["chat_messages"][:2] == _state(last).chat_messages, ( - "history is appended to, never rewritten" - ) - assert _update(result)["script"] is None - assert _update(result)["step_count"] == 1 - assert _update(result)["tool_calls_used_run"] == 2 and _update(result)["tool_budget_exhausted"] is False - - -@pytest.mark.asyncio -async def test_raw_provider_shaped_calls_with_json_string_args_are_decoded(monkeypatch): - """LangChain parses ``args`` into a dict before an AIMessage is built, so a - JSON-string payload only arrives in the raw provider shape (``function.arguments``). - ``model_construct`` bypasses validation to put exactly that on the message.""" - _caps(monkeypatch) - seen = [] - - async def add(a: int, b: int) -> int: - seen.append((a, b)) - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage.model_construct( - content="", - tool_calls=[ - {"id": "c1", "function": {"name": "add", "arguments": '{"a": 5, "b": 6}'}}, - {"id": "c2", "function": {"name": "add", "arguments": "{not json"}}, - ], - invalid_tool_calls=[], - additional_kwargs={}, - response_metadata={}, - ) - - result = await node(_state(last), config={"configurable": {}}) - - m1, m2 = _tool_messages(result) - assert (m1.tool_call_id, m1.content, m1.status) == ("c1", "11", "success") and seen == [(5, 6)] - assert m2.tool_call_id == "c2" and m2.status == "error" and "valid JSON" in m2.content - - -# ── every failure is a ToolMessage, never an exception ─────────────────────── - - -@pytest.mark.asyncio -async def test_unknown_tool_lists_the_known_names(monkeypatch): - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"add": add, "_internal": add})) - result = await node(_state(AIMessage(content="", tool_calls=[_call("subtract", {}, "c1")])), config=None) - - (m,) = _tool_messages(result) - assert m.status == "error" and m.tool_call_id == "c1" - assert "Unknown tool 'subtract'" in m.content and "add" in m.content - assert "_internal" not in m.content, "internal injections are not offered to the model" - - -@pytest.mark.asyncio -async def test_a_failing_call_does_not_abort_its_siblings(monkeypatch): - _caps(monkeypatch) - - async def boom(x: int) -> int: - raise ValueError("nope") - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"boom": boom, "add": add})) - last = AIMessage( - content="", - tool_calls=[ - _call("boom", {"x": 1}, "c1"), - _call("add", {"wrong": 1}, "c2"), - _call("add", {"a": 2, "b": 2}, "c3"), - ], - ) - - result = await node(_state(last), config=None) - - m1, m2, m3 = _tool_messages(result) - assert m1.status == "error" and "nope" in m1.content - assert m2.status == "error" and "rejected these arguments" in m2.content - assert (m3.content, m3.status) == ("4", "success") - - -@pytest.mark.asyncio -async def test_timeout_is_an_error_toolmessage(monkeypatch): - _caps(monkeypatch) - from cuga.config import settings - - monkeypatch.setattr(settings.advanced_features, "sandbox_execution_timeout", 0.05, raising=False) - - async def slow() -> str: - await asyncio.Event().wait() # never set: immune to any asyncio.sleep patch elsewhere in the suite - return "late" - - node = create_tool_exec_node(_Adapter({"slow": slow})) - result = await node(_state(AIMessage(content="", tool_calls=[_call("slow", {}, "c1")])), config=None) - - (m,) = _tool_messages(result) - assert m.status == "error" and "timed out" in m.content and m.tool_call_id == "c1" - - -@pytest.mark.asyncio -async def test_provider_side_invalid_tool_calls_are_answered(monkeypatch): - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage( - content="", - tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")], - invalid_tool_calls=[ - { - "name": "add", - "args": "{broken", - "id": "bad1", - "error": "unterminated", - "type": "invalid_tool_call", - } - ], - ) - - result = await node(_state(last), config=None) - - good, bad = _tool_messages(result) - assert (good.tool_call_id, good.content) == ("c1", "2") - assert bad.tool_call_id == "bad1" and bad.status == "error" and "unterminated" in bad.content - - -# ── guard envelope ─────────────────────────────────────────────────────────── - - -@pytest.mark.asyncio -async def test_bare_callables_are_still_budget_capped(monkeypatch): - """The callables in _tools_context carry no budget wrapper — the node must add it.""" - _caps(monkeypatch, run=1) - calls = [] - - async def add(a: int, b: int) -> int: - calls.append((a, b)) - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage( - content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] - ) - - result = await node(_state(last), config=None) - - m1, m2 = _tool_messages(result) - assert m1.content == "2" and calls == [(1, 1)], "the second call must be refused, not executed" - assert m2.status == "error" and "Tool call limit reached" in m2.content - assert _update(result)["tool_calls_used_run"] == 1, "a refused call never inflates the counter" - assert _update(result)["tool_budget_exhausted"] is True, ( - "call_model must see the exhaustion and end the turn" - ) - - -@pytest.mark.asyncio -async def test_tracking_session_wraps_the_calls(monkeypatch): - """A recorded tool lands in result['tool_calls'] when tracking is requested, and - the session is closed afterwards — the receipt cannot leak into the next node.""" - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"add": make_recording_awaitable(add, "add", app_name="calc")})) - last = AIMessage(content="", tool_calls=[_call("add", {"a": 3, "b": 4}, "c1")]) - - result = await node(_state(last), config={"configurable": {"track_tool_calls": True}}) - - assert [(c["name"], c["app_name"], c["result"]) for c in _update(result)["tool_calls"]] == [ - ("add", "calc", 7) - ] - assert ToolCallTracker.is_enabled() is False - - untracked = await node(_state(last), config={"configurable": {}}) - assert _update(untracked)["tool_calls"] == [] - - -# ── step discipline ────────────────────────────────────────────────────────── - - -@pytest.mark.asyncio -async def test_step_discipline_runs_only_the_first_call_and_defers_the_rest(monkeypatch): - _caps(monkeypatch) - calls = [] - - async def add(a: int, b: int) -> int: - calls.append((a, b)) - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage( - content="", - tool_calls=[ - _call("add", {"a": 1, "b": 1}, "c1"), - _call("add", {"a": 2, "b": 2}, "c2"), - _call("add", {"a": 3, "b": 3}, "c3"), - ], - ) - - result = await node( - _state(last), config={"configurable": {"cuga_lite_step_discipline": "one_tool_per_step"}} - ) - - m1, m2, m3 = _tool_messages(result) - assert calls == [(1, 1)] - assert (m1.tool_call_id, m1.content, m1.status) == ("c1", "2", "success") - assert (m2.tool_call_id, m2.content, m2.status) == ("c2", DEFERRED_CALL_MESSAGE, "error") - assert (m3.tool_call_id, m3.content, m3.status) == ("c3", DEFERRED_CALL_MESSAGE, "error") - assert _update(result)["tool_calls_used_run"] == 1 - - -@pytest.mark.asyncio -async def test_step_discipline_off_runs_every_call(monkeypatch): - _caps(monkeypatch) - calls = [] - - async def add(a: int, b: int) -> int: - calls.append((a, b)) - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage( - content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] - ) - - await node(_state(last), config={"configurable": {"cuga_lite_step_discipline": "off"}}) - assert calls == [(1, 1), (2, 2)] - - -# ── routing edge cases ─────────────────────────────────────────────────────── - - -@pytest.mark.asyncio -async def test_nothing_executable_still_advances_the_step(monkeypatch): - """A mis-route into tool_exec must not become an infinite no-op loop.""" - _caps(monkeypatch) - node = create_tool_exec_node(_Adapter({})) - - result = await node(_state(AIMessage(content="no calls here"), step_count=4), config=None) - - assert result.goto == "call_model" - assert _update(result)["step_count"] == 5 and _update(result)["script"] is None - assert "chat_messages" not in _update(result), "no ToolMessage is fabricated" - - -@pytest.mark.asyncio -async def test_step_limit_ends_the_run_with_an_error_command(monkeypatch): - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"add": add}, max_steps=3)) - last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) - - result = await node(_state(last, step_count=3), config=None) - - assert isinstance(result, Command) - assert "Maximum step limit" in result.update["error"] - assert result.update["execution_complete"] is True - assert "tool_calls_used_run" in result.update, "budget fields ride every exit" - - -# ── outer guard + batch output budget ──────────────────────────────────────── - - -@pytest.mark.asyncio -async def test_unexpected_failure_after_the_batch_still_reports_the_budget(monkeypatch): - """Same contract as the sandbox: whatever breaks after a call ran, the update - carries the spent budget, or the checkpoint keeps the pre-batch counts.""" - _caps(monkeypatch) - from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import tool_exec_node as mod - - async def add(a: int, b: int) -> int: - return a + b - - def boom(*args, **kwargs): - raise RuntimeError("history store unavailable") - - monkeypatch.setattr(mod, "core_append_with_step_limit", boom) - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) - - result = await node(_state(last), config=None) - - assert isinstance(result, Command) - assert "history store unavailable" in result.update["error"] - assert result.update["tool_calls_used_run"] == 1, "the call that ran must still be counted" - assert ToolCallTracker.is_enabled() is False - replies = [m for m in result.update["chat_messages"] if isinstance(m, ToolMessage)] - assert [(m.tool_call_id, m.content) for m in replies] == [("c1", "2")], ( - "the executed call's reply is kept" - ) - - -@pytest.mark.asyncio -async def test_outer_guard_answers_the_calls_the_batch_never_reached(monkeypatch): - _caps(monkeypatch) - from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import tool_exec_node as mod - - async def add(a: int, b: int) -> int: - return a + b - - async def broken_batch(adapter, calls, invalid, results, **kwargs): - results.append(ToolMessage(content="2", tool_call_id="c1", name="add")) - raise RuntimeError("executor died") - - monkeypatch.setattr(mod, "_run_batch", broken_batch) - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage( - content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")] - ) - - result = await node(_state(last), config=None) - - replies = {m.tool_call_id: m for m in result.update["chat_messages"] if isinstance(m, ToolMessage)} - assert replies["c1"].content == "2" - assert replies["c2"].status == "error" and "executor died" in replies["c2"].content, "no dangling id" - - -@pytest.mark.asyncio -async def test_one_batch_shares_one_output_limit(monkeypatch): - """N results share execution_output_max_length, like the N prints of one block.""" - _caps(monkeypatch) - from cuga.config import settings - - monkeypatch.setattr(settings.advanced_features, "execution_output_max_length", 10, raising=False) - - async def text(n: int) -> str: - return "x" * n - - node = create_tool_exec_node(_Adapter({"text": text})) - last = AIMessage( - content="", - tool_calls=[ - _call("text", {"n": 8}, "c1"), - _call("text", {"n": 8}, "c2"), - _call("text", {"n": 3}, "c3"), - ], - ) - - result = await node(_state(last), config=None) - - m1, m2, m3 = _tool_messages(result) - assert m1.content == "x" * 8, "under the limit: untouched" - assert m2.content.startswith("xx") and "truncated" in m2.content, "only the remaining budget is spent" - assert "truncated" in m3.content and not m3.content.startswith("x"), "nothing left for the third result" - - -@pytest.mark.asyncio -async def test_normal_exit_routes_to_call_model_and_error_exit_to_end(monkeypatch): - """One routing mechanism: no static edge, so a terminal error never schedules - one more model turn.""" - _caps(monkeypatch) - - async def add(a: int, b: int) -> int: - return a + b - - node = create_tool_exec_node(_Adapter({"add": add}, max_steps=3)) - last = AIMessage(content="", tool_calls=[_call("add", {"a": 1, "b": 1}, "c1")]) - - assert (await node(_state(last), config=None)).goto == "call_model" - from langgraph.graph import END - - assert (await node(_state(last, step_count=3), config=None)).goto == END - - -@pytest.mark.asyncio -async def test_step_discipline_defers_by_position_even_when_the_first_call_fails(monkeypatch): - """An error on the first call is the result the model must read before the next - call — the second call must not run just because the first one failed.""" - _caps(monkeypatch) - calls = [] - - async def add(a: int, b: int) -> int: - calls.append((a, b)) - return a + b - - node = create_tool_exec_node(_Adapter({"add": add})) - last = AIMessage(content="", tool_calls=[_call("nope", {}, "c1"), _call("add", {"a": 2, "b": 2}, "c2")]) - - result = await node( - _state(last), config={"configurable": {"cuga_lite_step_discipline": "one_tool_per_step"}} - ) - - m1, m2 = _tool_messages(result) - assert m1.status == "error" and "Unknown tool" in m1.content - assert (m2.tool_call_id, m2.content, m2.status) == ("c2", DEFERRED_CALL_MESSAGE, "error") - assert calls == [] diff --git a/src/cuga/backend/cuga_graph/policy/storage.py b/src/cuga/backend/cuga_graph/policy/storage.py index 38f342bf6..4a8d25f98 100644 --- a/src/cuga/backend/cuga_graph/policy/storage.py +++ b/src/cuga/backend/cuga_graph/policy/storage.py @@ -417,15 +417,7 @@ async def list_policies( policy_type: Optional[PolicyType] = None, enabled_only: bool = True, limit: int = 100, - *, - strict: bool = False, ) -> List[Policy]: - """List stored policies. - - ``strict=True`` re-raises a backend failure instead of returning ``[]``. - A safety guard that must fail closed cannot tell "no policies" from - "the backend is down" otherwise. - """ if not self._connected: await self.initialize_async() try: @@ -434,8 +426,6 @@ async def list_policies( policies.sort(key=lambda x: x.priority, reverse=True) return policies except Exception as e: - if strict: - raise logger.error(f"Failed to list policies: {e}") return [] diff --git a/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py b/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py deleted file mode 100644 index 0eb93115d..000000000 --- a/src/cuga/backend/cuga_graph/policy/tests/test_storage_strict_list.py +++ /dev/null @@ -1,34 +0,0 @@ -"""``PolicyStorage.list_policies(strict=True)`` surfaces backend failures. - -The default swallows them and returns ``[]`` — fine for lookups, wrong for a -guard that must fail closed and cannot tell "no policies" from "backend down". -""" - -from __future__ import annotations - -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from cuga.backend.cuga_graph.policy.storage import PolicyStorage - -pytestmark = pytest.mark.unit - - -def _storage_with_broken_backend() -> PolicyStorage: - storage = PolicyStorage.__new__(PolicyStorage) - storage._connected = True - storage._backend = MagicMock() - storage._backend.list_policies = AsyncMock(side_effect=RuntimeError("backend down")) - return storage - - -@pytest.mark.asyncio -async def test_default_swallows_backend_errors(): - assert await _storage_with_broken_backend().list_policies() == [] - - -@pytest.mark.asyncio -async def test_strict_raises_backend_errors(): - with pytest.raises(RuntimeError, match="backend down"): - await _storage_with_broken_backend().list_policies(strict=True) diff --git a/src/cuga/backend/cuga_graph/utils/agent_loop.py b/src/cuga/backend/cuga_graph/utils/agent_loop.py index b22ff9001..9be5c77e4 100644 --- a/src/cuga/backend/cuga_graph/utils/agent_loop.py +++ b/src/cuga/backend/cuga_graph/utils/agent_loop.py @@ -37,8 +37,10 @@ def _native_tool_calls(messages: Any) -> list: return list(calls or []) -def _tool_exec_step_output(messages: Any) -> str: - """Render the trailing run of tool results (one per native call) as execution output.""" +def _tool_results_step_output(messages: Any) -> str: + """Render the trailing run of tool results (one per native call) as execution output. + + Empty when the step did not end on tool results, e.g. a CodeAct block.""" lines: list = [] for msg in reversed(messages or []): role = msg.get("type") if isinstance(msg, dict) else getattr(msg, "type", None) @@ -477,14 +479,9 @@ def get_event_message(self, event) -> StreamEvent: f"Returning sandbox output with execution_output length: {len(execution_output)}" ) return StreamEvent(name="CodeAgent", data=json.dumps(output)) - else: - # Skip empty sandbox events - logger.debug("Skipping empty sandbox event") - return StreamEvent(name="", data="") - - # Function-calling mode: the step appended one ToolMessage per native call. - elif node_name == "tool_exec": - tool_output = _tool_exec_step_output(subgraph_messages or []) + # Function-calling turn through the sandbox: the step appended one + # ToolMessage per native call instead of an "Execution output:" message. + tool_output = _tool_results_step_output(subgraph_messages or []) if tool_output.strip(): output = { "code": "", @@ -494,7 +491,8 @@ def get_event_message(self, event) -> StreamEvent: "variables": state_data.get("variables_storage", {}), } return StreamEvent(name="CodeAgent", data=json.dumps(output)) - logger.debug("Skipping empty tool_exec event") + # Skip empty sandbox events + logger.debug("Skipping empty sandbox event") return StreamEvent(name="", data="") # Default handling for other subgraph nodes diff --git a/src/cuga/settings.toml b/src/cuga/settings.toml index f689500d9..5ecd3d0d5 100644 --- a/src/cuga/settings.toml +++ b/src/cuga/settings.toml @@ -55,7 +55,7 @@ cuga_lite_bind_tools_max_count = 128 # Pad shortlister output to fill the cap. Off by default for code-execution mode. cuga_lite_bind_tools_pad_to_cap = false # Execution mode: codeact (write Python, run it in the sandbox — default) | function_calling (the model -# issues native tool_calls, a tool_exec node runs them and replies with ToolMessages). Alias: fc. +# issues native tool_calls, the sandbox runs them and replies with ToolMessages). Alias: fc. # Per-invoke: config["configurable"]["cuga_lite_execution_mode"]. Env: DYNACONF_ADVANCED_FEATURES__CUGA_LITE_EXECUTION_MODE cuga_lite_execution_mode = "codeact" # Step discipline, both modes: off (default) | one_tool_per_step — observe one result before the next call. diff --git a/tests/unit/test_agent_loop_fc_events.py b/tests/unit/test_agent_loop_fc_events.py index 60c6308f7..47ebfac38 100644 --- a/tests/unit/test_agent_loop_fc_events.py +++ b/tests/unit/test_agent_loop_fc_events.py @@ -1,8 +1,9 @@ """Function-calling steps stream to the UI like code steps. -``get_event_message`` maps subgraph nodes to StreamEvents; without these branches a -native tool-call turn and its results were an empty ``tool_exec`` event and the UI -showed nothing between the question and the answer. +``get_event_message`` maps subgraph nodes to StreamEvents. A native tool-call turn +is shown as the code a block would have shown, and the sandbox step that answers it +appends ToolMessages instead of an "Execution output:" message — without the +tool-result branch the UI showed nothing between the question and the answer. """ from __future__ import annotations @@ -34,18 +35,31 @@ def test_native_tool_calls_stream_as_a_code_step(): assert json.loads(ev.data)["code"] == 'echo({"value": 7})' -def test_tool_results_stream_as_execution_output(): +def test_tool_results_from_the_sandbox_stream_as_execution_output(): messages = [ HumanMessage(content="q"), AIMessage(content="", tool_calls=[_CALL, {**_CALL, "id": "c2", "args": {"value": 8}}]), ToolMessage(content="7", tool_call_id="c1", name="echo"), ToolMessage(content="8", tool_call_id="c2", name="echo"), ] - ev = AgentLoop.get_event_message(MagicMock(), _event("tool_exec", messages)) + ev = AgentLoop.get_event_message(MagicMock(), _event("sandbox", messages)) assert ev.name == "CodeAgent" - assert json.loads(ev.data)["execution_output"] == "echo: 7\necho: 8" + data = json.loads(ev.data) + assert data["execution_output"] == "echo: 7\necho: 8" and data["summary"] == "Tool calls completed" -def test_tool_exec_without_results_is_an_empty_event(): - ev = AgentLoop.get_event_message(MagicMock(), _event("tool_exec", [HumanMessage(content="q")])) +def test_codeact_sandbox_output_is_unchanged(): + messages = [ + HumanMessage(content="q"), + AIMessage(content="```python\nprint(1)\n```"), + HumanMessage(content="Execution output:\n1"), + ] + ev = AgentLoop.get_event_message(MagicMock(), _event("sandbox", messages)) + assert ev.name == "CodeAgent" + data = json.loads(ev.data) + assert data["execution_output"] == "1" and data["summary"] == "Code execution completed" + + +def test_sandbox_step_without_results_is_an_empty_event(): + ev = AgentLoop.get_event_message(MagicMock(), _event("sandbox", [HumanMessage(content="q")])) assert ev.name == "" and ev.data == "" diff --git a/tests/unit/test_fc_sdk_fail_closed_approval.py b/tests/unit/test_fc_sdk_fail_closed_approval.py deleted file mode 100644 index 70f680712..000000000 --- a/tests/unit/test_fc_sdk_fail_closed_approval.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Function-calling mode refuses to start while a tool-approval policy is stored. - -Through the real SDK path: a policy added with ``agent.policies.add_tool_approval`` -lands in policy storage, and the function-calling turn queries that storage before -its first model call. CodeAct on the same agent still goes through the approval -prompt, which is the behaviour the refusal protects. -""" - -from __future__ import annotations - -import pytest -from langchain_core.messages import AIMessage -from langchain_core.tools import StructuredTool - -from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import FC_TOOL_APPROVAL_UNSUPPORTED - -pytestmark = pytest.mark.unit - - -class _Bound: - """What ``bind_tools`` returns: a runnable wrapping the model.""" - - def __init__(self, model): - self._model = model - - async def ainvoke(self, messages, config=None, **kwargs): - return await self._model.ainvoke(messages, config=config, **kwargs) - - -class _ScriptedModel: - def __init__(self, responses): - self._responses = list(responses) - self.invocations = 0 - - def bind_tools(self, tools, **kwargs): - return _Bound(self) # a real bind returns a new runnable, never the model itself - - async def ainvoke(self, messages, config=None, **kwargs): - self.invocations += 1 - if not self._responses: - raise AssertionError("model asked for more responses than scripted") - return self._responses.pop(0) - - -RAN: list = [] - - -def _delete_record(): - async def delete_record(record_id: int) -> str: - """Delete a record.""" - RAN.append(record_id) - return "deleted" - - return StructuredTool.from_function( - coroutine=delete_record, name="delete_record", description="Delete a record." - ) - - -@pytest.fixture(autouse=True) -def _settings(monkeypatch): - from cuga.config import settings - - monkeypatch.setattr(settings.policy, "enabled", True, raising=False) - monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) - RAN.clear() - yield - RAN.clear() - - -@pytest.mark.asyncio -async def test_fc_refuses_with_a_stored_tool_approval_policy_and_codeact_still_asks(): - from cuga.sdk import CugaAgent - - fc_model = _ScriptedModel( - [ - AIMessage( - content="", - tool_calls=[ - {"name": "delete_record", "args": {"record_id": 1}, "id": "c1", "type": "tool_call"} - ], - ) - ] - ) - agent = CugaAgent( - tools=[_delete_record()], model=fc_model, execution_mode="function_calling", reset_policy_storage=True - ) - await agent.policies.add_tool_approval( - name="Approve deletions", - required_tools=["delete_record"], - approval_message="Deleting needs approval.", - ) - - result = await agent.invoke("delete record 1", thread_id="fc-approval") - - assert fc_model.invocations == 0, "refused before the model was invoked" - assert ( - FC_TOOL_APPROVAL_UNSUPPORTED in (result.error or "") or FC_TOOL_APPROVAL_UNSUPPORTED in result.answer - ) - - # Same agent, CodeAct for this call: the guarded code goes to the approval prompt. - codeact_model = _ScriptedModel( - [AIMessage(content="```python\nprint(await delete_record(record_id=1))\n```")] - ) - agent._model = codeact_model - agent._graph = None - agent._compiled_graph = None - result = await agent.invoke("delete record 1", thread_id="codeact-approval", execution_mode="codeact") - - assert codeact_model.invocations == 1 - assert "Deleting needs approval." in result.answer, result.answer - assert RAN == [], "the guarded tool must not run on either path" diff --git a/tests/unit/test_fc_sdk_two_turn.py b/tests/unit/test_fc_sdk_two_turn.py index 3434b539b..4bbbdbab9 100644 --- a/tests/unit/test_fc_sdk_two_turn.py +++ b/tests/unit/test_fc_sdk_two_turn.py @@ -16,6 +16,8 @@ from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.tools import StructuredTool +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.graph_adapter import FC_VARIABLES_NOTE + pytestmark = pytest.mark.unit CALLS: list = [] @@ -117,4 +119,9 @@ async def test_second_turn_replays_a_valid_function_calling_transcript(): tool_replies = [m for m in outbound if isinstance(m, ToolMessage)] dangling = {c["id"] for m in ai_with_calls for c in m.tool_calls} - {t.tool_call_id for t in tool_replies} assert not dangling, f"tool_calls without a ToolMessage reply: {dangling} in {shapes}" - assert isinstance(outbound[-1], HumanMessage) and outbound[-1].content == "now echo 8" + # Turn 1's result survived as a variable, so turn 2 is told about it — in function-calling + # words, not the CodeAct invitation to write code. + assert isinstance(outbound[-1], HumanMessage) and outbound[-1].content.startswith("now echo 8") + assert "tool_result_call_1" in outbound[-1].content and FC_VARIABLES_NOTE in outbound[-1].content + assert "directly by their names" not in outbound[-1].content + assert turn1.variables.get("tool_result_call_1") == 7, turn1.variables From fd27513be070d775aa6e4c1fe78fabca2fe3f2b7 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:31:45 +0300 Subject: [PATCH 14/18] test(cuga-lite): translator, sandbox-reply and equivalence tests for function calling The 17 tool_exec fixtures run through the block on the real local executor with the same expected replies, plus a seeded 400-case random test compared reply for reply against a reference of the direct-call contract, and SDK tests for approval pause, approve-and-resume and denial. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/tests/test_fc_actions.py | 491 ++++++++++++++++++ tests/unit/test_fc_sdk_tool_approval.py | 144 +++++ 2 files changed, 635 insertions(+) create mode 100644 src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py create mode 100644 tests/unit/test_fc_sdk_tool_approval.py diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py new file mode 100644 index 000000000..58af97db2 --- /dev/null +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py @@ -0,0 +1,491 @@ +"""Function-calling through the sandbox: the translator, the reply builder, and equivalence. + +The direct-call contract these pin is the one ``tool_exec`` shipped with: ids +preserved, one failing call never aborts its siblings, provider-side invalid +calls answered, step discipline defers everything but the first call, one output +budget per batch, and the exact reply texts. ``_run_batch`` here is what the +sandbox node does with a plan — translate, execute the block on the real local +executor, read the replies back — so every fixture crosses the executor. + +The last test is the equivalence property: thousands of random tool names, +JSON arguments, results and exceptions through the block, compared reply for +reply against a reference of the direct-call contract. +""" + +from __future__ import annotations + +import asyncio +import json +import random +import string +from typing import Any, Dict, List, Tuple + +import pytest +from langchain_core.messages import ToolMessage + +from cuga.backend.cuga_graph.nodes.cuga_agent_core.execution.code_extraction import make_tool_awaitable +from cuga.backend.cuga_graph.nodes.cuga_lite.adapter.fc_actions import ( + DEFERRED_CALL_MESSAGE, + FC_TOOLS_KEY, + TRUNCATION_MARKER, + _coerce_args, + _stringify, + _tool_call_parts, + fc_context_overlay, + plan_tool_calls, + replies_from_execution, + replies_without_execution, +) +from cuga.backend.cuga_graph.nodes.cuga_lite.cuga_lite_graph import CugaLiteState +from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import CodeExecutor +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking import tracker as tracker_module +from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ToolCallTracker + +pytestmark = pytest.mark.unit + +RAN: list = [] + + +async def add(a: int, b: int) -> int: + RAN.append(("add", a, b)) + return a + b + + +async def boom(x: Any = None) -> None: + RAN.append(("boom", x)) + raise ValueError("bad x") + + +async def echo_args(**kwargs: Any) -> Dict[str, Any]: + RAN.append(("echo_args", kwargs)) + return kwargs + + +async def typed(a: int) -> Dict[str, int]: + RAN.append(("typed", a)) + return {"a": a} + + +async def dashed(q: str) -> str: + RAN.append(("dashed", q)) + return "dashed:" + q + + +async def slow(seconds: float) -> str: + await asyncio.sleep(seconds) + return "slept" + + +async def as_set() -> set: + return {1} + + +def _tools(*fns, **named) -> Dict[str, Any]: + tools = {fn.__name__: make_tool_awaitable(fn) for fn in fns} + tools.update({name: make_tool_awaitable(fn) for name, fn in named.items()}) + return tools + + +def _call(name: str, args: Any, call_id: str) -> dict: + return {"name": name, "args": args, "id": call_id, "type": "tool_call"} + + +@pytest.fixture(autouse=True) +def _reset(): + RAN.clear() + yield + RAN.clear() + tracker_module._tool_call_budget_context.set(None) + tracker_module._thread_tool_call_budget_context.set(None) + tracker_module._block_tool_call_budget_context.set(None) + tracker_module._block_tool_call_cap_override_context.set(None) + + +def _caps(monkeypatch, *, block=0, run=0, thread=0): + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_block", block, raising=False) + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_run", run, raising=False) + monkeypatch.setattr(settings.advanced_features, "max_tool_calls_per_thread", thread, raising=False) + + +async def _run_batch( + tools: Dict[str, Any], + calls: List[dict], + invalid: List[dict] | None = None, + *, + one_per_step: bool = False, + output_limit: int = 0, + timeout: float = 30, +) -> Tuple[List[ToolMessage], Dict[str, Any], str | None]: + """What the sandbox node does with a function-calling plan, on the real local executor.""" + block, plan = plan_tool_calls(calls, list(invalid or []), tools, one_per_step=one_per_step) + if block is None: + return replies_without_execution(plan, ""), {}, None + context = {**tools, **fc_context_overlay(tools, plan)} + ToolCallTracker.start_tracking(enabled=False, timings_only=False) + ToolCallTracker.seed_call_budget(0, 0) + try: + output, new_vars = await CodeExecutor.eval_with_tools_async( + code=block, + _locals=context, + state=CugaLiteState(chat_messages=[]), + mode="local", + keep_listing_vars=True, + ) + finally: + ToolCallTracker.stop_tracking() + replies, keep, _drop = replies_from_execution( + plan, new_vars, output, output_limit=output_limit, timeout=timeout + ) + return replies, keep, block + + +def _shape(replies: List[ToolMessage]) -> List[Tuple[str, str, str, str]]: + return [(m.tool_call_id, m.name, m.status, m.content) for m in replies] + + +# ── the direct-call contract, through the block ───────────────────────────── + + +@pytest.mark.asyncio +async def test_every_call_runs_and_ids_are_preserved(): + replies, kept, block = await _run_batch( + _tools(add), [_call("add", {"a": 1, "b": 2}, "c1"), _call("add", {"a": 3, "b": 4}, "c2")] + ) + assert _shape(replies) == [("c1", "add", "success", "3"), ("c2", "add", "success", "7")] + assert RAN == [("add", 1, 2), ("add", 3, 4)] + assert kept == {"tool_result_c1": 3, "tool_result_c2": 7}, "results persist as variables" + assert "await add(**" in block and FC_TOOLS_KEY not in block, "a real name is called as itself" + + +@pytest.mark.asyncio +async def test_raw_provider_shaped_calls_with_json_string_args_are_decoded(): + raw = {"id": "c1", "type": "function", "function": {"name": "add", "arguments": '{"a": 1, "b": 2}'}} + replies, _, _ = await _run_batch(_tools(add), [raw]) + assert _shape(replies) == [("c1", "add", "success", "3")] + + +@pytest.mark.asyncio +async def test_unknown_tool_lists_the_known_names(): + replies, _, block = await _run_batch(_tools(add, boom), [_call("nope", {}, "c1")]) + assert block is None, "nothing to execute" + assert _shape(replies) == [ + ("c1", "nope", "error", "Unknown tool 'nope'. Choose one of the provided tools: add, boom.") + ] + + +@pytest.mark.asyncio +async def test_a_failing_call_does_not_abort_its_siblings(): + replies, kept, _ = await _run_batch( + _tools(add, boom), [_call("boom", {"x": 1}, "c1"), _call("add", {"a": 1, "b": 2}, "c2")] + ) + assert _shape(replies) == [ + ("c1", "boom", "error", "Tool 'boom' failed: bad x"), + ("c2", "add", "success", "3"), + ] + assert kept == {"tool_result_c2": 3}, "the error marker is not kept as a variable" + + +@pytest.mark.asyncio +async def test_wrong_arguments_are_reported_as_rejected(): + replies, _, _ = await _run_batch(_tools(typed), [_call("typed", {"a": 1, "zzz": 2}, "c1")]) + ((call_id, name, status, content),) = _shape(replies) + assert (call_id, name, status) == ("c1", "typed", "error") + assert content.startswith("Tool 'typed' rejected these arguments: ") and content.endswith( + ". Check the parameter names and types." + ) + + +@pytest.mark.asyncio +async def test_unparsable_json_arguments_are_an_error_reply(): + raw = {"id": "c1", "type": "function", "function": {"name": "add", "arguments": "{not json"}} + replies, _, block = await _run_batch(_tools(add), [raw]) + assert block is None + ((_, _, status, content),) = _shape(replies) + assert status == "error" and content.startswith("Could not parse tool arguments as JSON") + + +@pytest.mark.asyncio +async def test_block_timeout_keeps_finished_results_and_answers_the_rest(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.advanced_features, "sandbox_execution_timeout", 0.3, raising=False) + replies, kept, _ = await _run_batch( + _tools(add, slow), + [ + _call("add", {"a": 1, "b": 2}, "c1"), + _call("slow", {"seconds": 5}, "c2"), + _call("add", {"a": 3, "b": 4}, "c3"), + ], + timeout=0.3, + ) + timed_out = "Tool '{name}' timed out after 0.3s. Try a narrower call or a different tool." + assert _shape(replies) == [ + ("c1", "add", "success", "3"), + ("c2", "slow", "error", timed_out.format(name="slow")), + ("c3", "add", "error", timed_out.format(name="add")), + ] + assert kept.get("tool_result_c1") == 3, "the result computed before the timeout is kept" + assert RAN == [("add", 1, 2)], "nothing after the stalled call ran" + + +@pytest.mark.asyncio +async def test_provider_side_invalid_tool_calls_are_answered(): + invalid = [ + { + "name": "add", + "args": "{bad", + "id": "bad1", + "error": "Unterminated string", + "type": "invalid_tool_call", + } + ] + replies, _, _ = await _run_batch(_tools(add), [_call("add", {"a": 1, "b": 2}, "c1")], invalid) + assert _shape(replies) == [ + ("c1", "add", "success", "3"), + ( + "bad1", + "add", + "error", + "The provider could not parse this tool call (Unterminated string). Re-issue it with valid arguments.", + ), + ] + + +@pytest.mark.asyncio +async def test_bare_callables_are_still_budget_capped(monkeypatch): + _caps(monkeypatch, run=1) + replies, _, _ = await _run_batch( + _tools(add), [_call("add", {"a": 1, "b": 2}, "c1"), _call("add", {"a": 3, "b": 4}, "c2")] + ) + assert _shape(replies)[0] == ("c1", "add", "success", "3") + c2 = replies[1] + assert c2.status == "error" and "limit" in c2.content.lower(), c2.content + assert RAN == [("add", 1, 2)], "the refused call never ran" + + +@pytest.mark.asyncio +async def test_step_discipline_runs_only_the_first_call_and_defers_the_rest(): + replies, _, _ = await _run_batch( + _tools(add), + [_call("add", {"a": 1, "b": 2}, "c1"), _call("add", {"a": 3, "b": 4}, "c2")], + one_per_step=True, + ) + assert _shape(replies) == [("c1", "add", "success", "3"), ("c2", "add", "error", DEFERRED_CALL_MESSAGE)] + assert RAN == [("add", 1, 2)] + + +@pytest.mark.asyncio +async def test_step_discipline_defers_by_position_even_when_the_first_call_fails(): + replies, _, _ = await _run_batch( + _tools(add, boom), [_call("boom", {}, "c1"), _call("add", {"a": 1, "b": 2}, "c2")], one_per_step=True + ) + assert _shape(replies) == [ + ("c1", "boom", "error", "Tool 'boom' failed: bad x"), + ("c2", "add", "error", DEFERRED_CALL_MESSAGE), + ] + + +@pytest.mark.asyncio +async def test_step_discipline_off_runs_every_call(): + replies, _, _ = await _run_batch( + _tools(add), [_call("add", {"a": 1, "b": 2}, "c1"), _call("add", {"a": 3, "b": 4}, "c2")] + ) + assert [m.content for m in replies] == ["3", "7"] + + +@pytest.mark.asyncio +async def test_one_batch_shares_one_output_limit(): + async def text(n: int) -> str: + return "x" * n + + replies, _, _ = await _run_batch( + _tools(text), [_call("text", {"n": 8}, "c1"), _call("text", {"n": 8}, "c2")], output_limit=10 + ) + assert replies[0].content == "x" * 8 + assert replies[1].content == "xx" + TRUNCATION_MARKER.format(limit=2) + + +def test_calls_the_block_never_reached_are_all_answered(): + _, plan = plan_tool_calls( + [_call("add", {"a": 1, "b": 2}, "c1"), _call("nope", {}, "c2")], [], _tools(add), one_per_step=False + ) + replies = replies_without_execution(plan, "Not executed: boom.") + assert _shape(replies) == [ + ("c1", "add", "error", "Not executed: boom."), + ("c2", "nope", "error", "Unknown tool 'nope'. Choose one of the provided tools: add."), + ], "pre-decided replies are delivered as they are" + + +# ── what only the block form has to get right ─────────────────────────────── + + +@pytest.mark.asyncio +async def test_a_tool_name_that_is_not_an_identifier_goes_through_the_lookup(): + tools = _tools(**{"my-tool": dashed}) + replies, kept, block = await _run_batch(tools, [_call("my-tool", {"q": "hi"}, "c1")]) + assert f'{FC_TOOLS_KEY}["my-tool"]' in block + assert _shape(replies) == [("c1", "my-tool", "success", "dashed:hi")] + assert kept == {"tool_result_c1": "dashed:hi"} + + +@pytest.mark.asyncio +async def test_python_keywords_as_tool_or_argument_names_work(): + tools = _tools(**{"class": echo_args}) + args = {"from": "2026-01-01", "user-id": 5, "class": True} + replies, _, block = await _run_batch(tools, [_call("class", args, "c1")]) + assert f'{FC_TOOLS_KEY}["class"]' in block + assert _shape(replies) == [("c1", "class", "success", json.dumps(args, ensure_ascii=False))] + + +@pytest.mark.asyncio +async def test_argument_values_survive_the_round_trip(): + args = { + "s": 'it\'s "quoted" \\ back\nslash\ttab 
 ünïcödé 🙂 {braces} ```', + "n": -3, + "f": 2.5, + "b": False, + "z": None, + "l": [1, "two", {"three": [3.0, None]}], + "d": {"k": {"kk": "v"}, "": "empty key"}, + } + replies, _, _ = await _run_batch(_tools(echo_args), [_call("echo_args", args, "c1")]) + assert json.loads(replies[0].content) == args + + +@pytest.mark.asyncio +async def test_non_json_numbers_are_refused_before_any_code_is_written(): + replies, _, block = await _run_batch(_tools(echo_args), [_call("echo_args", {"n": float("nan")}, "c1")]) + assert block is None and replies[0].status == "error" + assert replies[0].content.startswith("Tool arguments could not be encoded (non-JSON number)") + + +@pytest.mark.asyncio +async def test_result_variable_names_are_derived_from_the_call_id(): + _, plan = plan_tool_calls( + [_call("add", {"a": 1, "b": 2}, "toolu_01-AB"), _call("add", {"a": 1, "b": 2}, "toolu_01-AB")], + [], + _tools(add), + one_per_step=False, + ) + assert [e["var"] for e in plan] == ["tool_result_toolu_01_AB", "tool_result_toolu_01_AB_2"] + + +@pytest.mark.asyncio +async def test_a_non_json_result_is_kept_as_the_reply_text(): + replies, kept, _ = await _run_batch(_tools(as_set), [_call("as_set", {}, "c1")]) + assert replies[0].content == _stringify({1}) == '"{1}"', "json.dumps(default=str), as a direct call gave" + assert kept == {"tool_result_c1": '"{1}"'}, "kept as the text the reply carries, never dropped" + + +@pytest.mark.asyncio +async def test_a_find_tools_listing_reaches_the_reply(): + listing = "# Found 1 Matching Tool(s)\n**Query:** q\n## 1. `x`" + + async def find_tools(query: str, app_name: str) -> str: + return listing + + replies, _, _ = await _run_batch( + _tools(find_tools), [_call("find_tools", {"query": "q", "app_name": "a"}, "c1")] + ) + assert replies[0].content == listing + + +# ── equivalence with the direct-call contract ─────────────────────────────── + + +async def _reference(tools: Dict[str, Any], calls: List[dict]) -> List[Tuple[str, str, str, str]]: + """The reply each call got from a direct ``fn(**kwargs)``, as ``tool_exec`` produced it.""" + out = [] + for index, call in enumerate(calls): + call_id, name, raw = _tool_call_parts(call) + call_id = call_id or f"call_{index}" + if not name: + out.append((call_id, "unknown", "error", "Tool call carried no tool name.")) + continue + fn = tools.get(name) + if fn is None: + known = ", ".join(sorted(k for k in tools if not k.startswith("_"))) + out.append( + (call_id, name, "error", f"Unknown tool '{name}'. Choose one of the provided tools: {known}.") + ) + continue + kwargs, err = _coerce_args(raw) + if err: + out.append((call_id, name, "error", err)) + continue + try: + result = await fn(**kwargs) + except TypeError as exc: + out.append( + ( + call_id, + name, + "error", + f"Tool '{name}' rejected these arguments: {exc}. Check the parameter names and types.", + ) + ) + except Exception as exc: + out.append((call_id, name, "error", f"Tool '{name}' failed: {exc}")) + else: + out.append((call_id, name, "success", _stringify(result))) + return out + + +_ALPHABET = string.ascii_letters + string.digits + " \"'\\\n\t{}[]():,.-_ü🙂" + + +def _rand_str(rng: random.Random) -> str: + return "".join(rng.choice(_ALPHABET) for _ in range(rng.randint(0, 12))) + + +def _rand_json(rng: random.Random, depth: int = 0) -> Any: + kind = rng.randint(0, 8 if depth < 2 else 5) + if kind == 0: + return None + if kind == 1: + return rng.choice([True, False]) + if kind == 2: + return rng.randint(-(10**6), 10**6) + if kind == 3: + return rng.choice([0.0, -1.5, 3.25, 1e21, 1e-7]) + if kind in (4, 5): + return _rand_str(rng) + if kind == 6: + return [_rand_json(rng, depth + 1) for _ in range(rng.randint(0, 3))] + return {_rand_key(rng): _rand_json(rng, depth + 1) for _ in range(rng.randint(0, 3))} + + +_KEYS = ["a", "b", "x", "from", "class", "user-id", "q", "", "with space", "λ", "n"] + + +def _rand_key(rng: random.Random) -> str: + return rng.choice(_KEYS) + + +async def _fail(**kwargs: Any) -> None: + raise RuntimeError(json.dumps(kwargs, ensure_ascii=False, sort_keys=True)) + + +async def _ab(a: int, b: int = 0) -> Dict[str, Any]: + return {"sum": a + b if isinstance(a, (int, float)) and isinstance(b, (int, float)) else [a, b]} + + +@pytest.mark.asyncio +async def test_random_calls_get_the_same_replies_as_a_direct_call(): + rng = random.Random(777) + tools = _tools(echo_args, _fail, _ab, **{"odd-name": echo_args, "class": echo_args}) + names = list(tools) + ["missing", ""] + for case in range(400): + calls = [] + for i in range(rng.randint(1, 4)): + name = rng.choice(names) + args: Any = {_rand_key(rng): _rand_json(rng) for _ in range(rng.randint(0, 3))} + if rng.random() < 0.2: + args = json.dumps(args, ensure_ascii=False) # the raw provider form + call_id = rng.choice([f"call_{case}_{i}", f"toolu_{rng.randint(0, 99)}", ""]) + calls.append(_call(name, args, call_id)) + expected = await _reference(tools, calls) + replies, _, _ = await _run_batch(tools, calls) + assert _shape(replies) == expected, f"case {case}: {calls}" diff --git a/tests/unit/test_fc_sdk_tool_approval.py b/tests/unit/test_fc_sdk_tool_approval.py new file mode 100644 index 000000000..ddd5ff6b8 --- /dev/null +++ b/tests/unit/test_fc_sdk_tool_approval.py @@ -0,0 +1,144 @@ +"""Function-calling tool approval through the real SDK path. + +A policy stored with ``agent.policies.add_tool_approval`` interrupts a native tool +call exactly as it interrupts CodeAct code: the call is translated into a block, +the approval check matches the tool name in it, the run pauses before anything +runs, and resuming with the approval executes the call and answers its id. A +denial runs nothing. +""" + +from __future__ import annotations + +from datetime import datetime + +import pytest +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.tools import StructuredTool + +from cuga.backend.cuga_graph.nodes.human_in_the_loop.followup_model import ActionResponse, ActionType + +pytestmark = pytest.mark.unit + + +class _Bound: + """What ``bind_tools`` returns: a runnable wrapping the model.""" + + def __init__(self, model): + self._model = model + + async def ainvoke(self, messages, config=None, **kwargs): + return await self._model.ainvoke(messages, config=config, **kwargs) + + +class _ScriptedModel: + def __init__(self, responses): + self._responses = list(responses) + self.invocations = 0 + self.seen: list = [] + + def bind_tools(self, tools, **kwargs): + return _Bound(self) # a real bind returns a new runnable, never the model itself + + async def ainvoke(self, messages, config=None, **kwargs): + self.invocations += 1 + self.seen.append(list(messages)) + if not self._responses: + raise AssertionError("model asked for more responses than scripted") + return self._responses.pop(0) + + +RAN: list = [] +_CALL = {"name": "delete_record", "args": {"record_id": 1}, "id": "c1", "type": "tool_call"} + + +def _delete_record(): + async def delete_record(record_id: int) -> str: + """Delete a record.""" + RAN.append(record_id) + return "deleted" + + return StructuredTool.from_function( + coroutine=delete_record, name="delete_record", description="Delete a record." + ) + + +def _response(confirmed: bool) -> ActionResponse: + return ActionResponse( + action_id="tool_approval", + response_type=ActionType.CONFIRMATION, + confirmed=confirmed, + timestamp=datetime.now().isoformat(), + ) + + +async def _guarded_fc_agent(model): + from cuga.sdk import CugaAgent + + agent = CugaAgent( + tools=[_delete_record()], model=model, execution_mode="function_calling", reset_policy_storage=True + ) + await agent.policies.add_tool_approval( + name="Approve deletions", + required_tools=["delete_record"], + approval_message="Deleting needs approval.", + ) + return agent + + +@pytest.fixture(autouse=True) +def _settings(monkeypatch): + from cuga.config import settings + + monkeypatch.setattr(settings.policy, "enabled", True, raising=False) + monkeypatch.setattr(settings.evolve, "enabled", False, raising=False) + RAN.clear() + yield + RAN.clear() + + +@pytest.mark.asyncio +async def test_fc_tool_approval_pauses_then_the_approved_call_runs_and_answers_its_id(): + model = _ScriptedModel( + [AIMessage(content="", tool_calls=[_CALL]), AIMessage(content="Deleted record 1.")] + ) + agent = await _guarded_fc_agent(model) + + paused = await agent.invoke("delete record 1", thread_id="fc-approval") + + assert model.invocations == 1 and RAN == [], "paused before the tool ran" + assert "Deleting needs approval." in paused.answer, paused.answer + + resumed = await agent.invoke(None, thread_id="fc-approval", action_response=_response(True)) + + assert RAN == [1] and resumed.answer == "Deleted record 1." + second = model.seen[1] # the model turn after the approved call ran + ai = [m for m in second if isinstance(m, AIMessage) and m.tool_calls] + tool = [m for m in second if isinstance(m, ToolMessage)] + assert ai and ai[0].tool_calls[0]["id"] == "c1", "the paused assistant turn kept its tool_calls" + assert tool and (tool[0].tool_call_id, tool[0].content) == ("c1", "deleted") + + +@pytest.mark.asyncio +async def test_fc_tool_approval_denied_runs_nothing(): + model = _ScriptedModel([AIMessage(content="", tool_calls=[_CALL])]) + agent = await _guarded_fc_agent(model) + + await agent.invoke("delete record 1", thread_id="fc-denied") + denied = await agent.invoke(None, thread_id="fc-denied", action_response=_response(False)) + + assert RAN == [] and model.invocations == 1 + assert "cancelled" in denied.answer.lower(), denied.answer + + +@pytest.mark.asyncio +async def test_codeact_on_the_same_agent_still_asks(): + """Both modes meet the same approval check, on the block each one sends to the sandbox.""" + codeact_model = _ScriptedModel( + [AIMessage(content="```python\nprint(await delete_record(record_id=1))\n```")] + ) + agent = await _guarded_fc_agent(codeact_model) + + result = await agent.invoke("delete record 1", thread_id="codeact-approval", execution_mode="codeact") + + assert codeact_model.invocations == 1 and RAN == [] + assert "Deleting needs approval." in result.answer, result.answer From dad9b386404fe4f52519513347b08493b48da513 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:31:45 +0300 Subject: [PATCH 15/18] docs(cuga-lite): function calling through the sandbox README section updated; design note with the decision, the shape, the deliberate differences and the follow-ups. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 9 +++--- docs/design/fc-through-sandbox.md | 52 +++++++++++++++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) create mode 100644 docs/design/fc-through-sandbox.md diff --git a/README.md b/README.md index ac23faee3..30675cfc1 100644 --- a/README.md +++ b/README.md @@ -1178,7 +1178,7 @@ See `docs/design/pluggable-shortlister.md` for the full design. ## What it is -By default CUGA Lite works in **CodeAct**: the model writes a Python block and the sandbox runs it. `cuga_lite_execution_mode = "function_calling"` switches to native tool calls: the model emits `tool_calls`, a `tool_exec` node runs them under the same budgets, tracker and timeout as the sandbox and replies with `ToolMessage`s, and the model reads the results before it calls again or answers. Tools are advertised through `bind_tools`; a short dedicated prompt replaces the CodeAct prompt. +By default CUGA Lite works in **CodeAct**: the model writes a Python block and the sandbox runs it. `cuga_lite_execution_mode = "function_calling"` switches to native tool calls: the model emits `tool_calls`, CUGA translates them into a short internal block that the same sandbox runs (same approval check, VERIFY gate, budgets, tracker, timeout and variables), and replies with one `ToolMessage` per call id, so the model reads the results before it calls again or answers. The block is never shown to the model or persisted; its transcript stays native. Tools are advertised through `bind_tools`; a short dedicated prompt replaces the CodeAct prompt. **Step discipline** (`cuga_lite_step_discipline = "one_tool_per_step"`) works in both modes: only the first tool call of a step is attempted — success or error, that is the result the model reads before it decides the next call. CodeAct keeps the block's variables (including that first result); function-calling answers the extra calls with a "deferred" reply. @@ -1200,12 +1200,13 @@ result = await agent.invoke("...", execution_mode="codeact") # per-call overri ## Gotchas -- **Tool-approval policies**: function-calling mode has no approval interrupt yet. If an enabled tool-approval policy exists — or policy storage cannot be checked — the run stops with a clear error before any tool runs (fail closed); use `codeact` for that agent or disable the policy. +- **Tool-approval policies** apply exactly as in CodeAct: the approval check runs on the translated block, so a policy on a tool name pauses the run before the call and the approved call resumes into the sandbox. - Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call. Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. - The bundled CodeAct few-shot demos are not sent in function-calling mode; pass your own through `configurable["mcp_few_shot_examples"]`. -- Tools run in-process and sequentially. Each call is bounded by `sandbox_execution_timeout` (the timeout is per call, so a turn with N calls can take up to N × that). The pre-execute VERIFY gate, reflection and the E2B / OpenSandbox executors are CodeAct-only; step discipline caps local-executor blocks only. +- Tools run in-process and sequentially, on the local executor even when `e2b_sandbox` is on. One turn's calls are one block, bounded by `sandbox_execution_timeout` as a whole; results computed before a timeout are kept. The pre-execute VERIFY gate applies in both modes; reflection is CodeAct-only for now. +- Each result is kept as a variable (`tool_result_`): a large result is truncated in the reply to `execution_output_max_length` but survives whole for later turns, and a later CodeAct turn on the same thread can use it by name. - Function-calling needs tools bound natively: if `bind_tools` fails or the model does not support it, or the agent has no executable tools, the run stops with a clear error (fail closed) instead of silently answering without tools. A static prompt is ignored in this mode. -- A thread can switch modes between turns: results from function-calling turns are shown to a later CodeAct turn as text. +- A thread can switch modes between turns: results from function-calling turns are shown to a later CodeAct turn as text and as variables. diff --git a/docs/design/fc-through-sandbox.md b/docs/design/fc-through-sandbox.md new file mode 100644 index 000000000..d130ef8ae --- /dev/null +++ b/docs/design/fc-through-sandbox.md @@ -0,0 +1,52 @@ +# Function calling through the sandbox + +Status: implemented on the `feature/native-function-calling-mode` branch (PR #777), 2026-10-09. + +## Decision + +`cuga_lite_execution_mode = "function_calling"` keeps native function calling on the +model side (tools via `bind_tools`, `tool_calls` in, one `ToolMessage` per call id out) +and executes the calls through the existing CodeAct sandbox instead of a separate +`tool_exec` runner. A native turn is translated into a short internal Python block, +the sandbox runs it, and the block's variables are read back into the replies. + +Why: every CUGA guarantee already lives at the sandbox step — tool approval, the +pre-execute VERIFY gate, variables, output trimming, the block timeout, budgets, the +tracker, recovery. A second executor had to re-earn each one; after three review +rounds of #777 five gaps remained (approval refused to start, no recovery, no +variables, no observed shapes, no todos). One pipeline closes them at once. + +## Shape + +- `adapter/fc_actions.py` — `plan_tool_calls` (calls → block + plan) and + `replies_from_execution` (variables → `ToolMessage`s). Rules that keep the tool side + identical to a direct call: identifier names are called as themselves (approval and + VERIFY see the real name), other names go through an injected `_fc_tools` lookup, + arguments are `**{...}` of JSON values only, one `try` per call, results kept through + `_fc_keep` so nothing is dropped, non-executable calls answered at planning time. +- `graph_adapter.execute_call_model_fc` — routes to the sandbox with the block as + `script` and the plan under `cuga_lite_metadata["fc_pending"]`; runs the CodeAct + approval check on the block first. +- `sandbox_node` — when a plan is pending: local executor, listing variables kept, + replies instead of the "Execution output:" message at every exit (VERIFY revise, + normal, failure), plan cleared on every exit, reflection skipped. +- Feature off: no branch above is reachable (gated on a pending plan a CodeAct turn + never writes); every new parameter defaults to the old behavior. Golden: + `tests/snapshots/codeact_feature_off.json`. + +## Deliberate differences from the direct runner + +- Timeout is per block (whole batch), as for CodeAct; finished results are kept. +- Results persist as variables (`tool_result_`), visible to later turns and + in `InvokeResult.variables`. +- Approval follows CodeAct's check (text match on the block), including its + fail-open on infrastructure errors; the former refuse-to-start guard is gone. + +## Follow-ups + +1. Bind `find_tools` in function-calling mode (it is stripped unless + `cuga_lite_bind_tools_include_find_tools` is set, while the prompt advertises it). +2. A `run_python` tool backed by the same sandbox for data-heavy steps. +3. Per-thread dynamic binding: remember the shortlist and discovered tools, query on + the current task rather than the first message. +4. Function-calling wording for VERIFY, reflection and the empty-reply nudge. From 5a285f1850a7c4a38a22c71a53c0c031890b677d Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:40:06 +0300 Subject: [PATCH 16/18] feat(cuga-lite): aliases, find_tools and provider-error recovery in FC mode Three things a Vakra-sized domain needs from function calling: - A provider-safe alias (bind_tools/tool_names.py, #882 merged into this branch) is resolved to the real tool before the block is written, so the block, the approval scan and VERIFY see the real name while the reply keeps the alias the provider knows. 6 of 10 Vakra test domains carry a tool name over 64 characters; without this the whole request is rejected. - find_tools is bound in function-calling mode whenever prepare enabled it, since the FC prompt already tells the model to call it; the cap merge stripped it. An explicit per-invoke choice is respected. Its listing is replied, never kept as a variable. - The FC model call gets the recovery the CodeAct path has, in native terms: a Groq failed_generation comes back as a real tool call (or an invalid one when the arguments do not parse), a python attempt gets the mode-violation correction, a retryable rejection is retried once and a second failure ends the turn with a clear error. Anything else propagates, as before. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- README.md | 4 +- docs/design/fc-through-sandbox.md | 20 ++- .../nodes/cuga_lite/adapter/fc_actions.py | 62 ++++++--- .../nodes/cuga_lite/adapter/graph_adapter.py | 102 ++++++++++++++- .../nodes/cuga_lite/tests/test_fc_actions.py | 91 +++++++++++-- .../tests/test_fc_call_model_hook.py | 121 ++++++++++++++++++ .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 58 +++++++++ 7 files changed, 418 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 30675cfc1..035710ba4 100644 --- a/README.md +++ b/README.md @@ -1201,7 +1201,9 @@ result = await agent.invoke("...", execution_mode="codeact") # per-call overri ## Gotchas - **Tool-approval policies** apply exactly as in CodeAct: the approval check runs on the translated block, so a policy on a tool name pauses the run before the call and the approved call resumes into the sandbox. -- Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call. Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. +- Bind mode `none` (the default) is upgraded in function-calling mode to advertise exactly the tools the sandbox could call, plus `find_tools` whenever the prompt advertises it (set `configurable["cuga_lite_bind_tools_include_find_tools"] = False` to opt out). Past `cuga_lite_bind_tools_max_count` (128) the bind-cap shortlister runs every turn; `[shortlister.bind_cap] strategy = "embedding"` avoids the extra LLM call. +- Tool names providers reject (over 64 characters, or with characters outside `[A-Za-z0-9_-]`) are bound under a provider-safe alias; the model calls the alias, the real tool runs, and the reply keeps the alias. +- A provider that rejects a malformed native call but returns the attempt (Groq `tool_use_failed`) gets it back as a real call; a retryable rejection is retried once, then the turn ends with a clear error. - The bundled CodeAct few-shot demos are not sent in function-calling mode; pass your own through `configurable["mcp_few_shot_examples"]`. - Tools run in-process and sequentially, on the local executor even when `e2b_sandbox` is on. One turn's calls are one block, bounded by `sandbox_execution_timeout` as a whole; results computed before a timeout are kept. The pre-execute VERIFY gate applies in both modes; reflection is CodeAct-only for now. - Each result is kept as a variable (`tool_result_`): a large result is truncated in the reply to `execution_output_max_length` but survives whole for later turns, and a later CodeAct turn on the same thread can use it by name. diff --git a/docs/design/fc-through-sandbox.md b/docs/design/fc-through-sandbox.md index d130ef8ae..98b3b35d1 100644 --- a/docs/design/fc-through-sandbox.md +++ b/docs/design/fc-through-sandbox.md @@ -42,11 +42,21 @@ variables, no observed shapes, no todos). One pipeline closes them at once. - Approval follows CodeAct's check (text match on the block), including its fail-open on infrastructure errors; the former refuse-to-start guard is gone. +## Also in this round (Vakra readiness) + +- Provider-safe aliases (#882, merged into this branch): the translator resolves an + alias to the real name before planning, so the block, approval and VERIFY see the + real tool while the reply keeps the alias the provider knows. 6 of 10 Vakra test + domains have a tool name over 64 characters. +- `find_tools` is bound in function-calling mode whenever prepare enabled it, since the + FC prompt advertises it; its listing is replied, never kept as a variable. +- The FC model call has the recovery CodeAct has, in native terms: a Groq + `failed_generation` becomes a real call, a retryable rejection is retried once, a + second failure ends the turn with a clear error, anything else propagates. + ## Follow-ups -1. Bind `find_tools` in function-calling mode (it is stripped unless - `cuga_lite_bind_tools_include_find_tools` is set, while the prompt advertises it). -2. A `run_python` tool backed by the same sandbox for data-heavy steps. -3. Per-thread dynamic binding: remember the shortlist and discovered tools, query on +1. A `run_python` tool backed by the same sandbox for data-heavy steps. +2. Per-thread dynamic binding: remember the shortlist and discovered tools, query on the current task rather than the first message. -4. Function-calling wording for VERIFY, reflection and the empty-reply nudge. +3. Function-calling wording for VERIFY, reflection and the empty-reply nudge. diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py index 60b342fe3..58ee8d62c 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py @@ -13,9 +13,12 @@ What keeps the tool side identical to a direct call: -- a tool whose name is a Python identifier is called by that name, so the - approval text scan and VERIFY see the real tool name; any other name goes - through the injected ``_fc_tools`` lookup, budget-counted like the rest; +- a provider-safe alias (``bind_tools/tool_names.py``) is resolved to the real + tool name before anything else, so the block, the approval text scan and + VERIFY see the real name while the reply keeps the alias the provider knows; +- a tool whose name is a Python identifier is called by that name; any other + name goes through the injected ``_fc_tools`` lookup, budget-counted like the + rest; - arguments are always passed as ``**{...}`` so keys like ``from`` work, and only JSON values are encoded — providers deliver JSON, so nothing is lost; - each call sits in its own ``try``: one failing call never aborts its siblings, @@ -33,7 +36,7 @@ import keyword import math import re -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple from langchain_core.messages import ToolMessage @@ -193,14 +196,17 @@ def plan_tool_calls( tools_context: Dict[str, Any], *, one_per_step: bool, + resolve_name: Optional[Callable[[str], str]] = None, ) -> Tuple[Optional[str], List[Dict[str, Any]]]: """Translate one turn's calls into ``(block, plan)``. ``plan`` has one JSON-serialisable entry per issued id, in order: - ``{"id", "name", "var", "lookup", "reply"}``. ``var`` names the block - variable that will hold the call's result; ``reply`` is set instead when - the call was answered here and emits no code. ``block`` is ``None`` when - nothing is left to execute. + ``{"id", "name", "target", "var", "lookup", "reply"}``. ``name`` is the + name as the model issued it (an alias stays an alias, so the reply matches + the provider's tool list); ``target`` is the real tool the block calls, as + ``resolve_name`` maps it. ``var`` names the block variable that will hold + the call's result; ``reply`` is set instead when the call was answered here + and emits no code. ``block`` is ``None`` when nothing is left to execute. """ plan: List[Dict[str, Any]] = [] lines: List[str] = [] @@ -209,7 +215,14 @@ def plan_tool_calls( for index, call in enumerate(calls): call_id, name, raw_args = _tool_call_parts(call) call_id = call_id or f"call_{index}" - entry: Dict[str, Any] = {"id": call_id, "name": name, "var": None, "lookup": False, "reply": None} + entry: Dict[str, Any] = { + "id": call_id, + "name": name, + "target": name, + "var": None, + "lookup": False, + "reply": None, + } plan.append(entry) if not name: @@ -220,7 +233,15 @@ def plan_tool_calls( # the model must read before it decides the next one. entry["reply"] = _pre_decided(DEFERRED_CALL_MESSAGE) continue - fn = tools_context.get(name) + target = name + if resolve_name is not None: + try: + target = resolve_name(name) or name + except Exception as exc: # an ambiguous alias: answer it, never crash the turn + entry["reply"] = _pre_decided(f"Tool name {name!r} could not be resolved: {exc}") + continue + entry["target"] = target + fn = tools_context.get(target) if fn is None or not callable(fn): known = ", ".join(sorted(k for k in tools_context if not k.startswith("_"))) entry["reply"] = _pre_decided( @@ -241,14 +262,14 @@ def plan_tool_calls( var = _result_var(call_id, taken) entry["var"] = var - if name.isidentifier() and not keyword.iskeyword(name): - target = name + if target.isidentifier() and not keyword.iskeyword(target): + callee = target else: entry["lookup"] = True - target = f"{FC_TOOLS_KEY}[{json.dumps(name, ensure_ascii=False)}]" + callee = f"{FC_TOOLS_KEY}[{json.dumps(target, ensure_ascii=False)}]" lines += [ "try:", - f" {var} = {FC_KEEP_KEY}(await {target}(**{args_src}))", + f" {var} = {FC_KEEP_KEY}(await {callee}(**{args_src}))", "except Exception as _fc_e:", f" {var} = {_ERROR_MARKER_SRC}", ] @@ -260,6 +281,7 @@ def plan_tool_calls( { "id": call_id or f"invalid_{index}", "name": name, + "target": name, "var": None, "lookup": False, "reply": _pre_decided( @@ -300,9 +322,9 @@ def fc_context_overlay(tools_context: Dict[str, Any], plan: List[Dict[str, Any]] overlay: Dict[str, Any] = {FC_KEEP_KEY: fc_keep, FC_BUDGET_EXC_KEY: ToolCallBudgetExceeded} lookup = { - e["name"]: counted_tool_call(tools_context[e["name"]]) + e["target"]: counted_tool_call(tools_context[e["target"]]) for e in plan - if e.get("lookup") and e.get("name") in tools_context + if e.get("lookup") and e.get("target") in tools_context } if lookup: overlay[FC_TOOLS_KEY] = lookup @@ -339,8 +361,11 @@ def replies_from_execution( """One ``ToolMessage`` per issued id, read back from the block's variables. Returns ``(replies, variables_to_keep, variables_to_drop)``: error markers - are formatted into their reply and dropped from the variables. + are formatted into their reply and dropped from the variables, and so is a + find_tools listing — the reply carries it, CodeAct never keeps it either. """ + from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import is_find_tools_listing_markdown + budget = _BatchOutputBudget(output_limit) keep = dict(new_vars) drop: List[str] = [] @@ -360,6 +385,9 @@ def replies_from_execution( drop.append(var) replies.append(_error_message(_format_error(name, value), call_id=call_id, name=name)) else: + if is_find_tools_listing_markdown(value): + keep.pop(var, None) + drop.append(var) replies.append( ToolMessage(content=budget.take(_stringify(value)), tool_call_id=call_id, name=name) ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py index 0f1feca7b..761e386cb 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/graph_adapter.py @@ -62,7 +62,12 @@ normalize_assistant_text, ) from cuga.backend.cuga_graph.utils.token_counter import clamp_watsonx_completion_for_messages -from cuga.backend.llm.errors import extract_code_from_tool_use_failed +from cuga.backend.llm.errors import ( + extract_code_from_tool_use_failed, + is_ollama_tool_call_parse_error, + is_tool_choice_none_tool_use_failed, + parse_tool_use_failed_generation, +) from cuga.config import settings @@ -85,6 +90,11 @@ ) _CODEACT_VARIABLES_NOTE = "You can use these variables directly by their names." +FC_PROVIDER_REJECTED = ( + "The model provider rejected the request twice in a row ({error}). No tool ran. " + 'If this persists, use cuga_lite_execution_mode = "codeact" for this model.' +) + FC_STEP_LIMIT_CALL_REPLY = "Not executed: the step limit was reached before this call could run." FC_BUDGET_CALL_REPLY = ( "Not executed: the tool budget for this turn is spent. Answer from the data already retrieved." @@ -378,6 +388,13 @@ async def resolve_bind_tools( "cuga_lite_bind_tools_mode": "tools", "cuga_lite_bind_tools_tool_names": names, } + # The FC prompt tells the model to call find_tools whenever prepare + # enabled it, so it has to be bound too; the cap merge strips it + # otherwise. An explicit per-invoke choice is respected. + if (self._tools_context_ref or {}).get( + "_lc_bind_tools_find_tools" + ) is not None and "cuga_lite_bind_tools_include_find_tools" not in configurable: + configurable["cuga_lite_bind_tools_include_find_tools"] = True logger.info( "[fc] bind_tools mode was 'none'; advertising the {} executable tool(s)", len(names) ) @@ -494,6 +511,69 @@ def _resolved_bind_mode(self, configurable: dict) -> str: return str(candidate).strip().lower() return _bind_tools_mode_from_settings() + async def _invoke_fc_model( + self, bound: Any, msgs: list, invoke_config: dict + ) -> Tuple[Any, Optional[str]]: + """``bound.ainvoke`` with the recovery the CodeAct path has, in native terms. + + - a provider that rejects a malformed native call but returns the attempt + (Groq ``tool_use_failed`` + ``failed_generation``) gets it back as a real + tool call (or an invalid one when the arguments do not parse), so the + loop answers it instead of dying; a ``python`` attempt is handed to the + mode-violation correction; + - a retryable rejection (``tool_choice`` none, Ollama's tool-call parser) is + retried once; a second failure ends the turn with a clear error, returned + as ``(None, text)``; + - anything else propagates, exactly as on the CodeAct path. + """ + clamp_watsonx_completion_for_messages(bound, msgs) + try: + return await bound.ainvoke(msgs, config=invoke_config), None + except Exception as exc: + failed = parse_tool_use_failed_generation(exc) + name = str(failed.get("name") or "") if isinstance(failed, dict) else "" + if name == "python": + logger.warning( + "[fc] provider rejected a code attempt sent as a tool call; one corrective turn" + ) + return AIMessage(content=f"```python\n{failed.get('arguments') or ''}\n```"), None + if name: + args = failed.get("arguments") + if isinstance(args, str): + try: + args = json.loads(args) + except json.JSONDecodeError: + pass + logger.warning( + "[fc] provider rejected a malformed tool call; recovering it from failed_generation" + ) + if isinstance(args, dict) or args is None: + return AIMessage( + content="", + tool_calls=[ + {"name": name, "args": args or {}, "id": "call_recovered_1", "type": "tool_call"} + ], + ), None + return AIMessage( + content="", + invalid_tool_calls=[ + { + "name": name, + "args": str(args), + "id": "call_recovered_1", + "error": "arguments are not a JSON object", + "type": "invalid_tool_call", + } + ], + ), None + if is_tool_choice_none_tool_use_failed(exc) or is_ollama_tool_call_parse_error(exc): + logger.warning("[fc] retrying once after a retryable provider tool-call error: {}", exc) + try: + return await bound.ainvoke(msgs, config=invoke_config), None + except Exception as again: + return None, str(again) + raise + async def execute_call_model_fc( self, *, @@ -562,8 +642,15 @@ async def execute_call_model_fc( # Outbound only as well. msgs.append(HumanMessage(content=TOOL_BUDGET_EXHAUSTED_INSTRUCTION)) - clamp_watsonx_completion_for_messages(bound, msgs) - response = await bound.ainvoke(msgs, config=invoke_config) + response, provider_error = await self._invoke_fc_model(bound, msgs, invoke_config) + if response is None: + logger.error("[fc] provider rejected the request twice; ending the turn: {}", provider_error) + return create_error_command( + self, + history, + AIMessage(content=FC_PROVIDER_REJECTED.format(error=provider_error)), + state.step_count, + ) tool_calls = list(getattr(response, "tool_calls", None) or []) invalid_tool_calls = list(getattr(response, "invalid_tool_calls", None) or []) @@ -626,8 +713,15 @@ async def execute_call_model_fc( resolve_step_discipline(cfg, self._runtime_model_name(cfg)) == STEP_DISCIPLINE_ONE_TOOL_PER_STEP ) + tools = self._tools_context block, plan = plan_tool_calls( - tool_calls, invalid_tool_calls, self._tools_context, one_per_step=one_per_step + tool_calls, + invalid_tool_calls, + tools, + one_per_step=one_per_step, + # A provider-safe alias is what the model was shown; the block, and so + # the approval scan and VERIFY, call the real tool. The reply keeps the alias. + resolve_name=lambda issued: resolve_tool_names(issued, tools), ) if block is None: # Every call was answered at planning time (no name, unknown tool, bad diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py index 58af97db2..b43c7bbec 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py @@ -379,19 +379,6 @@ async def test_a_non_json_result_is_kept_as_the_reply_text(): assert kept == {"tool_result_c1": '"{1}"'}, "kept as the text the reply carries, never dropped" -@pytest.mark.asyncio -async def test_a_find_tools_listing_reaches_the_reply(): - listing = "# Found 1 Matching Tool(s)\n**Query:** q\n## 1. `x`" - - async def find_tools(query: str, app_name: str) -> str: - return listing - - replies, _, _ = await _run_batch( - _tools(find_tools), [_call("find_tools", {"query": "q", "app_name": "a"}, "c1")] - ) - assert replies[0].content == listing - - # ── equivalence with the direct-call contract ─────────────────────────────── @@ -489,3 +476,81 @@ async def test_random_calls_get_the_same_replies_as_a_direct_call(): expected = await _reference(tools, calls) replies, _, _ = await _run_batch(tools, calls) assert _shape(replies) == expected, f"case {case}: {calls}" + + +# ── provider-safe aliases ──────────────────────────────────────────────────── + +LONG_NAME = "public_review_platform_get_reviews_with_user_details_and_metadata_v2" # 68 chars, over the 64 limit + + +@pytest.mark.asyncio +async def test_an_alias_is_called_by_its_real_name_and_answered_under_the_alias(): + """What the model was shown is the provider-safe alias; the block, and so the approval + scan and VERIFY, call the real tool; the reply keeps the alias the provider knows.""" + from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import ( + provider_safe_tool_name, + resolve_tool_names, + ) + + async def long_tool(n: int) -> int: + RAN.append(("long", n)) + return n * 2 + + tools = _tools(**{LONG_NAME: long_tool}) + alias = provider_safe_tool_name(LONG_NAME) + assert alias != LONG_NAME and len(alias) == 64 + + block, plan = plan_tool_calls( + [_call(alias, {"n": 5}, "c1")], + [], + tools, + one_per_step=False, + resolve_name=lambda issued: resolve_tool_names(issued, tools), + ) + assert f"await {LONG_NAME}(**" in block and alias not in block, block + assert plan[0]["name"] == alias and plan[0]["target"] == LONG_NAME + + context = {**tools, **fc_context_overlay(tools, plan)} + ToolCallTracker.start_tracking(enabled=False, timings_only=False) + ToolCallTracker.seed_call_budget(0, 0) + try: + output, new_vars = await CodeExecutor.eval_with_tools_async( + code=block, + _locals=context, + state=CugaLiteState(chat_messages=[]), + mode="local", + keep_listing_vars=True, + ) + finally: + ToolCallTracker.stop_tracking() + replies, kept, _ = replies_from_execution(plan, new_vars, output, output_limit=0, timeout=30) + + assert _shape(replies) == [("c1", alias, "success", "10")] + assert RAN == [("long", 5)] and kept == {"tool_result_c1": 10} + + +def test_a_name_the_resolver_cannot_place_is_answered_not_raised(): + def ambiguous(_: str) -> str: + raise RuntimeError("Tool name 'x' is ambiguous") + + block, plan = plan_tool_calls( + [_call("x_deadbeef", {}, "c1")], [], _tools(add), one_per_step=False, resolve_name=ambiguous + ) + assert block is None + assert ( + plan[0]["reply"]["content"] + == "Tool name 'x_deadbeef' could not be resolved: Tool name 'x' is ambiguous" + ) + + +@pytest.mark.asyncio +async def test_a_find_tools_listing_is_replied_but_not_kept_as_a_variable(): + listing = "# Found 1 Matching Tool(s)\n**Query:** q\n## 1. `x`" + + async def find_tools(query: str, app_name: str) -> str: + return listing + + replies, kept, _ = await _run_batch( + _tools(find_tools), [_call("find_tools", {"query": "q", "app_name": "a"}, "c1")] + ) + assert replies[0].content == listing and kept == {} diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py index 1adcd83f3..e1c9bf889 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_call_model_hook.py @@ -516,3 +516,124 @@ async def test_watsonx_reasoning_key_is_read_like_the_codeact_path(): cmd = await _turn(_adapter(), _Model(response), state, FC) assert cmd.goto == END and cmd.update["final_answer"] == "The total is 42." + + +# ── find_tools is advertised when prepare enabled it ───────────────────────── + + +@pytest.mark.asyncio +async def test_bind_advertises_find_tools_when_prepare_enabled_it(): + """The FC prompt tells the model to call find_tools whenever prepare enabled it, so + the bind upgrade includes it; the cap merge would strip it otherwise. An explicit + per-invoke choice wins.""" + captured = [] + + async def fake_resolve(active_model, **kwargs): + captured.append(kwargs["configurable"]) + return "BOUND" + + with patch.object(ga, "resolve_model_with_bind_tools", new=AsyncMock(side_effect=fake_resolve)): + with_find = _adapter( + tools_context_ref={ + "_lc_bind_tools_executable_names": ["echo", "find_tools"], + "_lc_bind_tools_find_tools": object(), + } + ) + await with_find.resolve_bind_tools(_state(), object(), dict(FC), None) + await with_find.resolve_bind_tools( + _state(), object(), {**FC, "cuga_lite_bind_tools_include_find_tools": False}, None + ) + without = _adapter(tools_context_ref={"_lc_bind_tools_executable_names": ["echo"]}) + await without.resolve_bind_tools(_state(), object(), dict(FC), None) + + enabled, explicit_off, not_enabled = captured + assert enabled["cuga_lite_bind_tools_include_find_tools"] is True + assert explicit_off["cuga_lite_bind_tools_include_find_tools"] is False, "an explicit choice is respected" + assert "cuga_lite_bind_tools_include_find_tools" not in not_enabled + + +# ── provider-side tool-call failures ───────────────────────────────────────── + + +class _Flaky: + """Scripted model: an Exception entry is raised, anything else is returned.""" + + def __init__(self, script: list): + self.script = list(script) + self.seen: List[list] = [] + + async def ainvoke(self, messages, config=None, **kwargs): + self.seen.append(list(messages)) + item = self.script.pop(0) + if isinstance(item, Exception): + raise item + return item + + +class _GroqToolUseFailed(Exception): + def __init__(self, failed_generation: str): + super().__init__("Error code: 400 - tool_use_failed") + self.body = {"error": {"code": "tool_use_failed", "failed_generation": failed_generation}} + + +_TOOL_CHOICE_NONE = Exception( + "Error code: 400 - {'error': {'message': 'Failed to call a function. Tool choice is none.', 'code': 'tool_use_failed'}}" +) + + +@pytest.mark.asyncio +async def test_a_rejected_malformed_tool_call_is_recovered_as_a_native_call(): + model = _Flaky([_GroqToolUseFailed('{"name": "add", "arguments": {"a": 1, "b": 2}}')]) + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == "sandbox" and 'await add(**{"a": 1, "b": 2})' in cmd.update["script"] + (entry,) = cmd.update["cuga_lite_metadata"][FC_PENDING_KEY] + assert entry["id"] == "call_recovered_1" + assert cmd.update["chat_messages"][-1].tool_calls[0]["id"] == "call_recovered_1" + + +@pytest.mark.asyncio +async def test_a_rejected_call_with_unparsable_arguments_is_answered_as_invalid(): + model = _Flaky([_GroqToolUseFailed('{"name": "add", "arguments": "{not json"}')]) + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == "call_model", "nothing executable: answered without the sandbox" + reply = cmd.update["chat_messages"][-1] + assert ( + isinstance(reply, ToolMessage) + and reply.tool_call_id == "call_recovered_1" + and reply.status == "error" + ) + assert reply.content.startswith("The provider could not parse this tool call") + + +@pytest.mark.asyncio +async def test_a_rejected_code_attempt_gets_the_mode_violation_correction(): + model = _Flaky([_GroqToolUseFailed('{"name": "python", "arguments": "print(await add(a=1, b=2))"}')]) + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == "call_model" + assert cmd.update["chat_messages"][-1] == HumanMessage(content=FC_MODE_VIOLATION_CORRECTION) + + +@pytest.mark.asyncio +async def test_a_retryable_provider_rejection_is_retried_once(): + model = _Flaky([_TOOL_CHOICE_NONE, AIMessage(content="It is 3.")]) + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == END and cmd.update["final_answer"] == "It is 3." and len(model.seen) == 2 + + +@pytest.mark.asyncio +async def test_two_provider_rejections_end_the_turn_with_a_clear_error(): + model = _Flaky([_TOOL_CHOICE_NONE, _TOOL_CHOICE_NONE]) + cmd = await _turn(_adapter(), model, _state(), FC) + + assert cmd.goto == END and len(model.seen) == 2 + assert cmd.update["error"].startswith("The model provider rejected the request twice in a row") + + +@pytest.mark.asyncio +async def test_other_provider_errors_still_propagate(): + with pytest.raises(RuntimeError, match="boom"): + await _turn(_adapter(), _Flaky([RuntimeError("boom")]), _state(), FC) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index 68a4b9fb5..02315c67b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -423,3 +423,61 @@ async def dashed(q: str) -> str: assert CALLS == [("my-tool", "hi")] reply = [m for m in result["chat_messages"] if isinstance(m, ToolMessage)][0] assert (reply.tool_call_id, reply.name, reply.content) == ("c1", "my-tool", "ok:hi") + + +@pytest.mark.asyncio +async def test_a_tool_name_providers_reject_is_bound_under_an_alias_and_still_runs(): + """A registry name over 64 characters is bound under a provider-safe alias; the model calls + the alias, the real tool runs, and the reply carries the alias the provider knows.""" + from cuga.backend.cuga_graph.nodes.cuga_lite.bind_tools.tool_names import provider_safe_tool_name + + long_name = "public_review_platform_get_reviews_with_user_details_and_metadata_v2" + assert len(long_name) > 64 + alias = provider_safe_tool_name(long_name) + + async def reviews(limit: int) -> str: + """Reviews.""" + CALLS.append(("reviews", limit)) + return f"{limit} reviews" + + tool = StructuredTool.from_function(coroutine=reviews, name=long_name, description="Reviews.") + + def first(messages): + assert alias in model.bound_tool_names and long_name not in model.bound_tool_names + return AIMessage(content="", tool_calls=[_tc(alias, {"limit": 3}, "c1")]) + + model = _ScriptedModel([first, AIMessage(content="Three reviews.")]) + + result = await _run(model, _config("fc-alias", cuga_lite_execution_mode="function_calling"), tool) + + assert CALLS == [("reviews", 3)] and result["final_answer"] == "Three reviews." + reply = [m for m in result["chat_messages"] if isinstance(m, ToolMessage)][0] + assert (reply.tool_call_id, reply.name, reply.content) == ("c1", alias, "3 reviews") + + +@pytest.mark.asyncio +async def test_find_tools_is_bound_in_fc_mode_when_prepare_enabled_it(): + """Above the shortlisting threshold the FC prompt advertises find_tools, so it is + bound; calling it works and its listing is not kept as a variable.""" + + def first(messages): + assert "find_tools" in model.bound_tool_names, model.bound_tool_names + return AIMessage( + content="", tool_calls=[_tc("find_tools", {"query": "echo", "app_name": "test_app"}, "c1")] + ) + + def second(messages): + reply = _last_tool_message(messages) + assert reply.tool_call_id == "c1" and reply.content.strip() + return AIMessage(content="Looked it up.") + + model = _ScriptedModel([first, second]) + config = _config("fc-find", cuga_lite_execution_mode="function_calling", shortlisting_tool_threshold=0) + + result = await _run(model, config, _echo_tool()) + + assert result["final_answer"] == "Looked it up." + assert not any( + isinstance(v.get("value"), str) and "Matching Tool(s)" in v["value"] + for v in result["variables_storage"].values() + ), "a find_tools listing is replied, never kept as a variable" From 5d102eebc7283ad06f8a4b83b32c28be03e4b394 Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Sun, 11 Oct 2026 11:57:29 +0300 Subject: [PATCH 17/18] fix(cuga-lite): record function-calling outcomes in a node-owned channel Review (Nakagawa-master): a call's failure was recognised by a key inside the tool's return value, which is untrusted data, so a tool that legitimately returned {"__cuga_fc_error__": ...} became an error reply and lost its variable. The block no longer stores a marker. Each call records its outcome through two node-owned recorders, _fc_ok and _fc_failed, that close over a dict the sandbox node creates per step; what lands there is decided by the generated block's structure, never by what a tool returns. The exception kind is classified in that closure with the real classes. Outcomes recorded before a block timeout survive it, as the timeout test shows. Tests: the reviewer's case (such a result is a success with the intact JSON and the variable kept), its inverse (a raising tool is an error with no variable), and the reserved spellings in the random equivalence pool. Also the formatting the CI Lint job flagged. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/fc_actions.py | 146 +++++++++++------- .../nodes/cuga_lite/adapter/sandbox_node.py | 9 +- .../nodes/cuga_lite/tests/test_fc_actions.py | 58 ++++++- 3 files changed, 143 insertions(+), 70 deletions(-) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py index 58ee8d62c..cc6579345 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py @@ -22,7 +22,9 @@ - arguments are always passed as ``**{...}`` so keys like ``from`` work, and only JSON values are encoded — providers deliver JSON, so nothing is lost; - each call sits in its own ``try``: one failing call never aborts its siblings, - and the exception is kept as a marker the reply builder formats; + and every call records its outcome through ``_fc_ok`` / ``_fc_failed`` into a + dict the sandbox node owns — a trusted channel, never a spelling inside the + tool's own return value, which is untrusted data; - the result is kept through ``_fc_keep`` so a non-JSON value survives as the very text the reply carries instead of being dropped from the variables; - calls that cannot run (no name, unknown tool, unparsable arguments, @@ -44,10 +46,10 @@ """Metadata key: the turn's call plan while the sandbox runs the block; cleared on every sandbox exit.""" FC_KEEP_KEY = "_fc_keep" +FC_OK_KEY = "_fc_ok" +FC_FAILED_KEY = "_fc_failed" FC_TOOLS_KEY = "_fc_tools" -FC_BUDGET_EXC_KEY = "_fc_budget_exc" RESULT_VAR_PREFIX = "tool_result_" -_ERROR_KEY = "__cuga_fc_error__" _TIMEOUT_PREFIX = "Error during execution: Execution timed out" DEFERRED_CALL_MESSAGE = ( @@ -184,12 +186,6 @@ def _pre_decided(content: str) -> Dict[str, str]: return {"content": content, "status": "error"} -_ERROR_MARKER_SRC = ( - "{" + json.dumps(_ERROR_KEY) + ': str(_fc_e), "kind": "TypeError" if isinstance(_fc_e, TypeError) ' - 'else ("budget" if isinstance(_fc_e, ' + FC_BUDGET_EXC_KEY + ') else "error")}' -) - - def plan_tool_calls( calls: List[Any], invalid: List[Any], @@ -267,11 +263,13 @@ def plan_tool_calls( else: entry["lookup"] = True callee = f"{FC_TOOLS_KEY}[{json.dumps(target, ensure_ascii=False)}]" + id_src = json.dumps(call_id, ensure_ascii=False) lines += [ "try:", f" {var} = {FC_KEEP_KEY}(await {callee}(**{args_src}))", + f" {FC_OK_KEY}({id_src}, {var})", "except Exception as _fc_e:", - f" {var} = {_ERROR_MARKER_SRC}", + f" {FC_FAILED_KEY}({id_src}, _fc_e)", ] for index, bad in enumerate(invalid): @@ -308,19 +306,39 @@ def metadata_without_plan(adapter: Any, state: Any) -> Dict[str, Any]: return {k: v for k, v in (adapter.get_metadata(state) or {}).items() if k != FC_PENDING_KEY} -def fc_context_overlay(tools_context: Dict[str, Any], plan: List[Dict[str, Any]]) -> Dict[str, Any]: - """Names the block needs beyond the tools: the keeper, the budget exception and the lookup table. +def fc_context_overlay( + tools_context: Dict[str, Any], plan: List[Dict[str, Any]] +) -> Tuple[Dict[str, Any], Dict[str, Dict[str, Any]]]: + """Names the block needs beyond the tools, and the outcome channel they write to. - All ``_``-prefixed, so the executor neither counts them as tools nor keeps - them as variables. Lookup entries are budget-counted here because the - executor only wraps top-level callables. + Returns ``(overlay, outcomes)``. ``overlay`` holds the keeper, the two outcome + recorders and the lookup table for non-identifier names, all ``_``-prefixed so + the executor neither counts them as tools nor keeps them as variables. The + recorders close over ``outcomes``, the node's own dict: the block is generated + code, so what lands there is decided by its structure, not by anything a tool + returns. Lookup entries are budget-counted here because the executor only + wraps top-level callables. """ from cuga.backend.cuga_graph.nodes.cuga_lite.tracking.tracker import ( ToolCallBudgetExceeded, counted_tool_call, ) - overlay: Dict[str, Any] = {FC_KEEP_KEY: fc_keep, FC_BUDGET_EXC_KEY: ToolCallBudgetExceeded} + outcomes: Dict[str, Dict[str, Any]] = {} + + def _ok(call_id: str, value: Any) -> None: + outcomes[call_id] = {"ok": True, "value": value} + + def _failed(call_id: str, exc: BaseException) -> None: + if isinstance(exc, ToolCallBudgetExceeded): + kind = "budget" + elif isinstance(exc, TypeError): + kind = "TypeError" + else: + kind = "error" + outcomes[call_id] = {"ok": False, "kind": kind, "error": str(exc)} + + overlay: Dict[str, Any] = {FC_KEEP_KEY: fc_keep, FC_OK_KEY: _ok, FC_FAILED_KEY: _failed} lookup = { e["target"]: counted_tool_call(tools_context[e["target"]]) for e in plan @@ -328,12 +346,12 @@ def fc_context_overlay(tools_context: Dict[str, Any], plan: List[Dict[str, Any]] } if lookup: overlay[FC_TOOLS_KEY] = lookup - return overlay + return overlay, outcomes -def _format_error(name: str, marker: Dict[str, Any]) -> str: - text = str(marker.get(_ERROR_KEY, "")) - kind = marker.get("kind") +def _format_error(name: str, outcome: Dict[str, Any]) -> str: + text = str(outcome.get("error", "")) + kind = outcome.get("kind") if kind == "TypeError": return f"Tool '{name}' rejected these arguments: {text}. Check the parameter names and types." if kind == "budget": @@ -350,61 +368,69 @@ def _unreached_reason(output: str, timeout: Any) -> str: return NOT_EXECUTED_REPLY.format(reason=first_line.replace("{", "{{").replace("}", "}}")) +def _reply_for( + entry: Dict[str, Any], + outcome: Optional[Dict[str, Any]], + value_lookup: Dict[str, Any], + unreached: str, + budget: _BatchOutputBudget, + drop: List[str], +) -> ToolMessage: + call_id, name = entry["id"], entry.get("name") or "" + pre = entry.get("reply") + if pre: + return _error_message(pre["content"], call_id=call_id, name=name) + if outcome is None: + return _error_message(unreached.format(name=name), call_id=call_id, name=name) + if not outcome.get("ok"): + return _error_message(_format_error(name, outcome), call_id=call_id, name=name) + var = entry.get("var") + value = value_lookup[var] if var is not None and var in value_lookup else outcome.get("value") + if var is not None and _is_listing(value): + drop.append(var) + return ToolMessage(content=budget.take(_stringify(value)), tool_call_id=call_id, name=name) + + +def _is_listing(value: Any) -> bool: + from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import is_find_tools_listing_markdown + + return is_find_tools_listing_markdown(value) + + def replies_from_execution( plan: List[Dict[str, Any]], new_vars: Dict[str, Any], + outcomes: Dict[str, Dict[str, Any]], output: str, *, output_limit: Any, timeout: Any, ) -> Tuple[List[ToolMessage], Dict[str, Any], List[str]]: - """One ``ToolMessage`` per issued id, read back from the block's variables. + """One ``ToolMessage`` per issued id, from the outcome channel and the block's variables. - Returns ``(replies, variables_to_keep, variables_to_drop)``: error markers - are formatted into their reply and dropped from the variables, and so is a - find_tools listing — the reply carries it, CodeAct never keeps it either. + Returns ``(replies, variables_to_keep, variables_to_drop)``. A call with no + outcome never ran — the block timed out before it, or failed outside the + per-call guard — and is answered with why. A find_tools listing is replied but + dropped from the variables, as CodeAct never keeps it either. """ - from cuga.backend.cuga_graph.nodes.cuga_lite.executors.code_executor import is_find_tools_listing_markdown - budget = _BatchOutputBudget(output_limit) - keep = dict(new_vars) drop: List[str] = [] unreached = _unreached_reason(output, timeout) - replies: List[ToolMessage] = [] - for entry in plan: - call_id, name = entry["id"], entry.get("name") or "" - pre = entry.get("reply") - if pre: - replies.append(_error_message(pre["content"], call_id=call_id, name=name)) - continue - var = entry.get("var") - if var is not None and var in new_vars: - value = new_vars[var] - if isinstance(value, dict) and _ERROR_KEY in value: - keep.pop(var, None) - drop.append(var) - replies.append(_error_message(_format_error(name, value), call_id=call_id, name=name)) - else: - if is_find_tools_listing_markdown(value): - keep.pop(var, None) - drop.append(var) - replies.append( - ToolMessage(content=budget.take(_stringify(value)), tool_call_id=call_id, name=name) - ) - continue - replies.append(_error_message(unreached.format(name=name), call_id=call_id, name=name)) + replies = [_reply_for(e, outcomes.get(e["id"]), new_vars, unreached, budget, drop) for e in plan] + keep = {k: v for k, v in new_vars.items() if k not in drop} return replies, keep, drop -def replies_without_execution(plan: List[Dict[str, Any]], reason: str) -> List[ToolMessage]: - """Every id answered when the block did not run (VERIFY revise, a failure outside the block). +def replies_without_execution( + plan: List[Dict[str, Any]], + reason: str, + outcomes: Optional[Dict[str, Dict[str, Any]]] = None, +) -> List[ToolMessage]: + """Every id answered when the block did not complete (VERIFY revise, a failure outside it). - Pre-decided replies are delivered as they are; every call that would have - run gets ``reason``. + Pre-decided replies are delivered as they are; calls whose outcome was + recorded before the failure keep it; every other call gets ``reason``. """ - out: List[ToolMessage] = [] - for entry in plan: - pre = entry.get("reply") - text = pre["content"] if pre else reason - out.append(_error_message(text, call_id=entry["id"], name=entry.get("name") or "")) - return out + budget = _BatchOutputBudget(0) + drop: List[str] = [] + return [_reply_for(e, (outcomes or {}).get(e["id"]), {}, reason, budget, drop) for e in plan] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py index ada340a46..ee7c9a02b 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/sandbox_node.py @@ -161,8 +161,10 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): # instead of an "Execution output:" message (see fc_actions). Empty for a # CodeAct block, where nothing below changes. fc_plan = fc_pending_plan(adapter, state) + fc_outcomes: dict = {} if fc_plan: - context.update(fc_context_overlay(adapter._tools_context, fc_plan)) + fc_overlay, fc_outcomes = fc_context_overlay(adapter._tools_context, fc_plan) + context.update(fc_overlay) fc_meta = {adapter.metadata_key: metadata_without_plan(adapter, state)} if fc_plan else {} # Start tool call tracking (enabled via invoke parameter, or internally @@ -320,11 +322,12 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): fc_replies, new_vars, dropped = replies_from_execution( fc_plan, new_vars, + fc_outcomes, output, output_limit=settings.advanced_features.execution_output_max_length, timeout=settings.advanced_features.sandbox_execution_timeout, ) - # Error markers: the reply carries them, the variables do not. + # A find_tools listing: the reply carries it, the variables do not. for name in dropped: if name in state.variables_manager.get_variable_names(): state.variables_manager.remove_variable(name) @@ -490,7 +493,7 @@ async def sandbox(state: Any, config: Optional[RunnableConfig] = None): error_msg = f"Error during execution: {str(e)}" logger.error(error_msg) new_messages = ( - replies_without_execution(fc_plan, NOT_EXECUTED_REPLY.format(reason=e)) + replies_without_execution(fc_plan, NOT_EXECUTED_REPLY.format(reason=e), fc_outcomes) if fc_plan else [HumanMessage(content=error_msg)] ) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py index b43c7bbec..6c4455e17 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py @@ -122,7 +122,8 @@ async def _run_batch( block, plan = plan_tool_calls(calls, list(invalid or []), tools, one_per_step=one_per_step) if block is None: return replies_without_execution(plan, ""), {}, None - context = {**tools, **fc_context_overlay(tools, plan)} + overlay, outcomes = fc_context_overlay(tools, plan) + context = {**tools, **overlay} ToolCallTracker.start_tracking(enabled=False, timings_only=False) ToolCallTracker.seed_call_budget(0, 0) try: @@ -136,7 +137,7 @@ async def _run_batch( finally: ToolCallTracker.stop_tracking() replies, keep, _drop = replies_from_execution( - plan, new_vars, output, output_limit=output_limit, timeout=timeout + plan, new_vars, outcomes, output, output_limit=output_limit, timeout=timeout ) return replies, keep, block @@ -184,7 +185,7 @@ async def test_a_failing_call_does_not_abort_its_siblings(): ("c1", "boom", "error", "Tool 'boom' failed: bad x"), ("c2", "add", "success", "3"), ] - assert kept == {"tool_result_c2": 3}, "the error marker is not kept as a variable" + assert kept == {"tool_result_c2": 3}, "a failed call leaves no variable behind" @pytest.mark.asyncio @@ -379,6 +380,30 @@ async def test_a_non_json_result_is_kept_as_the_reply_text(): assert kept == {"tool_result_c1": '"{1}"'}, "kept as the text the reply carries, never dropped" +@pytest.mark.asyncio +async def test_a_result_that_spells_the_old_error_marker_is_still_a_success(): + """Review (Nakagawa-master): a tool's return value is untrusted data, so an error must + never be recognised by a key inside it. Outcomes travel through the node-owned + channel instead; this JSON object is an ordinary success.""" + payload = {"__cuga_fc_error__": "legitimate field", "kind": "error", "other": 1} + + async def looks_like_a_marker() -> dict: + return payload + + replies, kept, _ = await _run_batch(_tools(looks_like_a_marker), [_call("looks_like_a_marker", {}, "c1")]) + assert _shape(replies) == [ + ("c1", "looks_like_a_marker", "success", json.dumps(payload, ensure_ascii=False)) + ] + assert kept == {"tool_result_c1": payload}, "kept intact as a variable" + + +@pytest.mark.asyncio +async def test_a_call_that_actually_raises_is_an_error_and_leaves_no_variable(): + replies, kept, _ = await _run_batch(_tools(boom), [_call("boom", {"x": 1}, "c1")]) + assert _shape(replies) == [("c1", "boom", "error", "Tool 'boom' failed: bad x")] + assert kept == {} + + # ── equivalence with the direct-call contract ─────────────────────────────── @@ -444,7 +469,23 @@ def _rand_json(rng: random.Random, depth: int = 0) -> Any: return {_rand_key(rng): _rand_json(rng, depth + 1) for _ in range(rng.randint(0, 3))} -_KEYS = ["a", "b", "x", "from", "class", "user-id", "q", "", "with space", "λ", "n"] +_KEYS = [ + "a", + "b", + "x", + "from", + "class", + "user-id", + "q", + "", + "with space", + "λ", + "n", + "__cuga_fc_error__", + "kind", + "ok", + "error", +] def _rand_key(rng: random.Random) -> str: @@ -480,7 +521,9 @@ async def test_random_calls_get_the_same_replies_as_a_direct_call(): # ── provider-safe aliases ──────────────────────────────────────────────────── -LONG_NAME = "public_review_platform_get_reviews_with_user_details_and_metadata_v2" # 68 chars, over the 64 limit +LONG_NAME = ( + "public_review_platform_get_reviews_with_user_details_and_metadata_v2" # 68 chars, over the 64 limit +) @pytest.mark.asyncio @@ -510,7 +553,8 @@ async def long_tool(n: int) -> int: assert f"await {LONG_NAME}(**" in block and alias not in block, block assert plan[0]["name"] == alias and plan[0]["target"] == LONG_NAME - context = {**tools, **fc_context_overlay(tools, plan)} + overlay, outcomes = fc_context_overlay(tools, plan) + context = {**tools, **overlay} ToolCallTracker.start_tracking(enabled=False, timings_only=False) ToolCallTracker.seed_call_budget(0, 0) try: @@ -523,7 +567,7 @@ async def long_tool(n: int) -> int: ) finally: ToolCallTracker.stop_tracking() - replies, kept, _ = replies_from_execution(plan, new_vars, output, output_limit=0, timeout=30) + replies, kept, _ = replies_from_execution(plan, new_vars, outcomes, output, output_limit=0, timeout=30) assert _shape(replies) == [("c1", alias, "success", "10")] assert RAN == [("long", 5)] and kept == {"tool_result_c1": 10} From 43c97d81266bd5b77f61d20d03fc3342dfc4519b Mon Sep 17 00:00:00 2001 From: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> Date: Tue, 13 Oct 2026 07:09:02 +0300 Subject: [PATCH 18/18] fix(cuga-lite): CodeRabbit round on the sandbox path - replies_without_execution: the reason is literal text (a VERIFY alert, an exception message), not a template. Brace-escape it before the per-id formatting; a VERIFY alert quoting JSON raised KeyError and ended the run instead of giving the model its revise feedback. - prepare: when the knowledge scope gate removes the knowledge app, the function-calling prompt's find_tools catalogue drops it too. - stream: a function-calling turn carries the translated block in script; the call_model handler now checks native tool calls first so the UI shows the calls, never the internal block. Each reproduced before the change; tests for all three. Signed-off-by: Segev Shlomov <105229001+segevshlomovIBM@users.noreply.github.com> --- .../nodes/cuga_lite/adapter/fc_actions.py | 5 +- .../nodes/cuga_lite/adapter/prepare_node.py | 5 ++ .../nodes/cuga_lite/tests/test_fc_actions.py | 9 ++++ .../nodes/cuga_lite/tests/test_fc_mode_e2e.py | 47 ++++++++++++++++++- .../backend/cuga_graph/utils/agent_loop.py | 32 +++++++------ tests/unit/test_agent_loop_fc_events.py | 25 ++++++++-- 6 files changed, 100 insertions(+), 23 deletions(-) diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py index cc6579345..d5a238840 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/fc_actions.py @@ -433,4 +433,7 @@ def replies_without_execution( """ budget = _BatchOutputBudget(0) drop: List[str] = [] - return [_reply_for(e, (outcomes or {}).get(e["id"]), {}, reason, budget, drop) for e in plan] + # ``reason`` is literal text (a VERIFY alert, an exception message), not a template: + # escape it, since ``_reply_for`` formats the unreached text with the tool name. + literal = (reason or "").replace("{", "{{").replace("}", "}}") + return [_reply_for(e, (outcomes or {}).get(e["id"]), {}, literal, budget, drop) for e in plan] diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py index 823719e01..fccbdda97 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/adapter/prepare_node.py @@ -582,6 +582,11 @@ async def prepare_tools_and_apps(state: Any, config: Optional[RunnableConfig] = apps_for_prompt = [ app for app in (apps_for_prompt or []) if getattr(app, "name", "") != "knowledge" ] + if _find_tools_apps is not None: + # The function-calling prompt's find_tools catalogue was captured above. + _find_tools_apps = [ + app for app in _find_tools_apps if getattr(app, "name", "") != "knowledge" + ] for tool_name in knowledge_tool_names: adapter._tools_context.pop(tool_name, None) elif knowledge_tool_names: diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py index 6c4455e17..142a586e2 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_actions.py @@ -319,6 +319,15 @@ def test_calls_the_block_never_reached_are_all_answered(): ], "pre-decided replies are delivered as they are" +def test_a_reason_with_braces_is_delivered_verbatim(): + """CodeRabbit: a VERIFY alert or an exception text often quotes JSON; it is literal text, + never a template, so braces must survive on every id.""" + _, plan = plan_tool_calls([_call("add", {"a": 1, "b": 2}, "c1")], [], _tools(add), one_per_step=False) + reason = 'VERIFY blocked this code block before execution. Alert: write with {"value": 9}' + replies = replies_without_execution(plan, reason) + assert _shape(replies) == [("c1", "add", "error", reason)] + + # ── what only the block form has to get right ─────────────────────────────── diff --git a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py index 02315c67b..9a67be5a3 100644 --- a/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py +++ b/src/cuga/backend/cuga_graph/nodes/cuga_lite/tests/test_fc_mode_e2e.py @@ -334,13 +334,13 @@ async def test_verify_revise_answers_every_call_without_running_it(monkeypatch): monkeypatch.setattr( sandbox_node, "decide_pre_execute_verify", - AsyncMock(return_value=VerifyDecision(gate="revise", alert="echoing 9 looks wrong")), + AsyncMock(return_value=VerifyDecision(gate="revise", alert='echoing {"value": 9} looks wrong')), ) def after_revise(messages): reply = _last_tool_message(messages) assert reply.tool_call_id == "c1" and reply.status == "error" - assert reply.content == verify_blocked_message("echoing 9 looks wrong") + assert reply.content == verify_blocked_message('echoing {"value": 9} looks wrong') return AIMessage(content="Understood, not echoing.") model = _ScriptedModel( @@ -481,3 +481,46 @@ def second(messages): isinstance(v.get("value"), str) and "Matching Tool(s)" in v["value"] for v in result["variables_storage"].values() ), "a find_tools listing is replied, never kept as a variable" + + +@pytest.mark.asyncio +async def test_a_scope_gated_knowledge_app_is_not_advertised_to_find_tools(monkeypatch): + """CodeRabbit: when the knowledge gate removes the knowledge app and its tools, the + function-calling prompt's find_tools catalogue must drop it too.""" + from types import SimpleNamespace + + from cuga.backend.cuga_graph.nodes.cuga_lite.adapter import prepare_node + + monkeypatch.setattr( + prepare_node, "_get_knowledge_tool_scope_context", lambda engine, thread: (set(), None) + ) + + async def knowledge_search(query: str) -> str: + """Search the knowledge base.""" + return "nothing" + + knowledge_tool = StructuredTool.from_function( + coroutine=knowledge_search, name="knowledge_search", description="Search the knowledge base." + ) + provider = _provider(_echo_tool(), knowledge_tool) + provider.get_apps = AsyncMock( + return_value=[ + SimpleNamespace(name="knowledge", type="api", description="Knowledge base"), + SimpleNamespace(name="test_app", type="api", description="Echo app"), + ] + ) + model = _ScriptedModel([AIMessage(content="done")]) + graph = create_cuga_lite_graph(model=model, tool_provider=provider, apps_list=[], thread_id="t").compile( + checkpointer=MemorySaver() + ) + + result = await graph.ainvoke( + CugaLiteState(chat_messages=[HumanMessage(content="hi")]), + config=_config( + "fc-knowledge", cuga_lite_execution_mode="function_calling", shortlisting_tool_threshold=0 + ), + ) + + prompt = result["prepared_prompt"] + assert "# AVAILABLE APPS" in prompt and "- test_app (api): Echo app" in prompt + assert "- knowledge (" not in prompt, prompt diff --git a/src/cuga/backend/cuga_graph/utils/agent_loop.py b/src/cuga/backend/cuga_graph/utils/agent_loop.py index 9be5c77e4..0aaf9df0b 100644 --- a/src/cuga/backend/cuga_graph/utils/agent_loop.py +++ b/src/cuga/backend/cuga_graph/utils/agent_loop.py @@ -406,23 +406,14 @@ def get_event_message(self, event) -> StreamEvent: # Handle call_model node (shared by CugaLite and CugaSupervisor subgraphs) if node_name == "call_model": logger.debug("Detected call_model node") - # Check if it generated code or text - if "script" in state_data and state_data["script"] and state_data["script"].strip(): - # Code generated - show as CodeAgent with JSON - logger.info(f"call_model generated script of length {len(state_data['script'])}") - output = { - "code": state_data["script"], - "execution_output": "", - "steps_summary": [], - "summary": "Code generated, preparing to execute", - "variables": state_data.get("variables_storage", {}), - } - return StreamEvent(name="CodeAgent", data=json.dumps(output)) - elif _native_tool_calls(subgraph_messages): - # Function-calling mode: show the native calls the way a code block is shown. + # Check if it generated native tool calls, code, or text. Native calls come + # first: a function-calling turn also carries the translated block in + # ``script``, and that block is internal — the UI shows the calls, never it. + native_calls = _native_tool_calls(subgraph_messages) + if native_calls: calls_text = "\n".join( f"{c.get('name')}({json.dumps(c.get('args') or {}, ensure_ascii=False, default=str)})" - for c in _native_tool_calls(subgraph_messages) + for c in native_calls if isinstance(c, dict) ) output = { @@ -433,6 +424,17 @@ def get_event_message(self, event) -> StreamEvent: "variables": state_data.get("variables_storage", {}), } return StreamEvent(name="CodeAgent", data=json.dumps(output)) + elif "script" in state_data and state_data["script"] and state_data["script"].strip(): + # Code generated - show as CodeAgent with JSON + logger.info(f"call_model generated script of length {len(state_data['script'])}") + output = { + "code": state_data["script"], + "execution_output": "", + "steps_summary": [], + "summary": "Code generated, preparing to execute", + "variables": state_data.get("variables_storage", {}), + } + return StreamEvent(name="CodeAgent", data=json.dumps(output)) else: # Text/reasoning output - only when last chat turn is a non-empty assistant message logger.info("call_model generated text response (no code)") diff --git a/tests/unit/test_agent_loop_fc_events.py b/tests/unit/test_agent_loop_fc_events.py index 47ebfac38..736f442b7 100644 --- a/tests/unit/test_agent_loop_fc_events.py +++ b/tests/unit/test_agent_loop_fc_events.py @@ -21,18 +21,33 @@ _CALL = {"name": "echo", "args": {"value": 7}, "id": "c1", "type": "tool_call"} -def _event(node, messages): +_BLOCK = 'try:\n tool_result_c1 = _fc_keep(await echo(**{"value": 7}))\n' + + +def _event(node, messages, script=None): return ( ("CugaLiteSubgraph",), - {node: {"chat_messages": messages, "script": None, "variables_storage": {}}}, + {node: {"chat_messages": messages, "script": script, "variables_storage": {}}}, ) -def test_native_tool_calls_stream_as_a_code_step(): +def test_native_tool_calls_stream_as_a_code_step_never_the_internal_block(): + """A function-calling turn carries the translated block in ``script``; the UI shows the calls.""" messages = [HumanMessage(content="q"), AIMessage(content="", tool_calls=[_CALL])] - ev = AgentLoop.get_event_message(MagicMock(), _event("call_model", messages)) + ev = AgentLoop.get_event_message(MagicMock(), _event("call_model", messages, script=_BLOCK)) assert ev.name == "CodeAgent" - assert json.loads(ev.data)["code"] == 'echo({"value": 7})' + data = json.loads(ev.data) + assert ( + data["code"] == 'echo({"value": 7})' and data["summary"] == "Tool calls issued, preparing to execute" + ) + assert "_fc_keep" not in ev.data + + +def test_a_codeact_script_still_streams_as_generated_code(): + messages = [HumanMessage(content="q"), AIMessage(content="```python\nprint(1)\n```")] + ev = AgentLoop.get_event_message(MagicMock(), _event("call_model", messages, script="print(1)")) + data = json.loads(ev.data) + assert data["code"] == "print(1)" and data["summary"] == "Code generated, preparing to execute" def test_tool_results_from_the_sandbox_stream_as_execution_output():