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

@@ -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.
 

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
 

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(

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,

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 == []