Inject tool failures with environment simulation

"""Replay the recorded Gemini turns against the same injected faults, offline.

Run: python -m unittest -v test_simulated.py
The model is replaced by its recorded turns; the environment simulation plugin, the retry
plugin and the tools are real.
"""
import asyncio
import json
import logging
import unittest
import warnings
from pathlib import Path
from types import SimpleNamespace

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk.apps import App  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.tools import FunctionTool  # noqa: E402
from google.adk.tools.environment_simulation.environment_simulation_config import (  # noqa: E402
    EnvironmentSimulationConfig, ToolSimulationConfig)
from google.adk.tools.environment_simulation.environment_simulation_engine import EnvironmentSimulationEngine  # noqa: E402
from google.genai import types  # noqa: E402

import record  # noqa: E402
import simulated  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))
READ_PAGE = FunctionTool(simulated.read_page)


class ReplayLlm(BaseLlm):
    turns: list

    async def generate_content_async(self, llm_request, stream=False):
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)))


def replay(run: str):
    """Run the recorded model turns of one run through the same app; return its trail and error."""
    options = dict(record.RUNS[run])
    app = simulated.build_app(**options)
    model = ReplayLlm(model="replay", turns=list(RECORD["runs"][run]["model_turns"]))
    app = App(name=app.name, root_agent=app.root_agent.clone(update={"model": model}), plugins=app.plugins)
    events, error = [], None
    try:
        asyncio.run(simulated.ask(app, RECORD["question"], events))
    except RuntimeError as exc:
        error = exc
    return simulated.show(events), error


def engine(*injections) -> EnvironmentSimulationEngine:
    return EnvironmentSimulationEngine(EnvironmentSimulationConfig(tool_simulation_configs=[
        ToolSimulationConfig(tool_name="read_page", injection_configs=list(injections))]))


def simulate(sim, path):
    return asyncio.run(sim.simulate(READ_PAGE, {"path": path}, None))


def setUpModule():
    warnings.simplefilter("ignore")  # unittest resets the filter; ADK warns about experimental features


class Faults(unittest.TestCase):
    def test_the_outage_replays_the_recorded_tool_results(self):
        trail, _ = replay("outage")
        self.assertEqual(trail, RECORD["runs"]["outage"]["trail"])
        self.assertEqual(sum("error_code" in line for line in trail), RECORD["runs"]["outage"]["read_page_failures"])

    def test_without_an_error_rule_the_agent_sometimes_cites_a_page_it_never_read(self):
        runs = [RECORD["runs"]["outage"], *RECORD["outage_repeats"]]
        self.assertTrue(any(r["cites_a_page"] for r in runs))
        self.assertGreater(min(r["read_page_failures"] for r in runs), 1)

    def test_the_error_rule_stops_after_one_failure_and_cites_nothing(self):
        run = RECORD["runs"]["outage-with-error-rule"]
        self.assertEqual((run["read_page_failures"], run["cites_a_page"]), (1, False))
        trail, _ = replay("outage-with-error-rule")
        self.assertEqual(trail, run["trail"])

    def test_the_retry_plugin_raises_once_retries_are_used_up(self):
        trail, error = replay("outage-with-retry-plugin")
        self.assertIn("reflect_retry_tool_plugin", str(error))
        self.assertEqual(sum("ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN" in line for line in trail), 2)

    def test_the_retry_plugin_sees_results_that_the_simulator_injected(self):
        trail, error = replay("outage-with-retry-limit-answer")
        self.assertIsNone(error)
        self.assertTrue(any("ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN" in line for line in trail))

    def test_non_throwing_exhaustion_returns_guidance_without_blocking_more_calls(self):
        plugin = simulated.RetryOnErrorCode(max_retries=2, throw_exception_if_retry_exceeded=False)
        async def failures():
            return [await plugin.after_tool_callback(
                tool=READ_PAGE, tool_args={"path": "install.md"},
                tool_context=SimpleNamespace(invocation_id="retry-demo"),
                result={"error_code": 503, "error_message": "Unavailable"}) for _ in range(4)]
        results = asyncio.run(failures())
        self.assertEqual(len(results), 4)
        self.assertEqual(results[2], results[3])
        self.assertIn("exceeded", str(results[3]).lower())

    def test_match_args_fails_only_the_matching_page(self):
        sim = engine(*simulated.FAULTS["missing-page"])
        self.assertEqual(simulate(sim, "install.md"), {"error_code": 404, "error_message": "Page not found."})
        self.assertIsNone(simulate(sim, "cli.md"))  # None: the real tool runs

    def test_a_random_seed_is_reapplied_on_every_call(self):
        always = engine(*simulated.FAULTS["flaky"])  # probability 0.5, random_seed=1
        self.assertEqual(sum(simulate(always, "cli.md") is not None for _ in range(20)), 20)
        never = engine(simulated.fail(503, "Docs service unavailable.", injection_probability=0.5, random_seed=2))
        self.assertEqual(sum(simulate(never, "cli.md") is not None for _ in range(20)), 0)


