runtime: harden LLM retry handling and operator workflows

tests/unit/admin_console/test_step_repository.py

@@ -73,3 +73,42 @@ def test_retrying_llm_payload_keeps_user_visible_retry_details():
         "source": "provider_sdk",
         "recoverable": True,
     }
+
+
+def test_terminal_llm_payload_keeps_wait_and_retry_aggregation():
+    repository = StepRepository()
+    retries = [
+        {
+            "delay": 5.0,
+            "scheduled_at": 100.0,
+            "provider": "google",
+            "source": "provider_sdk",
+        }
+    ]
+    trace = {
+        "type": "llm_call",
+        "name": "llm_pause",
+        "status": "failed",
+        "payload": json.dumps(
+            {
+                "error": "503 UNAVAILABLE",
+                "pause": True,
+                "request_id": "request-1",
+                "provider": "google",
+                "waited_seconds": 18.5,
+                "retries": retries,
+                "messages": ["private prompt"],
+            }
+        ),
+    }
+
+    normalized = repository._normalize_display_trace(trace)
+
+    assert normalized["payload"] == {
+        "error": "503 UNAVAILABLE",
+        "provider": "google",
+        "pause": True,
+        "request_id": "request-1",
+        "waited_seconds": 18.5,
+        "retries": retries,
+    }

tests/unit/agents/test_operator_prompts.py

@@ -0,0 +1,55 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from artemis.agents.operator.prompt_builder import (
+    OperatorPromptBuilder,
+    load_operator_prompts,
+)
+from artemis.agents.operator.prompts import apply_operator_prompt_contract
+
+
+def test_operator_prompt_contract_matches_note_runtime_semantics():
+    prompt = apply_operator_prompt_contract(load_operator_prompts()["main_template"])
+
+    assert "write-through tools `update_note` and `append_note`" in prompt
+    assert "may accompany at most one Turn-Ending Action" in prompt
+    assert "`read_note`, `list_notes`, and `save_note`" in prompt
+    assert "Do NOT submit a Turn-Ending Action at the same time" not in prompt
+    assert "memory note tools (`read_note`, `list_notes`, `save_note`, `update_note`" not in prompt
+
+
+def test_troubleshooter_prompt_has_mutually_exclusive_recovery_branches():
+    prompt = apply_operator_prompt_contract(load_operator_prompts()["troubleshooter_template"])
+
+    assert "invoke `reply_to_checker` directly" in prompt
+    assert "This branch takes precedence over diagnosis" in prompt
+    assert "perform that correction directly" in prompt
+    assert "also takes precedence over generic diagnosis triggers" in prompt
+    assert "A Checker rejection alone is not a mandatory diagnosis trigger" in prompt
+    assert "**Validation/Checker Rejection**" not in prompt
+
+
+def test_operator_prompt_omits_environment_trust_and_explorer_directives():
+    for template in load_operator_prompts().values():
+        prompt = apply_operator_prompt_contract(template)
+        assert "Untrusted Screen Content & Instruction Priority" not in prompt
+        assert "Visual Explorer Rule" not in prompt
+
+
+def test_compatibility_builder_uses_canonical_operator_template():
+    expected = apply_operator_prompt_contract(load_operator_prompts()["main_template"])
+
+    assert OperatorPromptBuilder.build_system_message() == expected
+    assert "wait_for_delay(seconds=" not in expected
+    assert 'press_key(key="home"' not in expected

tests/unit/test_context.py

@@ -42,3 +42,36 @@ async def test_artemis_context_waits_for_background_tasks():
 
     # After exiting the context, the task should have run to completion
     assert task_run is True
+
+
+@pytest.mark.asyncio
+async def test_artemis_context_cancels_background_tasks_after_grace_period():
+    device = DeviceContext(
+        host_platform="LINUX",
+        mobile_platform=DevicePlatform.ANDROID,
+        device_id="dummy",
+        device_width=1080,
+        device_height=2400,
+    )
+    context = ArtemisContext(
+        device=device,
+        background_task_grace_period_seconds=0.01,
+    )
+
+    task_cancelled = False
+
+    async def stuck_summary():
+        nonlocal task_cancelled
+        try:
+            await asyncio.Event().wait()
+        finally:
+            task_cancelled = True
+
+    async with context:
+        task = asyncio.create_task(stuck_summary(), name="stuck_summary")
+        context.background_tasks.append(task)
+
+    assert task.done()
+    assert task.cancelled()
+    assert task_cancelled is True
+    assert context.background_tasks == []

