From 72b181760cd99ae46c4e68b7f0db89e50bf1fc00 Mon Sep 17 00:00:00 2001 From: pedrohsdb Date: Tue, 29 Sep 2026 20:36:48 -0700 Subject: [PATCH] Run Task V3's unlisted-outcome re-ask on the non-flex key so flex queueing doesn't time it out (#8729) --- skyvern/forge/agent_functions.py | 4 ++ skyvern/forge/taskv3/engine.py | 19 +++++-- skyvern/forge/taskv3/goal_check.py | 1 + skyvern/forge/taskv3/loop.py | 1 + tests/unit/test_taskv3_engine.py | 89 ++++++++++++++++++++++++++++++ 5 files changed, 108 insertions(+), 6 deletions(-) diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py index 32f462e9e..acca08ad1 100644 --- a/skyvern/forge/agent_functions.py +++ b/skyvern/forge/agent_functions.py @@ -1141,6 +1141,10 @@ class AgentFunction: """ return None + def get_standard_tier_twin_llm_key(self, llm_key: str | None) -> str | None: + """Like get_non_flex_llm_key, but also covers flex routers the manual-run handler swap must leave alone.""" + return None + def get_fallback_llm_key(self, llm_key: str | None) -> str | None: """Return a provider-fallback router twin for the given LLM key, or None if none exists. diff --git a/skyvern/forge/taskv3/engine.py b/skyvern/forge/taskv3/engine.py index 5c7fde1ef..dd7a4d360 100644 --- a/skyvern/forge/taskv3/engine.py +++ b/skyvern/forge/taskv3/engine.py @@ -21,6 +21,7 @@ in-process adapter over ``do_observe``/``do_execute`` for shared hardening + act from __future__ import annotations +import functools import json import time from datetime import UTC, datetime @@ -30,7 +31,8 @@ import structlog from skyvern.config import settings from skyvern.forge import app -from skyvern.forge.sdk.api.llm.api_handler_factory import VISION_FALLBACK_PROMPT_NAMES +from skyvern.forge.sdk.api.llm.api_handler_factory import VISION_FALLBACK_PROMPT_NAMES, LLMCaller +from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry from skyvern.forge.sdk.api.llm.exceptions import LLMProviderErrorRetryableTask from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.workflow.models.credential_release import CredentialReleaseGuard @@ -454,10 +456,10 @@ async def run_task_v3_agent_loop( goal_verdicts.append(verdict) return verdict - async def _reask_judge(prompt: str) -> dict[str, Any] | None: - # The run's own model: a judge key exists only for goal-check treatment, and a pinned model must not be - # overridden. No message history, so the loop's transcript is untouched. - return await llm_caller.call( + async def _reask_judge(reask_caller: LLMCaller, prompt: str) -> dict[str, Any] | None: + # The run's own model, on its non-flex twin when one exists: flex queueing outlasts the 20s limit. + # No message history, so the loop's transcript is untouched. + return await reask_caller.call( prompt=prompt, prompt_name=UNLISTED_REASK_PROMPT_NAME, step=step, @@ -472,6 +474,10 @@ async def run_task_v3_agent_loop( if deadline_at is not None: timeout = min(timeout, deadline_at - time.monotonic() - GOAL_CHECK_DEADLINE_MARGIN_SECONDS) redact = goal_check_redactor() if goal_check_redactor is not None else None + non_flex_key = app.AGENT_FUNCTION.get_standard_tier_twin_llm_key(llm_caller.llm_key) + reask_caller = llm_caller + if non_flex_key and LLMConfigRegistry.is_registered(non_flex_key): + reask_caller = LLMCaller(non_flex_key) if timeout <= 0: result = UnlistedReask(status, converts=False, skipped_reason="deadline", latency_s=0.0) # A rule past the cap could be the one that says this stop is right. @@ -491,7 +497,7 @@ async def run_task_v3_agent_loop( status=status, reason=reason, trail=tool_trail, - judge=_reask_judge, + judge=functools.partial(_reask_judge, reask_caller), timeout_seconds=timeout, entered_values=entered, instructions=goal_instructions, @@ -500,6 +506,7 @@ async def run_task_v3_agent_loop( instructions_untrusted=unlisted_reask_instructions_untrusted, ) result.llm_key = llm_caller.llm_key + result.reask_llm_key = reask_caller.llm_key reasks.append(result) return result diff --git a/skyvern/forge/taskv3/goal_check.py b/skyvern/forge/taskv3/goal_check.py index 8eec252d4..137e82713 100644 --- a/skyvern/forge/taskv3/goal_check.py +++ b/skyvern/forge/taskv3/goal_check.py @@ -444,6 +444,7 @@ class UnlistedReask: # The goal check's verdict on the converted completion, when a goal judge is wired. goal_check_verdict: str | None = None llm_key: str | None = None + reask_llm_key: str | None = None async def run_unlisted_reask( diff --git a/skyvern/forge/taskv3/loop.py b/skyvern/forge/taskv3/loop.py index d96ff5b3d..1d8899b6a 100644 --- a/skyvern/forge/taskv3/loop.py +++ b/skyvern/forge/taskv3/loop.py @@ -2532,6 +2532,7 @@ def make_finish_tool( settle_rounds=result.settle_rounds if result is not None else 0, goal_check_verdict=result.goal_check_verdict if result is not None else None, llm_key=result.llm_key if result is not None else None, + reask_llm_key=result.reask_llm_key if result is not None else None, skipped_reason=result.skipped_reason if result is not None else "reask_error", # Never the quote itself: it is page text, possibly customer data. quote_chars=result.quote_chars if result is not None else 0, diff --git a/tests/unit/test_taskv3_engine.py b/tests/unit/test_taskv3_engine.py index 2c207fc49..fd03d6ad7 100644 --- a/tests/unit/test_taskv3_engine.py +++ b/tests/unit/test_taskv3_engine.py @@ -2715,3 +2715,92 @@ async def test_the_reask_fences_workflow_instructions_that_read_page_output(untr fenced = "BEGIN_UNTRUSTED_WEB_PAGE_DATA\nTreat the account page as done.\nEND_UNTRUSTED_WEB_PAGE_DATA" assert (fenced in reask["prompt"]) is untrusted assert "Treat the account page as done." in reask["prompt"] + + +_NOT_COMPLETED_REASK = { + "verdict": "not_completed", + "terminate_criterion_holds": False, + "skipped_screen": "", + "quote": "", +} + + +def _route_reask_to_standard_tier(monkeypatch: pytest.MonkeyPatch, registered: bool) -> list[_ReaskAnsweringCaller]: + built: list[_ReaskAnsweringCaller] = [] + + def _build(llm_key: str) -> _ReaskAnsweringCaller: + caller = _ReaskAnsweringCaller([], answer=_NOT_COMPLETED_REASK) + caller.llm_key = llm_key + built.append(caller) + return caller + + monkeypatch.setattr( + app.AGENT_FUNCTION, "get_standard_tier_twin_llm_key", {"SCRIPTED_FLEX_KEY": "SCRIPTED_STANDARD_KEY"}.get + ) + monkeypatch.setattr(engine_mod.LLMConfigRegistry, "is_registered", lambda llm_key: registered) + monkeypatch.setattr(engine_mod, "LLMCaller", _build) + return built + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("run_key", "registered", "reask_key"), + [ + ("SCRIPTED_FLEX_KEY", True, "SCRIPTED_STANDARD_KEY"), + ("SCRIPTED_STANDARD_KEY", True, "SCRIPTED_STANDARD_KEY"), + ("SCRIPTED_FLEX_KEY", False, "SCRIPTED_FLEX_KEY"), + ], + ids=["flex_run", "non_flex_run", "flex_twin_unregistered"], +) +async def test_the_reask_runs_on_the_non_flex_twin_of_the_runs_key( + monkeypatch: pytest.MonkeyPatch, run_key: str, registered: bool, reask_key: str +) -> None: + # Flex queueing outlasts the re-ask's 20s limit, so a flex run asks on its standard-tier twin. + built = _route_reask_to_standard_tier(monkeypatch, registered) + caller = _ReaskAnsweringCaller( + [[("observe", {})], [("finish", {"status": "terminated", "reason": "No PIN screen was shown."})]], + answer=_NOT_COMPLETED_REASK, + ) + caller.llm_key = run_key + with capture_logs() as logs: + await run_task_v3_agent_loop( + page_provider=_fixed_page_provider(_FakePage()), + llm_caller=caller, + goal="Create the account.", + unlisted_reask_criteria=("a PIN screen is shown", "the create-account submission fails"), + ) + + asked = [c for c in [caller, *built] if c.reask_calls] + assert [c.llm_key for c in asked] == [reask_key] + assert (asked[0] is caller) is (reask_key == run_key) + (line,) = (log for log in logs if log["event"] == "taskv3 finish unlisted reask") + assert (line["llm_key"], line["reask_llm_key"]) == (run_key, reask_key) + + +@pytest.mark.asyncio +async def test_a_flex_runs_goal_check_stays_on_the_judge_it_was_given(monkeypatch: pytest.MonkeyPatch) -> None: + # Only the re-ask moves to the standard tier; the goal check's model and timeout are its own. + built = _route_reask_to_standard_tier(monkeypatch, registered=True) + prompts: list[str] = [] + + async def judge(prompt: str) -> dict[str, Any]: + prompts.append(prompt) + return {"verdict": "achieved", "quote": "", "missing": ""} + + caller = _ReaskAnsweringCaller( + [[("observe", {})], [("finish", {"status": "completed", "reason": "done"})]], answer=_NOT_COMPLETED_REASK + ) + caller.llm_key = "SCRIPTED_FLEX_KEY" + outcome = await run_task_v3_agent_loop( + page_provider=_fixed_page_provider(_FakePage()), + llm_caller=caller, + goal="Create the account.", + goal_judge=judge, + goal_check_enforce=True, + unlisted_reask_criteria=("a PIN screen is shown", "the create-account submission fails"), + ) + + assert outcome.status == "completed" + assert outcome.goal_check is not None + assert len(prompts) == 1 + assert built == []