class SelectedScenario(unittest.TestCase):
    def test_launcher_attaches_error_rule_and_selected_plugins(self):
        from unittest.mock import AsyncMock, patch
        import run
        marker = object()
        args = SimpleNamespace(question="chosen", fault="missing-page", retry="answer")
        with patch.object(simulated, "build_app", return_value=marker) as build, patch.object(simulated, "ask", new_callable=AsyncMock) as ask:
            asyncio.run(run.main(args))
        build.assert_called_once_with(fault="missing-page", instruction=simulated.INSTRUCTION + simulated.ON_ERROR, retry="answer")
        self.assertIs(ask.await_args.args[0], marker)
        self.assertEqual(ask.await_args.args[1], "chosen")


if __name__ == "__main__":
    unittest.main()

"""Record the docs assistant on Gemini under each injected fault.

Run from this folder with GOOGLE_API_KEY set: python record.py
Writes recorded-runs.json (questions, model turns, tool results, answers) and run-output.txt.
"""
import asyncio
import json
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")

from google.adk.apps import App  # noqa: E402

import simulated  # noqa: E402

HERE = Path(__file__).resolve().parent
QUESTION = "How do I make the assistant use a local model instead of Gemini?"
FIXED = simulated.INSTRUCTION + simulated.ON_ERROR
RUNS = {
    "baseline": dict(fault=None),
    "outage": dict(fault="outage"),
    "outage-with-error-rule": dict(fault="outage", instruction=FIXED),
    "outage-with-retry-plugin": dict(fault="outage", retry="raise"),
    "outage-with-retry-limit-answer": dict(fault="outage", retry="answer"),
    "missing-page-with-error-rule": dict(fault="missing-page", instruction=FIXED),
}

async def run(options: dict) -> dict:
    turns = []

    def keep_turn(callback_context, llm_response):
        if llm_response.content:
            content = llm_response.content.model_dump(mode="json", exclude_none=True)
            for part in content.get("parts", []):
                part.pop("thought_signature", None)
            turns.append(content)

    app = simulated.build_app(**options)
    app = App(name=app.name, root_agent=app.root_agent.clone(update={"after_model_callback": keep_turn}),
              plugins=app.plugins)
    events, error = [], None
    try:
        await simulated.ask(app, QUESTION, events)
    except Exception as exc:  # a plugin may end the run by raising
        error = f"{type(exc).__name__}: {exc}"
    lines = simulated.show(events)
    answer = next((l[len("answer "):] for l in reversed(lines) if l.startswith("answer ")), None)
    return {"options": {k: (v == FIXED and "INSTRUCTION + ON_ERROR" or v) for k, v in options.items()},
            "model_turns": turns, "trail": lines, "error": error, "answer": answer,
            "read_page_failures": sum("error_code" in l for l in lines if l.startswith("result read_page")),
            "cites_a_page": bool(answer and answer.rstrip(".").endswith(".md)"))}


async def main() -> None:
    record = {"question": QUESTION, "model": simulated.MODEL, "runs": {}}
    output = []
    for name, options in RUNS.items():
        r = record["runs"][name] = await run(options)
        output.append("\n".join([f"# Run {name}", f"> {QUESTION}", *r["trail"]]
                                + ([f"raised {r['error']}"] if r["error"] else [])) + "\n")
    # The same outage five more times: how often the agent cites a page it never read.
    repeats = [await run(RUNS["outage"]) for _ in range(5)]
    record["outage_repeats"] = [{"read_page_failures": r["read_page_failures"], "cites_a_page": r["cites_a_page"],
                                 "answer": r["answer"]} for r in repeats]
    output.append("# Outage repeated 5 times: read_page failures, cites a page\n" + "\n".join(
        f"{r['read_page_failures']} failures, cites a page: {r['cites_a_page']}  | {r['answer']}" for r in repeats) + "\n")
    (HERE / "recorded-runs.json").write_text(json.dumps(record, indent=2) + "\n")
    (HERE / "run-output.txt").write_text("\n".join(output))
    print("\n".join(output))


if __name__ == "__main__":
    asyncio.run(main())