tests/unit/test_llm_pause.py

@@ -57,24 +57,98 @@ def test_provider_sdk_retry_is_persisted_and_published(monkeypatch):
         exc_info=None,
     )
 
-    llm._ProviderRetryTelemetryHandler().emit(record)
-
-    assert engine.trace_kwargs == {
-        "type": "llm_call",
-        "name": "llm_retry",
-        "payload": {
-            "error": "ServerError: 503 UNAVAILABLE: high demand",
-            "delay": 1.18,
+    request_token = llm._ACTIVE_LLM_REQUEST.set(
+        {
+            "request_id": "request-1",
             "provider": "google",
-            "source": "provider_sdk",
-            "recoverable": True,
-        },
-        "step_id": "step-7",
-        "parent_trace_id": None,
-        "status": "retrying",
+            "started_at": 100.0,
+            "retries": [],
+        }
+    )
+    try:
+        llm._ProviderRetryTelemetryHandler().emit(record)
+    finally:
+        llm._ACTIVE_LLM_REQUEST.reset(request_token)
+
+    assert engine.trace_kwargs["type"] == "llm_call"
+    assert engine.trace_kwargs["name"] == "llm_retry"
+    assert engine.trace_kwargs["status"] == "retrying"
+    assert engine.trace_kwargs["payload"] == {
+        "error": "ServerError: 503 UNAVAILABLE: high demand",
+        "delay": 1.18,
+        "provider": "google",
+        "source": "provider_sdk",
+        "recoverable": True,
+        "request_id": "request-1",
+        "scheduled_at": engine.trace_kwargs["payload"]["scheduled_at"],
     }
     event_type, payload = engine.published
     assert event_type == "llm_retrying"
     assert payload["trace_id"] == "trace-1"
     assert payload["delay"] == 1.18
     assert payload["error"].startswith("ServerError: 503")
+
+
+def test_provider_sdk_retry_is_hidden_for_non_gemini_request(monkeypatch):
+    engine = _RecordingEngine()
+    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
+    record = logging.LogRecord(
+        name="google_genai._api_client",
+        level=logging.INFO,
+        pathname=__file__,
+        lineno=1,
+        msg="Retrying request in 2 seconds as it raised 503 UNAVAILABLE",
+        args=(),
+        exc_info=None,
+    )
+    request_token = llm._ACTIVE_LLM_REQUEST.set(
+        {
+            "request_id": "request-openai",
+            "provider": "openai",
+            "started_at": 100.0,
+            "retries": [],
+        }
+    )
+    try:
+        llm._ProviderRetryTelemetryHandler().emit(record)
+    finally:
+        llm._ACTIVE_LLM_REQUEST.reset(request_token)
+
+    assert engine.trace_kwargs is None
+    assert engine.published is None
+
+
+def test_terminal_failure_aggregates_observed_retries_and_total_wait(monkeypatch, tmp_path):
+    engine = _RecordingEngine()
+    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
+    monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces"))
+    monkeypatch.setattr(llm.time, "time", lambda: 120.0)
+    retries = [
+        {
+            "error": "503 high demand",
+            "delay": 10.0,
+            "provider": "google",
+            "source": "provider_sdk",
+            "request_id": "request-2",
+            "scheduled_at": 101.0,
+            "timestamp": 101.0,
+        }
+    ]
+    request_token = llm._ACTIVE_LLM_REQUEST.set(
+        {
+            "request_id": "request-2",
+            "provider": "google",
+            "started_at": 100.0,
+            "retries": retries,
+        }
+    )
+    try:
+        llm._handle_llm_pause_and_resume(RuntimeError("503 high demand"))
+    finally:
+        llm._ACTIVE_LLM_REQUEST.reset(request_token)
+
+    payload = engine.trace_kwargs["payload"]
+    assert payload["request_id"] == "request-2"
+    assert payload["provider"] == "google"
+    assert payload["waited_seconds"] == 20.0
+    assert payload["retries"] == retries