From e726c58e2d1b96e1ed3b6e88b8e092ca57b4d548 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 02:46:33 -0700 Subject: [PATCH 1/9] feat(optimization): add run-scoped usage, budget, and cancellation contracts (OptimizationRunContext, OptimizerCapabilities, terminal_status) RFC: google/adk-python#6357. PR 1 of 2 (contracts only; no optimizer behavior change). Adds the one-shot OptimizationRunContext ledger with truthful usage coverage (verified/partial/unreported, never zero-coerced), call/token budgets with commit-then-terminate overshoot semantics and configurable on_budget_exceeded (raise | return_partial), cooperative cancellation, conservative OptimizerCapabilities defaults on AgentOptimizer, an optional keyword-only run_context parameter on optimize(), and an experimental terminal_status field on OptimizerResult. Enforcement in the two prompt optimizers lands in the stacked PR 2. --- .../adk/optimization/agent_optimizer.py | 19 + src/google/adk/optimization/data_types.py | 11 + src/google/adk/optimization/run_context.py | 419 ++++++++++++++++++ .../optimization/run_context_test.py | 197 ++++++++ 4 files changed, 646 insertions(+) create mode 100644 src/google/adk/optimization/run_context.py create mode 100644 tests/unittests/optimization/run_context_test.py diff --git a/src/google/adk/optimization/agent_optimizer.py b/src/google/adk/optimization/agent_optimizer.py index 7a6da422b09..ff9b44cadca 100644 --- a/src/google/adk/optimization/agent_optimizer.py +++ b/src/google/adk/optimization/agent_optimizer.py @@ -17,22 +17,37 @@ from abc import ABC from abc import abstractmethod from typing import Generic +from typing import Optional from ..agents.llm_agent import Agent from .data_types import AgentWithScoresT from .data_types import OptimizerResult from .data_types import SamplingResultT +from .run_context import OptimizationRunContext +from .run_context import OptimizerCapabilities from .sampler import Sampler class AgentOptimizer(ABC, Generic[SamplingResultT, AgentWithScoresT]): """Base class for agent optimizers.""" + @property + def capabilities(self) -> OptimizerCapabilities: + """Run-context instrumentation this optimizer supports (experimental). + + Conservative defaults (all unsupported) are correct for optimizers that + do not implement the run-context seam. Callers should check capabilities + and call ``optimize`` without a ``run_context`` when unsupported. + """ + return OptimizerCapabilities() + @abstractmethod async def optimize( self, initial_agent: Agent, sampler: Sampler[SamplingResultT], + *, + run_context: Optional[OptimizationRunContext] = None, ) -> OptimizerResult[AgentWithScoresT]: """Runs the optimizer. @@ -40,6 +55,10 @@ async def optimize( initial_agent: The initial agent to be optimized. sampler: The interface used to get training and validation example UIDs, request agent evaluations, and get useful data for optimizing the agent. + run_context: Optional one-shot run-scoped usage/budget/cancellation + controls (experimental). Optimizers whose ``capabilities`` report no + support must raise ``UnsupportedOptimizationContextError`` when a + context is supplied, before any sampler or model work. Returns: The final result of the optimization process, containing the optimized diff --git a/src/google/adk/optimization/data_types.py b/src/google/adk/optimization/data_types.py index 603ba5a44ea..c6aa68153a7 100644 --- a/src/google/adk/optimization/data_types.py +++ b/src/google/adk/optimization/data_types.py @@ -75,6 +75,17 @@ class OptimizerResult(BaseModel, Generic[AgentWithScoresT]): ), ) + terminal_status: Optional[str] = Field( + default=None, + description=( + "Run-context terminal status (experimental). None or 'completed'" + " for a normal run; 'budget_exceeded' when an OptimizationRunContext" + " configured on_budget_exceeded='return_partial' stopped the run," + " so a governed partial result is never mistaken for an unmarked" + " success." + ), + ) + class UnstructuredSamplingResult(SamplingResult): """Evaluation result providing per-example unstructured evaluation data.""" diff --git a/src/google/adk/optimization/run_context.py b/src/google/adk/optimization/run_context.py new file mode 100644 index 00000000000..8b7400f7ef8 --- /dev/null +++ b/src/google/adk/optimization/run_context.py @@ -0,0 +1,419 @@ +# 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 +# +# http://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. + +"""Run-scoped usage, budget, and cancellation controls for agent optimizers. + +Experimental. All names in this module are subject to change until the +callback/error semantics have at least one downstream implementation. + +An :class:`OptimizationRunContext` is an optional, one-shot, caller-owned +object attached to a single ``AgentOptimizer.optimize(...)`` run. It records +every *logical* optimizer-owned model invocation (control-plane metadata only, +never prompt or response content), enforces configured call/token budgets, and +carries a cooperative cancellation signal that instrumented optimizers observe +between logical model invocations. + +The ledger records provider-reported usage truthfully: missing token counters +stay ``None`` and are classified via :class:`UsageCoverage`, never coerced to +zero. The final :class:`OptimizationRunSnapshot` is readable from the context +on success *and* failure, so a governance caller can persist the attempt in a +``finally`` block without parsing logs. +""" + +from __future__ import annotations + +import enum +import threading +import time +from typing import Any +from typing import Literal +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field + + +class UsageCoverage(str, enum.Enum): + """How completely the provider reported token usage for one logical call.""" + + VERIFIED = "verified" + """The provider supplied an authoritative total token count.""" + + PARTIAL = "partial" + """At least one token counter was supplied without an authoritative total.""" + + UNREPORTED = "unreported" + """No token counter was supplied.""" + + +class ModelCallStage(str, enum.Enum): + """Which optimizer stage issued the logical model invocation.""" + + CANDIDATE_GENERATION = "candidate_generation" + REFLECTION = "reflection" + + +class ModelCallState(str, enum.Enum): + """Terminal state of one logical model invocation.""" + + COMPLETED = "completed" + PROVIDER_ERROR = "provider_error" + CANCELLED = "cancelled" + BUDGET_EXCEEDED = "budget_exceeded" + + +class ModelCallEvent(BaseModel): + """Control-plane record of one logical model invocation. + + Prompt and response content are deliberately not part of this event; + content capture remains an explicit evaluator/sampler concern. + """ + + sequence: int + stage: ModelCallStage + requested_model: Optional[str] = None + returned_model_version: Optional[str] = None + start_time: float + end_time: Optional[float] = None + state: Optional[ModelCallState] = None + prompt_tokens: Optional[int] = None + output_tokens: Optional[int] = None + reasoning_tokens: Optional[int] = None + cached_tokens: Optional[int] = None + tool_use_tokens: Optional[int] = None + total_tokens: Optional[int] = None + usage_coverage: Optional[UsageCoverage] = None + error_message: Optional[str] = None + + +class OptimizationRunSnapshot(BaseModel): + """Immutable view of the run ledger at a point in time.""" + + events: list[ModelCallEvent] = Field(default_factory=list) + started_calls: int = 0 + completed_calls: int = 0 + cumulative_total_tokens: int = 0 + """Sum of provider-reported authoritative totals (verified calls only).""" + + usage_coverage: UsageCoverage = UsageCoverage.UNREPORTED + """VERIFIED only if every completed call was verified; PARTIAL if any call + reported any counter; UNREPORTED otherwise.""" + + cancel_requested: bool = False + cancel_reason: Optional[str] = None + terminal_control_state: Optional[str] = None + """Run-level control outcome (e.g. ``call_budget_rejected``, + ``budget_exceeded``, ``cancelled``); ``None`` while running or on a + normal completion.""" + + +class OptimizationBudgets(BaseModel): + """Configured ceilings for optimizer-owned logical model invocations.""" + + max_model_calls: Optional[int] = Field( + default=None, + description=( + "Maximum number of logical optimizer-owned model invocations that" + " may start. Checked before each call; a preflight rejection does" + " not create a model-call event." + ), + ) + max_total_tokens: Optional[int] = Field( + default=None, + description=( + "Hard ceiling on cumulative provider-reported total tokens." + " Checked after each terminal response is committed; an over-budget" + " final call is committed first, then the run terminates." + ), + ) + on_budget_exceeded: Literal["raise", "return_partial"] = Field( + default="raise", + description=( + "Terminal behavior on budget exhaustion. Both modes stop scheduling" + " immediately after the final in-flight call settles. 'raise' raises" + " OptimizationBudgetExceeded; 'return_partial' lets the optimizer" + " return its best-so-far result marked terminal_status=" + "'budget_exceeded' (never an unmarked success)." + ), + ) + + +class OptimizationRunContextError(Exception): + """Base class for run-context errors carrying the final snapshot.""" + + def __init__(self, message: str, snapshot: OptimizationRunSnapshot): + super().__init__(message) + self.snapshot = snapshot + + +class OptimizationBudgetExceeded(OptimizationRunContextError): + """A configured call or token budget was exhausted.""" + + +class OptimizationCancelledError(OptimizationRunContextError): + """The run stopped because the context's cancellation was requested.""" + + +class UnsupportedOptimizationContextError(Exception): + """The optimizer does not support run-context instrumentation.""" + + +class ContextAlreadyAttachedError(Exception): + """A one-shot context was attached to more than one optimization run.""" + + +class _CallHandle: + """Opaque handle for one in-flight logical model invocation.""" + + def __init__(self, event: ModelCallEvent): + self._event = event + self._closed = False + + +class OptimizationRunContext: + """One-shot, concurrency-safe ledger and control channel for one run. + + Instrumented optimizers call :meth:`begin_model_call` immediately before a + logical ``BaseLlm.generate_content_async()`` invocation and + :meth:`end_model_call` with the terminal response, in-band error, or raised + provider error. Governance callers own the context, may call + :meth:`request_cancel` from any thread, and read :meth:`snapshot` at any + time, including after a failure. + """ + + def __init__(self, budgets: Optional[OptimizationBudgets] = None): + self._budgets = budgets or OptimizationBudgets() + self._lock = threading.Lock() + self._events: list[ModelCallEvent] = [] + self._started_calls = 0 + self._completed_calls = 0 + self._cumulative_total_tokens = 0 + self._cancel_requested = False + self._cancel_reason: Optional[str] = None + self._terminal_control_state: Optional[str] = None + self._attached_owner: Optional[object] = None + + @property + def budgets(self) -> OptimizationBudgets: + return self._budgets + + def attach(self, owner: object) -> None: + """Binds the context to one optimization run; reuse is rejected.""" + with self._lock: + if self._attached_owner is not None: + raise ContextAlreadyAttachedError( + "OptimizationRunContext is one-shot: it is already attached to an" + " optimization run and cannot be reused or shared." + ) + self._attached_owner = owner + + def request_cancel(self, reason: str = "requested") -> None: + """Thread-safe, idempotent cooperative cancellation signal.""" + with self._lock: + if not self._cancel_requested: + self._cancel_requested = True + self._cancel_reason = reason + + @property + def cancel_requested(self) -> bool: + with self._lock: + return self._cancel_requested + + def raise_if_cancelled(self) -> None: + """Raises ``OptimizationCancelledError`` if cancellation was requested.""" + with self._lock: + if not self._cancel_requested: + return + self._terminal_control_state = "cancelled" + snapshot = self._snapshot_locked() + reason = self._cancel_reason + raise OptimizationCancelledError( + f"Optimization cancelled: {reason}", snapshot + ) + + def begin_model_call( + self, + stage: ModelCallStage, + requested_model: Optional[str] = None, + ) -> _CallHandle: + """Reserves and records the start of one logical model invocation. + + Checks cancellation and the logical-call budget *before* the invocation + starts. A preflight rejection does not create a model-call event and does + not increment the started-call count; it records the run-level control + state and raises. + """ + with self._lock: + if self._cancel_requested: + self._terminal_control_state = "cancelled" + snapshot = self._snapshot_locked() + reason = self._cancel_reason + raise OptimizationCancelledError( + f"Optimization cancelled: {reason}", snapshot + ) + max_calls = self._budgets.max_model_calls + if max_calls is not None and self._started_calls >= max_calls: + self._terminal_control_state = "call_budget_rejected" + snapshot = self._snapshot_locked() + raise OptimizationBudgetExceeded( + f"Logical model-call budget exhausted ({max_calls}).", snapshot + ) + self._started_calls += 1 + event = ModelCallEvent( + sequence=self._started_calls, + stage=stage, + requested_model=requested_model, + start_time=time.monotonic(), + ) + self._events.append(event) + return _CallHandle(event) + + def end_model_call( + self, + handle: _CallHandle, + *, + usage_metadata: Any = None, + returned_model_version: Optional[str] = None, + error_message: Optional[str] = None, + ) -> None: + """Commits the terminal outcome of one logical model invocation. + + Token usage is committed before budget enforcement: if the new cumulative + total crosses the hard ceiling, the completed call is persisted first and + ``OptimizationBudgetExceeded`` is raised immediately afterwards, so an + over-budget final call can never produce an unmarked success. A terminal + provider error closes the call as ``provider_error``, preserves any + reported usage, and re-raising the primary error remains the caller's + responsibility. + """ + if handle._closed: + return + handle._closed = True + event = handle._event + with self._lock: + event.end_time = time.monotonic() + event.returned_model_version = returned_model_version + _apply_usage(event, usage_metadata) + if error_message is not None: + event.state = ModelCallState.PROVIDER_ERROR + event.error_message = error_message + else: + event.state = ModelCallState.COMPLETED + self._completed_calls += 1 + if event.total_tokens is not None: + self._cumulative_total_tokens += event.total_tokens + max_tokens = self._budgets.max_total_tokens + over_budget = ( + max_tokens is not None + and self._cumulative_total_tokens > max_tokens + ) + if over_budget: + self._terminal_control_state = "budget_exceeded" + snapshot = self._snapshot_locked() + if over_budget: + raise OptimizationBudgetExceeded( + f"Token budget exhausted ({max_tokens}).", snapshot + ) + + def snapshot(self) -> OptimizationRunSnapshot: + with self._lock: + return self._snapshot_locked() + + def _snapshot_locked(self) -> OptimizationRunSnapshot: + completed = [ + e for e in self._events if e.state is not None + ] + if completed and all( + e.usage_coverage == UsageCoverage.VERIFIED for e in completed + ): + run_coverage = UsageCoverage.VERIFIED + elif any( + e.usage_coverage in (UsageCoverage.VERIFIED, UsageCoverage.PARTIAL) + for e in completed + ): + run_coverage = UsageCoverage.PARTIAL + else: + run_coverage = UsageCoverage.UNREPORTED + return OptimizationRunSnapshot( + events=[e.model_copy(deep=True) for e in self._events], + started_calls=self._started_calls, + completed_calls=self._completed_calls, + cumulative_total_tokens=self._cumulative_total_tokens, + usage_coverage=run_coverage, + cancel_requested=self._cancel_requested, + cancel_reason=self._cancel_reason, + terminal_control_state=self._terminal_control_state, + ) + + +def _apply_usage(event: ModelCallEvent, usage_metadata: Any) -> None: + """Copies provider-reported token counters onto the event, truthfully. + + Missing counters stay ``None``. Coverage: ``verified`` iff the provider + supplied ``total_token_count``; ``partial`` if any other counter was + supplied; ``unreported`` otherwise. + """ + if usage_metadata is None: + event.usage_coverage = UsageCoverage.UNREPORTED + return + + def _get(name: str) -> Optional[int]: + value = getattr(usage_metadata, name, None) + return int(value) if isinstance(value, (int, float)) else None + + event.prompt_tokens = _get("prompt_token_count") + event.output_tokens = _get("candidates_token_count") + event.reasoning_tokens = _get("thoughts_token_count") + event.cached_tokens = _get("cached_content_token_count") + event.tool_use_tokens = _get("tool_use_prompt_token_count") + event.total_tokens = _get("total_token_count") + if event.total_tokens is not None: + event.usage_coverage = UsageCoverage.VERIFIED + elif any( + v is not None + for v in ( + event.prompt_tokens, + event.output_tokens, + event.reasoning_tokens, + event.cached_tokens, + event.tool_use_tokens, + ) + ): + event.usage_coverage = UsageCoverage.PARTIAL + else: + event.usage_coverage = UsageCoverage.UNREPORTED + + +class OptimizerCapabilities(BaseModel): + """What run-context instrumentation an optimizer actually supports. + + Capabilities describe ADK instrumentation, not provider metadata guarantees + or transport-attempt visibility. Conservative defaults (all ``False``) are + correct for optimizers that predate or do not implement the run-context + seam, so a governance wrapper can reject an incompatible optimizer at + preflight instead of discovering opacity after spend occurs. + """ + + model_calls_observable: bool = False + """Optimizer-owned logical model invocations are recorded on the context.""" + + call_limits_enforceable: bool = False + """Configured logical-invocation limits stop the next call from starting.""" + + cooperative_cancellation: bool = False + """``request_cancel`` is observed at the documented boundaries.""" + + sampler_usage_included: bool = False + """Always ``False`` for this context: candidate execution and evaluator + inference belong to the supplied sampler/Runner, not the optimizer.""" diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py new file mode 100644 index 00000000000..ef0a412fc1e --- /dev/null +++ b/tests/unittests/optimization/run_context_test.py @@ -0,0 +1,197 @@ +# 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 +# +# http://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 __future__ import annotations + +import threading +from types import SimpleNamespace + +from google.adk.optimization.run_context import ContextAlreadyAttachedError +from google.adk.optimization.run_context import ModelCallStage +from google.adk.optimization.run_context import ModelCallState +from google.adk.optimization.run_context import OptimizationBudgetExceeded +from google.adk.optimization.run_context import OptimizationBudgets +from google.adk.optimization.run_context import OptimizationCancelledError +from google.adk.optimization.run_context import OptimizationRunContext +from google.adk.optimization.run_context import OptimizerCapabilities +from google.adk.optimization.run_context import UsageCoverage +import pytest + + +def _usage(**kwargs) -> SimpleNamespace: + return SimpleNamespace(**kwargs) + + +class TestOneShotAttachment: + + def test_attach_twice_rejected(self): + ctx = OptimizationRunContext() + ctx.attach(owner=object()) + with pytest.raises(ContextAlreadyAttachedError): + ctx.attach(owner=object()) + + def test_two_runs_have_isolated_contexts(self): + a, b = OptimizationRunContext(), OptimizationRunContext() + a.attach(owner="run-a") + b.attach(owner="run-b") + h = a.begin_model_call(ModelCallStage.REFLECTION) + a.end_model_call(h, usage_metadata=_usage(total_token_count=7)) + assert a.snapshot().completed_calls == 1 + assert b.snapshot().completed_calls == 0 + + +class TestCallBudget: + + def test_exactly_n_calls_may_start(self): + ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=2)) + for _ in range(2): + h = ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + ctx.end_model_call(h, usage_metadata=None) + with pytest.raises(OptimizationBudgetExceeded) as exc: + ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + snap = exc.value.snapshot + # The rejected reservation is not a call event and did not start a call. + assert snap.started_calls == 2 + assert len(snap.events) == 2 + assert snap.terminal_control_state == "call_budget_rejected" + + +class TestTokenBudget: + + def test_overshoot_commits_then_raises(self): + ctx = OptimizationRunContext(OptimizationBudgets(max_total_tokens=100)) + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + with pytest.raises(OptimizationBudgetExceeded) as exc: + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=150)) + snap = exc.value.snapshot + # The over-budget final call is committed before the raise. + assert snap.completed_calls == 1 + assert snap.events[0].state == ModelCallState.COMPLETED + assert snap.cumulative_total_tokens == 150 + assert snap.terminal_control_state == "budget_exceeded" + + def test_unreported_usage_does_not_consume_token_budget(self): + ctx = OptimizationRunContext(OptimizationBudgets(max_total_tokens=10)) + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call(h, usage_metadata=None) # no raise + assert ctx.snapshot().cumulative_total_tokens == 0 + + +class TestUsageClassification: + + def test_verified_partial_unreported(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call( + h, usage_metadata=_usage(prompt_token_count=5, total_token_count=9) + ) + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(prompt_token_count=5)) + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call(h, usage_metadata=None) + events = ctx.snapshot().events + assert events[0].usage_coverage == UsageCoverage.VERIFIED + assert events[1].usage_coverage == UsageCoverage.PARTIAL + assert events[1].total_tokens is None # never coerced to zero + assert events[2].usage_coverage == UsageCoverage.UNREPORTED + # Run-level coverage degrades to partial, not verified. + assert ctx.snapshot().usage_coverage == UsageCoverage.PARTIAL + + def test_run_coverage_verified_only_when_all_verified(self): + ctx = OptimizationRunContext() + for _ in range(2): + h = ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=3)) + assert ctx.snapshot().usage_coverage == UsageCoverage.VERIFIED + + +class TestProviderError: + + def test_error_preserves_usage_so_far(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call( + h, + usage_metadata=_usage(total_token_count=42), + error_message="RESOURCE_EXHAUSTED", + ) + event = ctx.snapshot().events[0] + assert event.state == ModelCallState.PROVIDER_ERROR + assert event.total_tokens == 42 + assert event.error_message == "RESOURCE_EXHAUSTED" + + +class TestCancellation: + + def test_cancel_is_idempotent_and_first_reason_wins(self): + ctx = OptimizationRunContext() + ctx.request_cancel("deadline") + ctx.request_cancel("other") + with pytest.raises(OptimizationCancelledError) as exc: + ctx.begin_model_call(ModelCallStage.REFLECTION) + assert "deadline" in str(exc.value) + assert exc.value.snapshot.cancel_reason == "deadline" + assert exc.value.snapshot.terminal_control_state == "cancelled" + + def test_raise_if_cancelled_noop_when_not_cancelled(self): + OptimizationRunContext().raise_if_cancelled() + + def test_cancel_from_another_thread_observed(self): + ctx = OptimizationRunContext() + t = threading.Thread(target=ctx.request_cancel, args=("watchdog",)) + t.start() + t.join() + with pytest.raises(OptimizationCancelledError): + ctx.raise_if_cancelled() + + def test_snapshot_readable_after_failure(self): + ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=0)) + with pytest.raises(OptimizationBudgetExceeded): + ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + # A governance caller can persist the attempt from a finally block. + assert ctx.snapshot().terminal_control_state == "call_budget_rejected" + + +class TestCapabilitiesDefaults: + + def test_conservative_defaults(self): + caps = OptimizerCapabilities() + assert not caps.model_calls_observable + assert not caps.call_limits_enforceable + assert not caps.cooperative_cancellation + assert not caps.sampler_usage_included + + def test_base_optimizer_reports_conservative_capabilities(self): + from google.adk.optimization.agent_optimizer import AgentOptimizer + + class _Impl(AgentOptimizer): + + async def optimize(self, initial_agent, sampler, *, run_context=None): + raise NotImplementedError() + + assert _Impl().capabilities == OptimizerCapabilities() + + +class TestResultTerminalStatus: + + def test_default_none(self): + from google.adk.optimization.data_types import AgentWithScores + from google.adk.optimization.data_types import OptimizerResult + + result = OptimizerResult(optimized_agents=[]) + assert result.terminal_status is None + result = OptimizerResult( + optimized_agents=[], terminal_status="budget_exceeded" + ) + assert result.terminal_status == "budget_exceeded" From f345eee1e2bdb171abd58ca0fde8c5a8359b0956 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 03:36:13 -0700 Subject: [PATCH 2/9] align contracts with revised RFC: call/run status split, indeterminate token compliance, granular capabilities, OptimizerResult unchanged Per the precision pass on RFC #6357: logical-call limits are the only hard enforcement (atomic slot admission); reported-token ceilings terminate reactively over authoritative totals with within_limit/exceeded/indeterminate compliance (missing totals are never proof of compliance); call status is split from run status; capabilities gain accepts_run_context and split logical-call vs reported-token enforceability; terminal_status is removed from OptimizerResult -- the caller-owned snapshot is authoritative; adds OptimizationProviderError for governed in-band error termination. --- src/google/adk/optimization/data_types.py | 11 --- src/google/adk/optimization/run_context.py | 98 +++++++++++++++---- .../optimization/run_context_test.py | 55 +++++++---- 3 files changed, 116 insertions(+), 48 deletions(-) diff --git a/src/google/adk/optimization/data_types.py b/src/google/adk/optimization/data_types.py index c6aa68153a7..603ba5a44ea 100644 --- a/src/google/adk/optimization/data_types.py +++ b/src/google/adk/optimization/data_types.py @@ -75,17 +75,6 @@ class OptimizerResult(BaseModel, Generic[AgentWithScoresT]): ), ) - terminal_status: Optional[str] = Field( - default=None, - description=( - "Run-context terminal status (experimental). None or 'completed'" - " for a normal run; 'budget_exceeded' when an OptimizationRunContext" - " configured on_budget_exceeded='return_partial' stopped the run," - " so a governed partial result is never mistaken for an unmarked" - " success." - ), - ) - class UnstructuredSamplingResult(SamplingResult): """Evaluation result providing per-example unstructured evaluation data.""" diff --git a/src/google/adk/optimization/run_context.py b/src/google/adk/optimization/run_context.py index 8b7400f7ef8..de0f30a3297 100644 --- a/src/google/adk/optimization/run_context.py +++ b/src/google/adk/optimization/run_context.py @@ -65,12 +65,38 @@ class ModelCallStage(str, enum.Enum): class ModelCallState(str, enum.Enum): - """Terminal state of one logical model invocation.""" + """Terminal state of one logical model invocation (call status). + + Budget exhaustion is a *run* status, not a call status: the triggering + successful call remains ``completed`` while the run becomes + ``budget_exceeded``. + """ COMPLETED = "completed" PROVIDER_ERROR = "provider_error" CANCELLED = "cancelled" + + +class RunStatus(str, enum.Enum): + """Terminal status of the optimization run (run status).""" + + COMPLETED = "completed" BUDGET_EXCEEDED = "budget_exceeded" + CANCELLED = "cancelled" + FAILED = "failed" + + +class TokenBudgetStatus(str, enum.Enum): + """Compliance of actual token usage with the configured reported-token limit. + + ``indeterminate`` whenever any completed call's usage coverage is partial or + unreported: missing totals are never interpreted as proof that actual usage + stayed below the ceiling. Downstream policy may fail closed on that state. + """ + + WITHIN_LIMIT = "within_limit" + EXCEEDED = "exceeded" + INDETERMINATE = "indeterminate" class ModelCallEvent(BaseModel): @@ -110,12 +136,13 @@ class OptimizationRunSnapshot(BaseModel): """VERIFIED only if every completed call was verified; PARTIAL if any call reported any counter; UNREPORTED otherwise.""" + token_budget_status: Optional[TokenBudgetStatus] = None + """Only set when a reported-token limit is configured.""" + cancel_requested: bool = False cancel_reason: Optional[str] = None - terminal_control_state: Optional[str] = None - """Run-level control outcome (e.g. ``call_budget_rejected``, - ``budget_exceeded``, ``cancelled``); ``None`` while running or on a - normal completion.""" + run_status: Optional[RunStatus] = None + """Terminal run status; ``None`` while the run is still in progress.""" class OptimizationBudgets(BaseModel): @@ -129,12 +156,15 @@ class OptimizationBudgets(BaseModel): " not create a model-call event." ), ) - max_total_tokens: Optional[int] = Field( + max_provider_reported_tokens: Optional[int] = Field( default=None, description=( - "Hard ceiling on cumulative provider-reported total tokens." - " Checked after each terminal response is committed; an over-budget" - " final call is committed first, then the run terminates." + "Ceiling on cumulative provider-reported authoritative total" + " tokens. Checked after each terminal response is committed; an" + " over-budget final call is committed first, then the run" + " terminates. Deliberately not a hard bound on billing or physical" + " tokens: when any call's usage is partial or unreported," + " actual-token compliance is indeterminate." ), ) on_budget_exceeded: Literal["raise", "return_partial"] = Field( @@ -165,6 +195,14 @@ class OptimizationCancelledError(OptimizationRunContextError): """The run stopped because the context's cancellation was requested.""" +class OptimizationProviderError(OptimizationRunContextError): + """A governed run terminated on an in-band provider error. + + Usage reported before the error is preserved on the snapshot; a governed + run must not silently succeed past an in-band ``LlmResponse.error_code``. + """ + + class UnsupportedOptimizationContextError(Exception): """The optimizer does not support run-context instrumentation.""" @@ -201,7 +239,7 @@ def __init__(self, budgets: Optional[OptimizationBudgets] = None): self._cumulative_total_tokens = 0 self._cancel_requested = False self._cancel_reason: Optional[str] = None - self._terminal_control_state: Optional[str] = None + self._run_status: Optional[RunStatus] = None self._attached_owner: Optional[object] = None @property @@ -235,7 +273,7 @@ def raise_if_cancelled(self) -> None: with self._lock: if not self._cancel_requested: return - self._terminal_control_state = "cancelled" + self._run_status = RunStatus.CANCELLED snapshot = self._snapshot_locked() reason = self._cancel_reason raise OptimizationCancelledError( @@ -256,7 +294,7 @@ def begin_model_call( """ with self._lock: if self._cancel_requested: - self._terminal_control_state = "cancelled" + self._run_status = RunStatus.CANCELLED snapshot = self._snapshot_locked() reason = self._cancel_reason raise OptimizationCancelledError( @@ -264,7 +302,9 @@ def begin_model_call( ) max_calls = self._budgets.max_model_calls if max_calls is not None and self._started_calls >= max_calls: - self._terminal_control_state = "call_budget_rejected" + # Atomic slot admission: the rejected reservation is not a call + # event; the run status records why the run ended. + self._run_status = RunStatus.BUDGET_EXCEEDED snapshot = self._snapshot_locked() raise OptimizationBudgetExceeded( f"Logical model-call budget exhausted ({max_calls}).", snapshot @@ -313,13 +353,13 @@ def end_model_call( self._completed_calls += 1 if event.total_tokens is not None: self._cumulative_total_tokens += event.total_tokens - max_tokens = self._budgets.max_total_tokens + max_tokens = self._budgets.max_provider_reported_tokens over_budget = ( max_tokens is not None and self._cumulative_total_tokens > max_tokens ) if over_budget: - self._terminal_control_state = "budget_exceeded" + self._run_status = RunStatus.BUDGET_EXCEEDED snapshot = self._snapshot_locked() if over_budget: raise OptimizationBudgetExceeded( @@ -345,15 +385,28 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: run_coverage = UsageCoverage.PARTIAL else: run_coverage = UsageCoverage.UNREPORTED + token_status = None + if self._budgets.max_provider_reported_tokens is not None: + if ( + self._cumulative_total_tokens + > self._budgets.max_provider_reported_tokens + ): + token_status = TokenBudgetStatus.EXCEEDED + elif run_coverage == UsageCoverage.VERIFIED: + token_status = TokenBudgetStatus.WITHIN_LIMIT + else: + # Missing totals are never proof of compliance. + token_status = TokenBudgetStatus.INDETERMINATE return OptimizationRunSnapshot( events=[e.model_copy(deep=True) for e in self._events], started_calls=self._started_calls, completed_calls=self._completed_calls, cumulative_total_tokens=self._cumulative_total_tokens, usage_coverage=run_coverage, + token_budget_status=token_status, cancel_requested=self._cancel_requested, cancel_reason=self._cancel_reason, - terminal_control_state=self._terminal_control_state, + run_status=self._run_status, ) @@ -405,11 +458,20 @@ class OptimizerCapabilities(BaseModel): preflight instead of discovering opacity after spend occurs. """ + accepts_run_context: bool = False + """``optimize`` accepts a run context at all. When ``False``, callers must + omit the keyword entirely (protects pre-existing third-party overrides).""" + model_calls_observable: bool = False """Optimizer-owned logical model invocations are recorded on the context.""" - call_limits_enforceable: bool = False - """Configured logical-invocation limits stop the next call from starting.""" + logical_call_limits_enforceable: bool = False + """Configured logical-invocation limits stop the next call from starting + (hard: atomic slot admission).""" + + reported_token_limits_enforceable: bool = False + """Reported-token ceilings terminate the run reactively after the + triggering call commits. Not a bound on unreported usage.""" cooperative_cancellation: bool = False """``request_cancel`` is observed at the documented boundaries.""" diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index ef0a412fc1e..c461c81421b 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -25,6 +25,8 @@ from google.adk.optimization.run_context import OptimizationCancelledError from google.adk.optimization.run_context import OptimizationRunContext from google.adk.optimization.run_context import OptimizerCapabilities +from google.adk.optimization.run_context import RunStatus +from google.adk.optimization.run_context import TokenBudgetStatus from google.adk.optimization.run_context import UsageCoverage import pytest @@ -64,28 +66,45 @@ def test_exactly_n_calls_may_start(self): # The rejected reservation is not a call event and did not start a call. assert snap.started_calls == 2 assert len(snap.events) == 2 - assert snap.terminal_control_state == "call_budget_rejected" + assert snap.run_status == RunStatus.BUDGET_EXCEEDED class TestTokenBudget: def test_overshoot_commits_then_raises(self): - ctx = OptimizationRunContext(OptimizationBudgets(max_total_tokens=100)) + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=100) + ) h = ctx.begin_model_call(ModelCallStage.REFLECTION) with pytest.raises(OptimizationBudgetExceeded) as exc: ctx.end_model_call(h, usage_metadata=_usage(total_token_count=150)) snap = exc.value.snapshot - # The over-budget final call is committed before the raise. + # The over-budget final call is committed before the raise; the CALL + # status stays completed while the RUN status becomes budget_exceeded. assert snap.completed_calls == 1 assert snap.events[0].state == ModelCallState.COMPLETED assert snap.cumulative_total_tokens == 150 - assert snap.terminal_control_state == "budget_exceeded" + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + assert snap.token_budget_status == TokenBudgetStatus.EXCEEDED - def test_unreported_usage_does_not_consume_token_budget(self): - ctx = OptimizationRunContext(OptimizationBudgets(max_total_tokens=10)) + def test_unreported_usage_is_indeterminate_not_compliant(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) h = ctx.begin_model_call(ModelCallStage.REFLECTION) ctx.end_model_call(h, usage_metadata=None) # no raise - assert ctx.snapshot().cumulative_total_tokens == 0 + snap = ctx.snapshot() + assert snap.cumulative_total_tokens == 0 + # Missing totals are never proof of compliance. + assert snap.token_budget_status == TokenBudgetStatus.INDETERMINATE + + def test_all_verified_under_limit_is_within_limit(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=100) + ) + h = ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=40)) + assert ctx.snapshot().token_budget_status == TokenBudgetStatus.WITHIN_LIMIT class TestUsageClassification: @@ -142,7 +161,7 @@ def test_cancel_is_idempotent_and_first_reason_wins(self): ctx.begin_model_call(ModelCallStage.REFLECTION) assert "deadline" in str(exc.value) assert exc.value.snapshot.cancel_reason == "deadline" - assert exc.value.snapshot.terminal_control_state == "cancelled" + assert exc.value.snapshot.run_status == RunStatus.CANCELLED def test_raise_if_cancelled_noop_when_not_cancelled(self): OptimizationRunContext().raise_if_cancelled() @@ -160,15 +179,17 @@ def test_snapshot_readable_after_failure(self): with pytest.raises(OptimizationBudgetExceeded): ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) # A governance caller can persist the attempt from a finally block. - assert ctx.snapshot().terminal_control_state == "call_budget_rejected" + assert ctx.snapshot().run_status == RunStatus.BUDGET_EXCEEDED class TestCapabilitiesDefaults: def test_conservative_defaults(self): caps = OptimizerCapabilities() + assert not caps.accepts_run_context assert not caps.model_calls_observable - assert not caps.call_limits_enforceable + assert not caps.logical_call_limits_enforceable + assert not caps.reported_token_limits_enforceable assert not caps.cooperative_cancellation assert not caps.sampler_usage_included @@ -183,15 +204,11 @@ async def optimize(self, initial_agent, sampler, *, run_context=None): assert _Impl().capabilities == OptimizerCapabilities() -class TestResultTerminalStatus: +class TestResultShapeUnchanged: - def test_default_none(self): - from google.adk.optimization.data_types import AgentWithScores + def test_optimizer_result_has_no_run_context_fields(self): + # The caller-owned snapshot is authoritative for run status; the public + # OptimizerResult schema stays untouched for observable equivalence. from google.adk.optimization.data_types import OptimizerResult - result = OptimizerResult(optimized_agents=[]) - assert result.terminal_status is None - result = OptimizerResult( - optimized_agents=[], terminal_status="budget_exceeded" - ) - assert result.terminal_status == "budget_exceeded" + assert "terminal_status" not in OptimizerResult.model_fields From bd5b94d3f123602bc524bd66c48b376b72d39cf5 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 04:24:11 -0700 Subject: [PATCH 3/9] harden contracts per adversarial review: terminal state machine, frozen ledger, provider-failure precedence, independent type-cleanliness - Explicit first-terminal-wins state machine (completed/budget_exceeded/ cancelled/failed) with finalize_success/finalize_cancelled/finalize_failed; begin_model_call rejects any terminal context; a late cancellation cannot overwrite an earlier terminal state. - Provider failure is a governed terminal: end_model_call with sanitized error_code/error_type commits usage, transitions the run to failed, and raises OptimizationProviderError; provider failure takes precedence over a simultaneous token overshoot while preserving usage and token-compliance evidence. Raw exception text never enters the ledger. - Ledger integrity: handles bound to their context (foreign commit is a typed error), atomic close under the context lock (concurrent double-close commits once), private unattached sentinel (attach(None) is one-shot), frozen OptimizationBudgets (ge=0 validated), frozen snapshots/events (tuples), snapshot carries budgets + terminal sequence + terminal error metadata; finalize_cancelled closes open call events. - Stage is an extensible string with documented constants. - Built-in optimizers accept the keyword and reject a supplied context with UnsupportedOptimizationContextError in this PR, so the abstract-signature change is independently type-clean; real support lands in the stacked enforcement PR. - isort + pyink applied. --- .../optimization/gepa_root_agent_optimizer.py | 10 + .../gepa_root_agent_prompt_optimizer.py | 11 + src/google/adk/optimization/run_context.py | 556 ++++++++++++------ .../optimization/simple_prompt_optimizer.py | 12 + .../optimization/run_context_test.py | 238 +++++++- 5 files changed, 630 insertions(+), 197 deletions(-) diff --git a/src/google/adk/optimization/gepa_root_agent_optimizer.py b/src/google/adk/optimization/gepa_root_agent_optimizer.py index 01d93aae029..a062860eda7 100644 --- a/src/google/adk/optimization/gepa_root_agent_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_optimizer.py @@ -333,6 +333,8 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context=None, ) -> GEPARootAgentOptimizerResult: """Runs the GEPARootAgentOptimizer. @@ -346,6 +348,14 @@ async def optimize( The final result of the optimization process, containing the optimized agent instance, its scores on the validation examples, and other metrics. """ + if run_context is not None: + from .run_context import UnsupportedOptimizationContextError + + raise UnsupportedOptimizationContextError( + "GEPARootAgentOptimizer does not support OptimizationRunContext in" + " this release; check optimizer.capabilities before passing a" + " context." + ) if initial_agent.sub_agents: logger.warning( "The GEPARootAgentOptimizer will not optimize prompts for sub-agents." diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py index e9b82cdd504..5adebfb140a 100644 --- a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -207,6 +207,8 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context=None, ) -> GEPARootAgentPromptOptimizerResult: """Runs the GEPARootAgentPromptOptimizer. @@ -220,6 +222,15 @@ async def optimize( The final result of the optimization process, containing the optimized agent instance, its scores on the validation examples, and other metrics. """ + if run_context is not None: + from .run_context import UnsupportedOptimizationContextError + + raise UnsupportedOptimizationContextError( + "GEPARootAgentPromptOptimizer does not support OptimizationRunContext" + " in this release; check optimizer.capabilities before passing a" + " context." + ) + if initial_agent.sub_agents: _logger.warning( "The GEPARootAgentPromptOptimizer will not optimize prompts for" diff --git a/src/google/adk/optimization/run_context.py b/src/google/adk/optimization/run_context.py index de0f30a3297..75aedbccfb3 100644 --- a/src/google/adk/optimization/run_context.py +++ b/src/google/adk/optimization/run_context.py @@ -20,15 +20,17 @@ An :class:`OptimizationRunContext` is an optional, one-shot, caller-owned object attached to a single ``AgentOptimizer.optimize(...)`` run. It records every *logical* optimizer-owned model invocation (control-plane metadata only, -never prompt or response content), enforces configured call/token budgets, and -carries a cooperative cancellation signal that instrumented optimizers observe -between logical model invocations. +never prompt or response content), enforces configured call/token budgets, +carries a cooperative cancellation signal, and owns an explicit terminal +state machine: every governed run finalizes exactly once as ``completed``, +``budget_exceeded``, ``cancelled``, or ``failed`` (first terminal transition +wins). The ledger records provider-reported usage truthfully: missing token counters -stay ``None`` and are classified via :class:`UsageCoverage`, never coerced to -zero. The final :class:`OptimizationRunSnapshot` is readable from the context -on success *and* failure, so a governance caller can persist the attempt in a -``finally`` block without parsing logs. +stay ``None`` and are classified via usage coverage, never coerced to zero. +Snapshots are immutable values; the final snapshot is readable from the +context on success *and* failure, so a governance caller can persist the +attempt in a ``finally`` block without parsing logs. """ from __future__ import annotations @@ -41,8 +43,15 @@ from typing import Optional from pydantic import BaseModel +from pydantic import ConfigDict from pydantic import Field +# Initially defined optimizer stages. The stage field is an extensible string +# so third-party optimizers can record their own stages without an ADK +# release; these constants cover the built-in optimizers. +STAGE_CANDIDATE_GENERATION = "candidate_generation" +STAGE_REFLECTION = "reflection" + class UsageCoverage(str, enum.Enum): """How completely the provider reported token usage for one logical call.""" @@ -57,13 +66,6 @@ class UsageCoverage(str, enum.Enum): """No token counter was supplied.""" -class ModelCallStage(str, enum.Enum): - """Which optimizer stage issued the logical model invocation.""" - - CANDIDATE_GENERATION = "candidate_generation" - REFLECTION = "reflection" - - class ModelCallState(str, enum.Enum): """Terminal state of one logical model invocation (call status). @@ -100,14 +102,18 @@ class TokenBudgetStatus(str, enum.Enum): class ModelCallEvent(BaseModel): - """Control-plane record of one logical model invocation. + """Immutable control-plane record of one logical model invocation. Prompt and response content are deliberately not part of this event; - content capture remains an explicit evaluator/sampler concern. + content capture remains an explicit evaluator/sampler concern. Error + metadata is structured and sanitized: a short provider error code and the + exception type name, never raw exception or payload text. """ + model_config = ConfigDict(frozen=True) + sequence: int - stage: ModelCallStage + stage: str requested_model: Optional[str] = None returned_model_version: Optional[str] = None start_time: float @@ -120,47 +126,33 @@ class ModelCallEvent(BaseModel): tool_use_tokens: Optional[int] = None total_tokens: Optional[int] = None usage_coverage: Optional[UsageCoverage] = None - error_message: Optional[str] = None - - -class OptimizationRunSnapshot(BaseModel): - """Immutable view of the run ledger at a point in time.""" + error_code: Optional[str] = None + error_type: Optional[str] = None - events: list[ModelCallEvent] = Field(default_factory=list) - started_calls: int = 0 - completed_calls: int = 0 - cumulative_total_tokens: int = 0 - """Sum of provider-reported authoritative totals (verified calls only).""" - usage_coverage: UsageCoverage = UsageCoverage.UNREPORTED - """VERIFIED only if every completed call was verified; PARTIAL if any call - reported any counter; UNREPORTED otherwise.""" - - token_budget_status: Optional[TokenBudgetStatus] = None - """Only set when a reported-token limit is configured.""" - - cancel_requested: bool = False - cancel_reason: Optional[str] = None - run_status: Optional[RunStatus] = None - """Terminal run status; ``None`` while the run is still in progress.""" +class OptimizationBudgets(BaseModel): + """Configured ceilings for optimizer-owned logical model invocations. + Immutable by construction, so limits cannot change mid-run. + """ -class OptimizationBudgets(BaseModel): - """Configured ceilings for optimizer-owned logical model invocations.""" + model_config = ConfigDict(frozen=True) max_model_calls: Optional[int] = Field( default=None, + ge=0, description=( "Maximum number of logical optimizer-owned model invocations that" - " may start. Checked before each call; a preflight rejection does" - " not create a model-call event." + " may start. Checked before each call via atomic slot admission; a" + " rejected reservation is not a call event." ), ) max_provider_reported_tokens: Optional[int] = Field( default=None, + ge=0, description=( "Ceiling on cumulative provider-reported authoritative total" - " tokens. Checked after each terminal response is committed; an" + " tokens. Checked after each terminal response is committed; the" " over-budget final call is committed first, then the run" " terminates. Deliberately not a hard bound on billing or physical" " tokens: when any call's usage is partial or unreported," @@ -171,14 +163,47 @@ class OptimizationBudgets(BaseModel): default="raise", description=( "Terminal behavior on budget exhaustion. Both modes stop scheduling" - " immediately after the final in-flight call settles. 'raise' raises" - " OptimizationBudgetExceeded; 'return_partial' lets the optimizer" - " return its best-so-far result marked terminal_status=" - "'budget_exceeded' (never an unmarked success)." + " immediately after the final in-flight call settles. 'raise'" + " raises OptimizationBudgetExceeded; 'return_partial' lets the" + " optimizer return its best-so-far result while the caller-owned" + " snapshot (run_status='budget_exceeded') remains the authoritative" + " record." ), ) +class OptimizationRunSnapshot(BaseModel): + """Immutable view of the run ledger at a point in time.""" + + model_config = ConfigDict(frozen=True) + + events: tuple[ModelCallEvent, ...] = () + budgets: OptimizationBudgets = Field(default_factory=OptimizationBudgets) + started_calls: int = 0 + completed_calls: int = 0 + cumulative_total_tokens: int = 0 + """Sum of provider-reported authoritative totals (verified calls only).""" + + usage_coverage: UsageCoverage = UsageCoverage.UNREPORTED + """VERIFIED only if every closed call was verified; PARTIAL if any call + reported any counter; UNREPORTED otherwise.""" + + token_budget_status: Optional[TokenBudgetStatus] = None + """Only set when a reported-token limit is configured.""" + + cancel_requested: bool = False + cancel_reason: Optional[str] = None + run_status: Optional[RunStatus] = None + """Terminal run status; ``None`` only while the run is still in progress.""" + + terminal_sequence: Optional[int] = None + """Sequence of the logical invocation that triggered the terminal state, + when one did.""" + + terminal_error_code: Optional[str] = None + terminal_error_type: Optional[str] = None + + class OptimizationRunContextError(Exception): """Base class for run-context errors carrying the final snapshot.""" @@ -196,13 +221,19 @@ class OptimizationCancelledError(OptimizationRunContextError): class OptimizationProviderError(OptimizationRunContextError): - """A governed run terminated on an in-band provider error. + """A governed run terminated on a provider failure. - Usage reported before the error is preserved on the snapshot; a governed - run must not silently succeed past an in-band ``LlmResponse.error_code``. + Usage reported before the failure is preserved on the snapshot. Provider + failure takes precedence over a simultaneous token overshoot: both the + committed usage and the token-compliance evidence are preserved, but the + run terminates ``failed`` and this error is raised. """ +class OptimizationRunFinalizedError(Exception): + """A terminal context was used where an in-progress context is required.""" + + class UnsupportedOptimizationContextError(Exception): """The optimizer does not support run-context instrumentation.""" @@ -211,12 +242,50 @@ class ContextAlreadyAttachedError(Exception): """A one-shot context was attached to more than one optimization run.""" +class _CallRecord: + """Mutable internal record for one logical model invocation.""" + + __slots__ = ( + "sequence", + "stage", + "requested_model", + "returned_model_version", + "start_time", + "end_time", + "state", + "usage", + "error_code", + "error_type", + "closed", + ) + + def __init__(self, sequence: int, stage: str, requested_model): + self.sequence = sequence + self.stage = stage + self.requested_model = requested_model + self.returned_model_version = None + self.start_time = time.monotonic() + self.end_time = None + self.state: Optional[ModelCallState] = None + self.usage: dict[str, Optional[int]] = {} + self.error_code: Optional[str] = None + self.error_type: Optional[str] = None + self.closed = False + + class _CallHandle: - """Opaque handle for one in-flight logical model invocation.""" + """Opaque handle for one in-flight logical model invocation. - def __init__(self, event: ModelCallEvent): - self._event = event - self._closed = False + Bound to exactly one context; committing it into a different context is a + typed misuse error. + """ + + def __init__(self, context: "OptimizationRunContext", record: _CallRecord): + self._context = context + self._record = record + + +_UNATTACHED = object() class OptimizationRunContext: @@ -224,40 +293,57 @@ class OptimizationRunContext: Instrumented optimizers call :meth:`begin_model_call` immediately before a logical ``BaseLlm.generate_content_async()`` invocation and - :meth:`end_model_call` with the terminal response, in-band error, or raised - provider error. Governance callers own the context, may call - :meth:`request_cancel` from any thread, and read :meth:`snapshot` at any - time, including after a failure. + :meth:`end_model_call` with the terminal outcome. Every governed run must + finalize exactly once (first terminal transition wins): optimizers call + :meth:`finalize_success` on a normal return; budget/cancel/provider + terminals are committed by the corresponding ledger operations or by + :meth:`finalize_cancelled` / :meth:`finalize_failed` on exceptional exits. + Governance callers own the context, may call :meth:`request_cancel` from + any thread, and read :meth:`snapshot` at any time, including after failure. """ def __init__(self, budgets: Optional[OptimizationBudgets] = None): + # OptimizationBudgets is frozen; keep a private reference and never + # expose a mutable path to it. self._budgets = budgets or OptimizationBudgets() self._lock = threading.Lock() - self._events: list[ModelCallEvent] = [] + self._records: list[_CallRecord] = [] self._started_calls = 0 self._completed_calls = 0 self._cumulative_total_tokens = 0 self._cancel_requested = False self._cancel_reason: Optional[str] = None self._run_status: Optional[RunStatus] = None - self._attached_owner: Optional[object] = None + self._terminal_sequence: Optional[int] = None + self._terminal_error_code: Optional[str] = None + self._terminal_error_type: Optional[str] = None + self._attached_owner: object = _UNATTACHED @property def budgets(self) -> OptimizationBudgets: + """The configured (immutable) budgets.""" return self._budgets def attach(self, owner: object) -> None: """Binds the context to one optimization run; reuse is rejected.""" with self._lock: - if self._attached_owner is not None: + if self._attached_owner is not _UNATTACHED: raise ContextAlreadyAttachedError( "OptimizationRunContext is one-shot: it is already attached to an" " optimization run and cannot be reused or shared." ) self._attached_owner = owner + # --- cancellation --------------------------------------------------------- + def request_cancel(self, reason: str = "requested") -> None: - """Thread-safe, idempotent cooperative cancellation signal.""" + """Thread-safe, idempotent cooperative cancellation signal. + + Requesting cancellation does not itself finalize the run; the terminal + ``cancelled`` transition is committed when the signal is observed at a + documented boundary (or via :meth:`finalize_cancelled`), and never + overwrites an earlier terminal state. + """ with self._lock: if not self._cancel_requested: self._cancel_requested = True @@ -269,55 +355,137 @@ def cancel_requested(self) -> bool: return self._cancel_requested def raise_if_cancelled(self) -> None: - """Raises ``OptimizationCancelledError`` if cancellation was requested.""" + """Raises ``OptimizationCancelledError`` if cancellation was requested. + + First terminal wins: if the run already terminated for another reason, + this re-raises that terminal outcome instead of overwriting it. + """ with self._lock: - if not self._cancel_requested: + if self._run_status is not None: + error = self._terminal_error_locked() + elif self._cancel_requested: + self._transition_locked(RunStatus.CANCELLED) + error = OptimizationCancelledError( + f"Optimization cancelled: {self._cancel_reason}", + self._snapshot_locked(), + ) + else: return - self._run_status = RunStatus.CANCELLED - snapshot = self._snapshot_locked() - reason = self._cancel_reason - raise OptimizationCancelledError( - f"Optimization cancelled: {reason}", snapshot + raise error + + # --- terminal transitions ------------------------------------------------- + + def finalize_success(self) -> None: + """Marks a normal optimizer return ``completed`` (first terminal wins).""" + with self._lock: + self._transition_locked(RunStatus.COMPLETED) + + def finalize_cancelled(self, reason: str = "cancelled") -> None: + """Commits the ``cancelled`` terminal (e.g. on native task cancellation). + + Closes any open call event as ``cancelled`` first, so the final snapshot + carries no open invocations. First terminal wins. + """ + with self._lock: + if not self._cancel_requested: + self._cancel_requested = True + self._cancel_reason = reason + for record in self._records: + if not record.closed: + record.closed = True + record.end_time = time.monotonic() + record.state = ModelCallState.CANCELLED + self._transition_locked(RunStatus.CANCELLED) + + def finalize_failed( + self, + *, + error_code: Optional[str] = None, + error_type: Optional[str] = None, + ) -> None: + """Commits the ``failed`` terminal for a non-provider-call failure path. + + First terminal wins. Error metadata must be sanitized identifiers (a + short code and an exception type name), never raw payload text. + """ + with self._lock: + self._transition_locked( + RunStatus.FAILED, error_code=error_code, error_type=error_type + ) + + def _transition_locked( + self, + status: RunStatus, + *, + sequence: Optional[int] = None, + error_code: Optional[str] = None, + error_type: Optional[str] = None, + ) -> bool: + """First-terminal-wins transition; returns True if this call won.""" + if self._run_status is not None: + return False + self._run_status = status + self._terminal_sequence = sequence + self._terminal_error_code = error_code + self._terminal_error_type = error_type + return True + + def _terminal_error_locked(self) -> Exception: + """The typed error corresponding to an existing terminal state.""" + snapshot = self._snapshot_locked() + if self._run_status == RunStatus.BUDGET_EXCEEDED: + return OptimizationBudgetExceeded( + "Optimization budget exhausted.", snapshot + ) + if self._run_status == RunStatus.CANCELLED: + return OptimizationCancelledError( + f"Optimization cancelled: {self._cancel_reason}", snapshot + ) + if self._run_status == RunStatus.FAILED: + return OptimizationProviderError( + f"Optimization failed: {self._terminal_error_code}", snapshot + ) + return OptimizationRunFinalizedError( + f"OptimizationRunContext already finalized as {self._run_status.value}." ) + # --- the ledger ----------------------------------------------------------- + def begin_model_call( self, - stage: ModelCallStage, + stage: str, requested_model: Optional[str] = None, ) -> _CallHandle: - """Reserves and records the start of one logical model invocation. + """Atomically admits and records the start of one logical invocation. - Checks cancellation and the logical-call budget *before* the invocation - starts. A preflight rejection does not create a model-call event and does - not increment the started-call count; it records the run-level control - state and raises. + Under one lock acquisition: rejects a terminal context, observes + cancellation, checks the logical-call budget, reserves the slot, assigns + the sequence ID, and creates the ledger entry. A preflight rejection is + not a call event and does not increment the started-call count. """ with self._lock: - if self._cancel_requested: - self._run_status = RunStatus.CANCELLED - snapshot = self._snapshot_locked() - reason = self._cancel_reason - raise OptimizationCancelledError( - f"Optimization cancelled: {reason}", snapshot - ) - max_calls = self._budgets.max_model_calls - if max_calls is not None and self._started_calls >= max_calls: - # Atomic slot admission: the rejected reservation is not a call - # event; the run status records why the run ended. - self._run_status = RunStatus.BUDGET_EXCEEDED - snapshot = self._snapshot_locked() - raise OptimizationBudgetExceeded( - f"Logical model-call budget exhausted ({max_calls}).", snapshot + if self._run_status is not None: + error = self._terminal_error_locked() + elif self._cancel_requested: + self._transition_locked(RunStatus.CANCELLED) + error = OptimizationCancelledError( + f"Optimization cancelled: {self._cancel_reason}", + self._snapshot_locked(), ) - self._started_calls += 1 - event = ModelCallEvent( - sequence=self._started_calls, - stage=stage, - requested_model=requested_model, - start_time=time.monotonic(), - ) - self._events.append(event) - return _CallHandle(event) + else: + max_calls = self._budgets.max_model_calls + if max_calls is not None and self._started_calls >= max_calls: + self._transition_locked(RunStatus.BUDGET_EXCEEDED) + error = OptimizationBudgetExceeded( + f"Logical model-call budget exhausted ({max_calls}).", + self._snapshot_locked(), + ) + else: + self._started_calls += 1 + record = _CallRecord(self._started_calls, stage, requested_model) + self._records.append(record) + return _CallHandle(self, record) + raise error def end_model_call( self, @@ -325,62 +493,109 @@ def end_model_call( *, usage_metadata: Any = None, returned_model_version: Optional[str] = None, - error_message: Optional[str] = None, + error_code: Optional[str] = None, + error_type: Optional[str] = None, ) -> None: """Commits the terminal outcome of one logical model invocation. - Token usage is committed before budget enforcement: if the new cumulative - total crosses the hard ceiling, the completed call is persisted first and - ``OptimizationBudgetExceeded`` is raised immediately afterwards, so an - over-budget final call can never produce an unmarked success. A terminal - provider error closes the call as ``provider_error``, preserves any - reported usage, and re-raising the primary error remains the caller's - responsibility. + Ownership and close are validated atomically under the context lock: a + handle from another context is a typed misuse error, and a concurrent + double-close commits exactly once. + + Provider failure takes precedence over token overshoot: usage is + committed either way, but a call ending in a provider error transitions + the run to ``failed`` and raises ``OptimizationProviderError``. A clean + over-budget final call is committed with call status ``completed`` while + the run transitions to ``budget_exceeded`` and + ``OptimizationBudgetExceeded`` is raised. Error metadata must be + sanitized identifiers, never raw exception/payload text. """ - if handle._closed: - return - handle._closed = True - event = handle._event + record = handle._record with self._lock: - event.end_time = time.monotonic() - event.returned_model_version = returned_model_version - _apply_usage(event, usage_metadata) - if error_message is not None: - event.state = ModelCallState.PROVIDER_ERROR - event.error_message = error_message - else: - event.state = ModelCallState.COMPLETED + if handle._context is not self: + raise OptimizationRunFinalizedError( + "Call handle belongs to a different OptimizationRunContext." + ) + if record.closed: + return + record.closed = True + record.end_time = time.monotonic() + record.returned_model_version = returned_model_version + record.usage = _extract_usage(usage_metadata) self._completed_calls += 1 - if event.total_tokens is not None: - self._cumulative_total_tokens += event.total_tokens - max_tokens = self._budgets.max_provider_reported_tokens - over_budget = ( - max_tokens is not None - and self._cumulative_total_tokens > max_tokens - ) - if over_budget: - self._run_status = RunStatus.BUDGET_EXCEEDED - snapshot = self._snapshot_locked() - if over_budget: - raise OptimizationBudgetExceeded( - f"Token budget exhausted ({max_tokens}).", snapshot - ) + total = record.usage.get("total_tokens") + if total is not None: + self._cumulative_total_tokens += total + + error: Optional[Exception] = None + if error_code is not None or error_type is not None: + record.state = ModelCallState.PROVIDER_ERROR + record.error_code = error_code + record.error_type = error_type + # Provider failure is primary even when the same terminal response + # also crossed the token ceiling; both usage and token-compliance + # evidence are preserved on the snapshot. + if self._transition_locked( + RunStatus.FAILED, + sequence=record.sequence, + error_code=error_code, + error_type=error_type, + ): + error = OptimizationProviderError( + f"Provider failure on call {record.sequence}: {error_code}", + self._snapshot_locked(), + ) + else: + record.state = ModelCallState.COMPLETED + max_tokens = self._budgets.max_provider_reported_tokens + if ( + max_tokens is not None + and self._cumulative_total_tokens > max_tokens + ): + if self._transition_locked( + RunStatus.BUDGET_EXCEEDED, sequence=record.sequence + ): + error = OptimizationBudgetExceeded( + f"Token budget exhausted ({max_tokens}).", + self._snapshot_locked(), + ) + if error is not None: + raise error + + # --- snapshots -------------------------------------------------------------- def snapshot(self) -> OptimizationRunSnapshot: with self._lock: return self._snapshot_locked() def _snapshot_locked(self) -> OptimizationRunSnapshot: - completed = [ - e for e in self._events if e.state is not None - ] - if completed and all( - e.usage_coverage == UsageCoverage.VERIFIED for e in completed - ): + events = tuple( + ModelCallEvent( + sequence=r.sequence, + stage=r.stage, + requested_model=r.requested_model, + returned_model_version=r.returned_model_version, + start_time=r.start_time, + end_time=r.end_time, + state=r.state, + prompt_tokens=r.usage.get("prompt_tokens"), + output_tokens=r.usage.get("output_tokens"), + reasoning_tokens=r.usage.get("reasoning_tokens"), + cached_tokens=r.usage.get("cached_tokens"), + tool_use_tokens=r.usage.get("tool_use_tokens"), + total_tokens=r.usage.get("total_tokens"), + usage_coverage=_coverage_of(r) if r.closed else None, + error_code=r.error_code, + error_type=r.error_type, + ) + for r in self._records + ) + closed = [r for r in self._records if r.closed] + coverages = [_coverage_of(r) for r in closed] + if coverages and all(c == UsageCoverage.VERIFIED for c in coverages): run_coverage = UsageCoverage.VERIFIED elif any( - e.usage_coverage in (UsageCoverage.VERIFIED, UsageCoverage.PARTIAL) - for e in completed + c in (UsageCoverage.VERIFIED, UsageCoverage.PARTIAL) for c in coverages ): run_coverage = UsageCoverage.PARTIAL else: @@ -392,13 +607,14 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: > self._budgets.max_provider_reported_tokens ): token_status = TokenBudgetStatus.EXCEEDED - elif run_coverage == UsageCoverage.VERIFIED: + elif not closed or run_coverage == UsageCoverage.VERIFIED: token_status = TokenBudgetStatus.WITHIN_LIMIT else: # Missing totals are never proof of compliance. token_status = TokenBudgetStatus.INDETERMINATE return OptimizationRunSnapshot( - events=[e.model_copy(deep=True) for e in self._events], + events=events, + budgets=self._budgets, started_calls=self._started_calls, completed_calls=self._completed_calls, cumulative_total_tokens=self._cumulative_total_tokens, @@ -407,45 +623,37 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: cancel_requested=self._cancel_requested, cancel_reason=self._cancel_reason, run_status=self._run_status, + terminal_sequence=self._terminal_sequence, + terminal_error_code=self._terminal_error_code, + terminal_error_type=self._terminal_error_type, ) -def _apply_usage(event: ModelCallEvent, usage_metadata: Any) -> None: - """Copies provider-reported token counters onto the event, truthfully. - - Missing counters stay ``None``. Coverage: ``verified`` iff the provider - supplied ``total_token_count``; ``partial`` if any other counter was - supplied; ``unreported`` otherwise. - """ +def _extract_usage(usage_metadata: Any) -> dict[str, Optional[int]]: + """Copies provider-reported token counters, truthfully (no zero-coercion).""" if usage_metadata is None: - event.usage_coverage = UsageCoverage.UNREPORTED - return + return {} def _get(name: str) -> Optional[int]: value = getattr(usage_metadata, name, None) return int(value) if isinstance(value, (int, float)) else None - event.prompt_tokens = _get("prompt_token_count") - event.output_tokens = _get("candidates_token_count") - event.reasoning_tokens = _get("thoughts_token_count") - event.cached_tokens = _get("cached_content_token_count") - event.tool_use_tokens = _get("tool_use_prompt_token_count") - event.total_tokens = _get("total_token_count") - if event.total_tokens is not None: - event.usage_coverage = UsageCoverage.VERIFIED - elif any( - v is not None - for v in ( - event.prompt_tokens, - event.output_tokens, - event.reasoning_tokens, - event.cached_tokens, - event.tool_use_tokens, - ) - ): - event.usage_coverage = UsageCoverage.PARTIAL - else: - event.usage_coverage = UsageCoverage.UNREPORTED + return { + "prompt_tokens": _get("prompt_token_count"), + "output_tokens": _get("candidates_token_count"), + "reasoning_tokens": _get("thoughts_token_count"), + "cached_tokens": _get("cached_content_token_count"), + "tool_use_tokens": _get("tool_use_prompt_token_count"), + "total_tokens": _get("total_token_count"), + } + + +def _coverage_of(record: _CallRecord) -> UsageCoverage: + if record.usage.get("total_tokens") is not None: + return UsageCoverage.VERIFIED + if any(v is not None for v in record.usage.values()): + return UsageCoverage.PARTIAL + return UsageCoverage.UNREPORTED class OptimizerCapabilities(BaseModel): @@ -458,6 +666,8 @@ class OptimizerCapabilities(BaseModel): preflight instead of discovering opacity after spend occurs. """ + model_config = ConfigDict(frozen=True) + accepts_run_context: bool = False """``optimize`` accepts a run context at all. When ``False``, callers must omit the keyword entirely (protects pre-existing third-party overrides).""" diff --git a/src/google/adk/optimization/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index 6a1c7398179..20f2a26d8b0 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -18,6 +18,7 @@ import logging import random +from typing import Optional from google.adk.agents.llm_agent import Agent from google.adk.evaluation._retry_options_utils import add_default_retry_options_if_not_present @@ -203,7 +204,18 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context=None, ) -> OptimizerResult[AgentWithScores]: + if run_context is not None: + from .run_context import UnsupportedOptimizationContextError + + raise UnsupportedOptimizationContextError( + "SimplePromptOptimizer does not support OptimizationRunContext in" + " this release; check optimizer.capabilities before passing a" + " context." + ) + train_example_ids = sampler.get_train_example_ids() if self._config.batch_size > len(train_example_ids): diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index c461c81421b..f3676581429 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -18,16 +18,20 @@ from types import SimpleNamespace from google.adk.optimization.run_context import ContextAlreadyAttachedError -from google.adk.optimization.run_context import ModelCallStage from google.adk.optimization.run_context import ModelCallState from google.adk.optimization.run_context import OptimizationBudgetExceeded from google.adk.optimization.run_context import OptimizationBudgets from google.adk.optimization.run_context import OptimizationCancelledError +from google.adk.optimization.run_context import OptimizationProviderError from google.adk.optimization.run_context import OptimizationRunContext +from google.adk.optimization.run_context import OptimizationRunFinalizedError from google.adk.optimization.run_context import OptimizerCapabilities from google.adk.optimization.run_context import RunStatus +from google.adk.optimization.run_context import STAGE_CANDIDATE_GENERATION +from google.adk.optimization.run_context import STAGE_REFLECTION from google.adk.optimization.run_context import TokenBudgetStatus from google.adk.optimization.run_context import UsageCoverage +import pydantic import pytest @@ -43,25 +47,157 @@ def test_attach_twice_rejected(self): with pytest.raises(ContextAlreadyAttachedError): ctx.attach(owner=object()) + def test_attach_none_owner_still_one_shot(self): + # None must not double as the unattached sentinel. + ctx = OptimizationRunContext() + ctx.attach(owner=None) + with pytest.raises(ContextAlreadyAttachedError): + ctx.attach(owner=None) + def test_two_runs_have_isolated_contexts(self): a, b = OptimizationRunContext(), OptimizationRunContext() a.attach(owner="run-a") b.attach(owner="run-b") - h = a.begin_model_call(ModelCallStage.REFLECTION) + h = a.begin_model_call(STAGE_REFLECTION) a.end_model_call(h, usage_metadata=_usage(total_token_count=7)) assert a.snapshot().completed_calls == 1 assert b.snapshot().completed_calls == 0 +class TestLedgerIntegrity: + + def test_foreign_handle_rejected(self): + a, b = OptimizationRunContext(), OptimizationRunContext() + handle_from_a = a.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationRunFinalizedError): + b.end_model_call(handle_from_a, usage_metadata=None) + # B's ledger is untouched by the attempt. + assert b.snapshot().completed_calls == 0 + assert b.snapshot().started_calls == 0 + + def test_concurrent_double_close_commits_once(self): + ctx = OptimizationRunContext() + handle = ctx.begin_model_call(STAGE_REFLECTION) + errors: list[Exception] = [] + + def close(): + try: + ctx.end_model_call(handle, usage_metadata=_usage(total_token_count=10)) + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=close) for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + assert not errors + snap = ctx.snapshot() + assert snap.completed_calls == 1 + assert snap.cumulative_total_tokens == 10 + + def test_budgets_are_immutable(self): + budgets = OptimizationBudgets(max_model_calls=3) + with pytest.raises(pydantic.ValidationError): + budgets.max_model_calls = 100 + ctx = OptimizationRunContext(budgets) + with pytest.raises(pydantic.ValidationError): + ctx.budgets.max_model_calls = 100 + + def test_negative_budgets_rejected(self): + with pytest.raises(pydantic.ValidationError): + OptimizationBudgets(max_model_calls=-1) + with pytest.raises(pydantic.ValidationError): + OptimizationBudgets(max_provider_reported_tokens=-5) + + def test_snapshot_and_events_are_immutable(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=1)) + snap = ctx.snapshot() + with pytest.raises(pydantic.ValidationError): + snap.started_calls = 99 + with pytest.raises(pydantic.ValidationError): + snap.events[0].total_tokens = 99 + assert isinstance(snap.events, tuple) + + def test_extensible_stage_strings(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call("my_custom_stage") + ctx.end_model_call(h, usage_metadata=None) + assert ctx.snapshot().events[0].stage == "my_custom_stage" + assert STAGE_CANDIDATE_GENERATION == "candidate_generation" + + +class TestTerminalStateMachine: + + def test_success_finalizer(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=1)) + ctx.finalize_success() + assert ctx.snapshot().run_status == RunStatus.COMPLETED + + def test_first_terminal_wins_late_cancel_does_not_overwrite(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + ctx.request_cancel("late") + with pytest.raises(OptimizationBudgetExceeded): + # Re-raises the EARLIER terminal outcome, not cancellation. + ctx.raise_if_cancelled() + assert ctx.snapshot().run_status == RunStatus.BUDGET_EXCEEDED + + def test_no_call_admitted_after_terminal(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + with pytest.raises(OptimizationBudgetExceeded): + ctx.begin_model_call(STAGE_REFLECTION) + assert ctx.snapshot().started_calls == 1 + + def test_no_call_admitted_after_completed(self): + ctx = OptimizationRunContext() + ctx.finalize_success() + with pytest.raises(OptimizationRunFinalizedError): + ctx.begin_model_call(STAGE_REFLECTION) + + def test_finalize_cancelled_closes_open_events(self): + ctx = OptimizationRunContext() + ctx.begin_model_call(STAGE_REFLECTION) # left open (native cancel) + ctx.finalize_cancelled("task_cancelled") + snap = ctx.snapshot() + assert snap.run_status == RunStatus.CANCELLED + assert snap.events[0].state == ModelCallState.CANCELLED + assert snap.events[0].end_time is not None + + def test_snapshot_carries_limits_and_terminal_metadata(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + snap = ctx.snapshot() + assert snap.budgets.max_provider_reported_tokens == 10 + assert snap.terminal_sequence == 1 + + class TestCallBudget: def test_exactly_n_calls_may_start(self): ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=2)) for _ in range(2): - h = ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + h = ctx.begin_model_call(STAGE_CANDIDATE_GENERATION) ctx.end_model_call(h, usage_metadata=None) with pytest.raises(OptimizationBudgetExceeded) as exc: - ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + ctx.begin_model_call(STAGE_CANDIDATE_GENERATION) snap = exc.value.snapshot # The rejected reservation is not a call event and did not start a call. assert snap.started_calls == 2 @@ -75,7 +211,7 @@ def test_overshoot_commits_then_raises(self): ctx = OptimizationRunContext( OptimizationBudgets(max_provider_reported_tokens=100) ) - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) with pytest.raises(OptimizationBudgetExceeded) as exc: ctx.end_model_call(h, usage_metadata=_usage(total_token_count=150)) snap = exc.value.snapshot @@ -91,7 +227,7 @@ def test_unreported_usage_is_indeterminate_not_compliant(self): ctx = OptimizationRunContext( OptimizationBudgets(max_provider_reported_tokens=10) ) - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) ctx.end_model_call(h, usage_metadata=None) # no raise snap = ctx.snapshot() assert snap.cumulative_total_tokens == 0 @@ -102,7 +238,7 @@ def test_all_verified_under_limit_is_within_limit(self): ctx = OptimizationRunContext( OptimizationBudgets(max_provider_reported_tokens=100) ) - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) ctx.end_model_call(h, usage_metadata=_usage(total_token_count=40)) assert ctx.snapshot().token_budget_status == TokenBudgetStatus.WITHIN_LIMIT @@ -111,13 +247,13 @@ class TestUsageClassification: def test_verified_partial_unreported(self): ctx = OptimizationRunContext() - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) ctx.end_model_call( h, usage_metadata=_usage(prompt_token_count=5, total_token_count=9) ) - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) ctx.end_model_call(h, usage_metadata=_usage(prompt_token_count=5)) - h = ctx.begin_model_call(ModelCallStage.REFLECTION) + h = ctx.begin_model_call(STAGE_REFLECTION) ctx.end_model_call(h, usage_metadata=None) events = ctx.snapshot().events assert events[0].usage_coverage == UsageCoverage.VERIFIED @@ -130,25 +266,55 @@ def test_verified_partial_unreported(self): def test_run_coverage_verified_only_when_all_verified(self): ctx = OptimizationRunContext() for _ in range(2): - h = ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + h = ctx.begin_model_call(STAGE_CANDIDATE_GENERATION) ctx.end_model_call(h, usage_metadata=_usage(total_token_count=3)) assert ctx.snapshot().usage_coverage == UsageCoverage.VERIFIED -class TestProviderError: +class TestProviderFailure: - def test_error_preserves_usage_so_far(self): + def test_provider_error_is_a_governed_terminal(self): ctx = OptimizationRunContext() - h = ctx.begin_model_call(ModelCallStage.REFLECTION) - ctx.end_model_call( - h, - usage_metadata=_usage(total_token_count=42), - error_message="RESOURCE_EXHAUSTED", - ) - event = ctx.snapshot().events[0] + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call( + h, + usage_metadata=_usage(total_token_count=42), + error_code="RESOURCE_EXHAUSTED", + error_type="ClientError", + ) + snap = exc.value.snapshot + event = snap.events[0] assert event.state == ModelCallState.PROVIDER_ERROR - assert event.total_tokens == 42 - assert event.error_message == "RESOURCE_EXHAUSTED" + assert event.total_tokens == 42 # usage preserved + assert event.error_code == "RESOURCE_EXHAUSTED" + assert event.error_type == "ClientError" + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "RESOURCE_EXHAUSTED" + + def test_provider_error_precedence_over_token_overshoot(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call( + h, + usage_metadata=_usage(total_token_count=50), + error_code="INTERNAL", + ) + snap = exc.value.snapshot + # Provider failure is primary; usage and compliance evidence preserved. + assert snap.run_status == RunStatus.FAILED + assert snap.cumulative_total_tokens == 50 + assert snap.token_budget_status == TokenBudgetStatus.EXCEEDED + + def test_finalize_failed_for_non_call_failures(self): + ctx = OptimizationRunContext() + ctx.finalize_failed(error_code="ADAPTER_CRASH", error_type="ValueError") + snap = ctx.snapshot() + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_type == "ValueError" class TestCancellation: @@ -158,7 +324,7 @@ def test_cancel_is_idempotent_and_first_reason_wins(self): ctx.request_cancel("deadline") ctx.request_cancel("other") with pytest.raises(OptimizationCancelledError) as exc: - ctx.begin_model_call(ModelCallStage.REFLECTION) + ctx.begin_model_call(STAGE_REFLECTION) assert "deadline" in str(exc.value) assert exc.value.snapshot.cancel_reason == "deadline" assert exc.value.snapshot.run_status == RunStatus.CANCELLED @@ -177,7 +343,7 @@ def test_cancel_from_another_thread_observed(self): def test_snapshot_readable_after_failure(self): ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=0)) with pytest.raises(OptimizationBudgetExceeded): - ctx.begin_model_call(ModelCallStage.CANDIDATE_GENERATION) + ctx.begin_model_call(STAGE_CANDIDATE_GENERATION) # A governance caller can persist the attempt from a finally block. assert ctx.snapshot().run_status == RunStatus.BUDGET_EXCEEDED @@ -212,3 +378,27 @@ def test_optimizer_result_has_no_run_context_fields(self): from google.adk.optimization.data_types import OptimizerResult assert "terminal_status" not in OptimizerResult.model_fields + + +class TestBuiltInsRejectContextInContractsPr: + + @pytest.mark.asyncio + async def test_simple_prompt_optimizer_rejects_context(self): + from unittest import mock + + from google.adk.optimization.run_context import UnsupportedOptimizationContextError + from google.adk.optimization.sampler import Sampler + from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer + from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig + + with mock.patch( + "google.adk.optimization.simple_prompt_optimizer.LLMRegistry.resolve" + ): + optimizer = SimplePromptOptimizer(SimplePromptOptimizerConfig()) + sampler = mock.MagicMock(spec=Sampler) + with pytest.raises(UnsupportedOptimizationContextError): + await optimizer.optimize( + mock.MagicMock(), sampler, run_context=OptimizationRunContext() + ) + # Rejected before any sampler work. + sampler.get_train_example_ids.assert_not_called() From 9ed15ac7f91763ccf888ddf10de87e00bc27e3a0 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 00:46:51 -0700 Subject: [PATCH 4/9] atomic settlement, complete terminal observation, truthful failure types (round-2 adversarial findings) - Validate-before-commit: usage counters must be finite non-negative numbers (NaN/inf/negative -> not reported, never a crash on a half-closed record); error metadata is normalized to bounded single-token identifiers at the context boundary (numeric provider codes become strings; multiline/path text is scrubbed); the event, counters, and terminal transition then commit atomically under one lock acquisition. - Every caller observes the committed terminal: an already-admitted call settles for truthful accounting and then receives the existing terminal error; finalize_success observes pending cancellation and commits/raises CANCELLED instead of recording a false success. - OptimizationFailedError added for non-provider failures; OptimizationProviderError is reserved for governed provider-call terminals (terminal provenance tracked). - end_model_call(cancelled=True) settles a call as cancelled while preserving usage evidence; completed_calls documented and counted as settled calls regardless of call status (finalize_cancelled counts its closures). - Event timestamps are wall-clock epoch seconds for durable, cross-worker attempt records. - Built-in optimize() signatures fully typed (Optional[OptimizationRunContext]). - 10 new regression tests reproduce every round-2 finding. Suite: 70/70. --- .../optimization/gepa_root_agent_optimizer.py | 4 +- .../gepa_root_agent_prompt_optimizer.py | 3 +- src/google/adk/optimization/run_context.py | 148 +++++++++++++++--- .../optimization/simple_prompt_optimizer.py | 3 +- .../optimization/run_context_test.py | 124 +++++++++++++++ 5 files changed, 253 insertions(+), 29 deletions(-) diff --git a/src/google/adk/optimization/gepa_root_agent_optimizer.py b/src/google/adk/optimization/gepa_root_agent_optimizer.py index a062860eda7..fc08ed5cb12 100644 --- a/src/google/adk/optimization/gepa_root_agent_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_optimizer.py @@ -19,6 +19,7 @@ import logging from typing import Any from typing import Callable +from typing import Optional from google.genai import types as genai_types from pydantic import BaseModel @@ -36,6 +37,7 @@ from .data_types import AgentWithScores from .data_types import OptimizerResult from .data_types import UnstructuredSamplingResult +from .run_context import OptimizationRunContext from .sampler import Sampler logger = logging.getLogger("google_adk." + __name__) @@ -334,7 +336,7 @@ async def optimize( initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], *, - run_context=None, + run_context: Optional[OptimizationRunContext] = None, ) -> GEPARootAgentOptimizerResult: """Runs the GEPARootAgentOptimizer. diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py index 5adebfb140a..de78e709fce 100644 --- a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -35,6 +35,7 @@ from .data_types import AgentWithScores from .data_types import OptimizerResult from .data_types import UnstructuredSamplingResult +from .run_context import OptimizationRunContext from .sampler import Sampler _logger = logging.getLogger("google_adk." + __name__) @@ -208,7 +209,7 @@ async def optimize( initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], *, - run_context=None, + run_context: Optional[OptimizationRunContext] = None, ) -> GEPARootAgentPromptOptimizerResult: """Runs the GEPARootAgentPromptOptimizer. diff --git a/src/google/adk/optimization/run_context.py b/src/google/adk/optimization/run_context.py index 75aedbccfb3..4bcb3cd7cf6 100644 --- a/src/google/adk/optimization/run_context.py +++ b/src/google/adk/optimization/run_context.py @@ -36,6 +36,8 @@ from __future__ import annotations import enum +import math +import re import threading import time from typing import Any @@ -117,6 +119,8 @@ class ModelCallEvent(BaseModel): requested_model: Optional[str] = None returned_model_version: Optional[str] = None start_time: float + """Wall-clock (epoch seconds): durable and comparable across workers.""" + end_time: Optional[float] = None state: Optional[ModelCallState] = None prompt_tokens: Optional[int] = None @@ -181,6 +185,9 @@ class OptimizationRunSnapshot(BaseModel): budgets: OptimizationBudgets = Field(default_factory=OptimizationBudgets) started_calls: int = 0 completed_calls: int = 0 + """Count of settled (closed) calls, whatever their terminal call status -- + completed, provider_error, or cancelled.""" + cumulative_total_tokens: int = 0 """Sum of provider-reported authoritative totals (verified calls only).""" @@ -221,7 +228,7 @@ class OptimizationCancelledError(OptimizationRunContextError): class OptimizationProviderError(OptimizationRunContextError): - """A governed run terminated on a provider failure. + """A governed run terminated on a provider-call failure. Usage reported before the failure is preserved on the snapshot. Provider failure takes precedence over a simultaneous token overshoot: both the @@ -230,6 +237,16 @@ class OptimizationProviderError(OptimizationRunContextError): """ +class OptimizationFailedError(OptimizationRunContextError): + """A governed run terminated on a non-provider failure. + + Sampler, adapter, validation, and optimizer-internal failures recorded via + ``finalize_failed`` surface as this generic type; + ``OptimizationProviderError`` is reserved for failures of governed + provider calls. + """ + + class OptimizationRunFinalizedError(Exception): """A terminal context was used where an in-progress context is required.""" @@ -264,7 +281,7 @@ def __init__(self, sequence: int, stage: str, requested_model): self.stage = stage self.requested_model = requested_model self.returned_model_version = None - self.start_time = time.monotonic() + self.start_time = time.time() self.end_time = None self.state: Optional[ModelCallState] = None self.usage: dict[str, Optional[int]] = {} @@ -309,7 +326,7 @@ def __init__(self, budgets: Optional[OptimizationBudgets] = None): self._lock = threading.Lock() self._records: list[_CallRecord] = [] self._started_calls = 0 - self._completed_calls = 0 + self._settled_calls = 0 self._cumulative_total_tokens = 0 self._cancel_requested = False self._cancel_reason: Optional[str] = None @@ -317,6 +334,7 @@ def __init__(self, budgets: Optional[OptimizationBudgets] = None): self._terminal_sequence: Optional[int] = None self._terminal_error_code: Optional[str] = None self._terminal_error_type: Optional[str] = None + self._terminal_from_provider_call = False self._attached_owner: object = _UNATTACHED @property @@ -376,9 +394,25 @@ def raise_if_cancelled(self) -> None: # --- terminal transitions ------------------------------------------------- def finalize_success(self) -> None: - """Marks a normal optimizer return ``completed`` (first terminal wins).""" + """Marks a normal optimizer return ``completed`` (first terminal wins). + + Defensive: a pending cancellation observed here commits ``cancelled`` + and raises the typed error -- a cancelled run can never be recorded as a + success, no matter how late the signal arrived. + """ with self._lock: - self._transition_locked(RunStatus.COMPLETED) + if self._run_status is not None: + return + if self._cancel_requested: + self._transition_locked(RunStatus.CANCELLED) + error = OptimizationCancelledError( + f"Optimization cancelled: {self._cancel_reason}", + self._snapshot_locked(), + ) + else: + self._transition_locked(RunStatus.COMPLETED) + return + raise error def finalize_cancelled(self, reason: str = "cancelled") -> None: """Commits the ``cancelled`` terminal (e.g. on native task cancellation). @@ -393,8 +427,9 @@ def finalize_cancelled(self, reason: str = "cancelled") -> None: for record in self._records: if not record.closed: record.closed = True - record.end_time = time.monotonic() + record.end_time = time.time() record.state = ModelCallState.CANCELLED + self._settled_calls += 1 self._transition_locked(RunStatus.CANCELLED) def finalize_failed( @@ -410,7 +445,9 @@ def finalize_failed( """ with self._lock: self._transition_locked( - RunStatus.FAILED, error_code=error_code, error_type=error_type + RunStatus.FAILED, + error_code=_sanitize_error_meta(error_code), + error_type=_sanitize_error_meta(error_type), ) def _transition_locked( @@ -420,6 +457,7 @@ def _transition_locked( sequence: Optional[int] = None, error_code: Optional[str] = None, error_type: Optional[str] = None, + provider: bool = False, ) -> bool: """First-terminal-wins transition; returns True if this call won.""" if self._run_status is not None: @@ -428,6 +466,7 @@ def _transition_locked( self._terminal_sequence = sequence self._terminal_error_code = error_code self._terminal_error_type = error_type + self._terminal_from_provider_call = provider return True def _terminal_error_locked(self) -> Exception: @@ -442,7 +481,11 @@ def _terminal_error_locked(self) -> Exception: f"Optimization cancelled: {self._cancel_reason}", snapshot ) if self._run_status == RunStatus.FAILED: - return OptimizationProviderError( + if self._terminal_from_provider_call: + return OptimizationProviderError( + f"Provider failure: {self._terminal_error_code}", snapshot + ) + return OptimizationFailedError( f"Optimization failed: {self._terminal_error_code}", snapshot ) return OptimizationRunFinalizedError( @@ -495,21 +538,38 @@ def end_model_call( returned_model_version: Optional[str] = None, error_code: Optional[str] = None, error_type: Optional[str] = None, + cancelled: bool = False, ) -> None: """Commits the terminal outcome of one logical model invocation. - Ownership and close are validated atomically under the context lock: a - handle from another context is a typed misuse error, and a concurrent - double-close commits exactly once. + All usage and error metadata is normalized and validated *before* the + record is mutated, then the event, counters, and any terminal transition + commit atomically under the context lock -- no exception path can expose + a half-closed event. Ownership and close are validated under the same + lock: a handle from another context is a typed misuse error, and a + concurrent double-close commits exactly once. Provider failure takes precedence over token overshoot: usage is committed either way, but a call ending in a provider error transitions the run to ``failed`` and raises ``OptimizationProviderError``. A clean over-budget final call is committed with call status ``completed`` while the run transitions to ``budget_exceeded`` and - ``OptimizationBudgetExceeded`` is raised. Error metadata must be - sanitized identifiers, never raw exception/payload text. + ``OptimizationBudgetExceeded`` is raised. ``cancelled=True`` settles the + call as ``cancelled`` (preserving usage evidence) without raising. + + If the run was already terminated by another call, this call still + settles for truthful accounting, and the existing terminal outcome is + then raised so no caller can continue scheduling work past a committed + terminal. """ + # Normalize/validate everything BEFORE touching the record. + usage = _extract_usage(usage_metadata) + clean_code = _sanitize_error_meta(error_code) + clean_type = _sanitize_error_meta(error_type) + is_provider_error = error_code is not None or error_type is not None + if is_provider_error and clean_code is None: + clean_code = "PROVIDER_ERROR" + record = handle._record with self._lock: if handle._context is not self: @@ -518,31 +578,35 @@ def end_model_call( ) if record.closed: return + was_terminal = self._run_status is not None record.closed = True - record.end_time = time.monotonic() + record.end_time = time.time() record.returned_model_version = returned_model_version - record.usage = _extract_usage(usage_metadata) - self._completed_calls += 1 - total = record.usage.get("total_tokens") + record.usage = usage + self._settled_calls += 1 + total = usage.get("total_tokens") if total is not None: self._cumulative_total_tokens += total error: Optional[Exception] = None - if error_code is not None or error_type is not None: + if cancelled: + record.state = ModelCallState.CANCELLED + elif is_provider_error: record.state = ModelCallState.PROVIDER_ERROR - record.error_code = error_code - record.error_type = error_type + record.error_code = clean_code + record.error_type = clean_type # Provider failure is primary even when the same terminal response # also crossed the token ceiling; both usage and token-compliance # evidence are preserved on the snapshot. if self._transition_locked( RunStatus.FAILED, sequence=record.sequence, - error_code=error_code, - error_type=error_type, + error_code=clean_code, + error_type=clean_type, + provider=True, ): error = OptimizationProviderError( - f"Provider failure on call {record.sequence}: {error_code}", + f"Provider failure on call {record.sequence}: {clean_code}", self._snapshot_locked(), ) else: @@ -559,6 +623,10 @@ def end_model_call( f"Token budget exhausted ({max_tokens}).", self._snapshot_locked(), ) + if error is None and was_terminal and not cancelled: + # Another call already terminated the run; the settling caller must + # observe that terminal instead of continuing. + error = self._terminal_error_locked() if error is not None: raise error @@ -616,7 +684,7 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: events=events, budgets=self._budgets, started_calls=self._started_calls, - completed_calls=self._completed_calls, + completed_calls=self._settled_calls, cumulative_total_tokens=self._cumulative_total_tokens, usage_coverage=run_coverage, token_budget_status=token_status, @@ -629,14 +697,42 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: ) +_ERROR_META_MAX_LEN = 128 +_ERROR_META_SAFE = re.compile(r"[^A-Za-z0-9_.:-]") + + +def _sanitize_error_meta(value: Any) -> Optional[str]: + """Normalizes error metadata to a bounded, single-token identifier. + + Accepts any input (providers report numeric codes, enums, or strings) and + never lets raw multiline/path-like text into the ledger. + """ + if value is None: + return None + text = str(value).strip() + if not text: + return None + text = _ERROR_META_SAFE.sub("_", text) + return text[:_ERROR_META_MAX_LEN] + + def _extract_usage(usage_metadata: Any) -> dict[str, Optional[int]]: - """Copies provider-reported token counters, truthfully (no zero-coercion).""" + """Copies provider-reported token counters, truthfully (no zero-coercion). + + Counters must be finite, non-negative numbers; anything else (NaN, inf, + negative, non-numeric) is treated as not reported rather than corrupting + the ledger. + """ if usage_metadata is None: return {} def _get(name: str) -> Optional[int]: value = getattr(usage_metadata, name, None) - return int(value) if isinstance(value, (int, float)) else None + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + if not math.isfinite(value) or value < 0: + return None + return int(value) return { "prompt_tokens": _get("prompt_token_count"), diff --git a/src/google/adk/optimization/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index 20f2a26d8b0..511f3e4b5cf 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -28,6 +28,7 @@ from google.adk.optimization.data_types import AgentWithScores from google.adk.optimization.data_types import OptimizerResult from google.adk.optimization.data_types import UnstructuredSamplingResult +from google.adk.optimization.run_context import OptimizationRunContext from google.adk.optimization.sampler import Sampler from google.genai import types as genai_types from pydantic import BaseModel @@ -205,7 +206,7 @@ async def optimize( initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], *, - run_context=None, + run_context: Optional[OptimizationRunContext] = None, ) -> OptimizerResult[AgentWithScores]: if run_context is not None: from .run_context import UnsupportedOptimizationContextError diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index f3676581429..982d92e3d30 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -402,3 +402,127 @@ async def test_simple_prompt_optimizer_rejects_context(self): ) # Rejected before any sampler work. sampler.get_train_example_ids.assert_not_called() + + +class TestAtomicSettlement: + """Round-2 adversarial findings: validate-before-commit.""" + + def test_nan_usage_does_not_corrupt_the_record(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=float("nan"))) + event = ctx.snapshot().events[0] + # NaN/inf/negative counters are "not reported", never a crash or a + # half-closed event. + assert event.state == ModelCallState.COMPLETED + assert event.total_tokens is None + assert event.usage_coverage == UsageCoverage.UNREPORTED + + def test_negative_and_inf_usage_treated_as_unreported(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, + usage_metadata=_usage( + total_token_count=float("inf"), prompt_token_count=-3 + ), + ) + event = ctx.snapshot().events[0] + assert event.total_tokens is None + assert event.prompt_tokens is None + + def test_numeric_error_code_is_normalized_to_string(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call(h, usage_metadata=None, error_code=429) + snap = exc.value.snapshot + assert snap.events[0].error_code == "429" + assert snap.terminal_error_code == "429" + + def test_error_metadata_is_sanitized_and_bounded(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call( + h, + error_code="multi\nline /etc/passwd " + "x" * 500, + error_type="/usr/lib/module.py", + ) + event = exc.value.snapshot.events[0] + assert "\n" not in event.error_code + assert "/" not in event.error_code + assert len(event.error_code) <= 128 + assert "/" not in event.error_type + + +class TestTerminalCompleteness: + """Round-2 adversarial findings: every caller observes the terminal.""" + + def test_finalize_success_observes_pending_cancellation(self): + ctx = OptimizationRunContext() + ctx.request_cancel("deadline") + with pytest.raises(OptimizationCancelledError): + ctx.finalize_success() + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + def test_inflight_settlement_after_terminal_raises_existing_terminal(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h1 = ctx.begin_model_call(STAGE_REFLECTION) + h2 = ctx.begin_model_call(STAGE_REFLECTION) # admitted before terminal + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h1, usage_metadata=_usage(total_token_count=50)) + # The second in-flight call settles for truthful accounting, but its + # caller receives the existing terminal and cannot continue. + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h2, usage_metadata=_usage(total_token_count=5)) + snap = ctx.snapshot() + assert snap.completed_calls == 2 # both settled + assert snap.cumulative_total_tokens == 55 # evidence preserved + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + + def test_cancelled_settlement_preserves_usage_without_raising(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, usage_metadata=_usage(total_token_count=7), cancelled=True + ) + snap = ctx.snapshot() + assert snap.events[0].state == ModelCallState.CANCELLED + assert snap.events[0].total_tokens == 7 + assert snap.completed_calls == 1 # settled calls, any call status + + +class TestFailureTypes: + + def test_finalize_failed_is_generic_not_provider(self): + from google.adk.optimization.run_context import OptimizationFailedError + + ctx = OptimizationRunContext() + ctx.finalize_failed(error_code="SAMPLER_CRASH", error_type="ValueError") + with pytest.raises(OptimizationFailedError): + ctx.begin_model_call(STAGE_REFLECTION) + + def test_provider_call_failure_stays_provider_typed(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError): + ctx.end_model_call(h, error_code="INTERNAL") + with pytest.raises(OptimizationProviderError): + ctx.begin_model_call(STAGE_REFLECTION) + + +class TestDurableTimestamps: + + def test_event_times_are_wall_clock(self): + import time as _time + + before = _time.time() + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=None) + event = ctx.snapshot().events[0] + after = _time.time() + assert before <= event.start_time <= event.end_time <= after From 815eded288d288458c59cca652408158377f3eaf Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 01:40:33 -0700 Subject: [PATCH 5/9] round-3 contracts: success invariant boundary, complete field normalization, local-abort settlement, private module packaging, typed internals - finalize_success is an invariant boundary: idempotent only on COMPLETED; any other existing terminal re-raises its typed outcome; a pending cancellation commits CANCELLED and raises; an open (unsettled) call commits FAILED (OPEN_MODEL_CALLS) and raises OptimizationFailedError so a run can never escape nonterminal with an open event. - returned_model_version is normalized (bounded optional string) before mutation -- a malformed value from a custom BaseLlm can no longer make every subsequent snapshot raise ValidationError. - Fractional token counters are unreported, not silently truncated (ints excluding bool; floats only when integral); zero and very large integers accepted -- tested. - abort_model_call settles an admitted call that failed locally as call-status ABORTED with run-status FAILED via OptimizationFailedError: truthfully distinct from provider failures and cancellation (needed by the enforcement PR's post-admission exception boundary). - ADK visibility rule: module renamed to _run_context.py; the supported experimental surface is exported from google.adk.optimization with explicit __all__; internal imports and tests use the package surface. - Typed _CallRecord/_CallHandle constructors and attributes; optional .value access narrowed (the five repo-config mypy errors named in review; exact gate reproduction is blocked in this environment by a numpy-stub / python_version=3.11 syntax failure, noted honestly). - New tests: success-with-open-call, success-after-budget/provider/cancel, malformed model version, fractional/integral/zero/large usage, local abort, concurrent attach (1 of 16 wins), concurrent admission (exactly k of 16, unique sequences, shared terminal), success-vs-settlement race. Suite: 54 contract tests, full optimization directory green. --- src/google/adk/optimization/__init__.py | 48 ++++ .../{run_context.py => _run_context.py} | 129 +++++++++-- .../adk/optimization/agent_optimizer.py | 4 +- .../optimization/gepa_root_agent_optimizer.py | 4 +- .../gepa_root_agent_prompt_optimizer.py | 4 +- .../optimization/simple_prompt_optimizer.py | 4 +- .../optimization/run_context_test.py | 214 ++++++++++++++++-- 7 files changed, 363 insertions(+), 44 deletions(-) rename src/google/adk/optimization/{run_context.py => _run_context.py} (86%) diff --git a/src/google/adk/optimization/__init__.py b/src/google/adk/optimization/__init__.py index 58d482ea386..c99bbe6438e 100644 --- a/src/google/adk/optimization/__init__.py +++ b/src/google/adk/optimization/__init__.py @@ -11,3 +11,51 @@ # 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. + +"""Agent optimization: run-scoped governance surface (experimental). + +The run-context contract is exported here as the supported public surface; +optimizer and sampler classes remain importable from their submodules. +""" + +from ._run_context import ContextAlreadyAttachedError +from ._run_context import ModelCallEvent +from ._run_context import ModelCallState +from ._run_context import OptimizationBudgetExceeded +from ._run_context import OptimizationBudgets +from ._run_context import OptimizationCancelledError +from ._run_context import OptimizationFailedError +from ._run_context import OptimizationProviderError +from ._run_context import OptimizationRunContext +from ._run_context import OptimizationRunContextError +from ._run_context import OptimizationRunFinalizedError +from ._run_context import OptimizationRunSnapshot +from ._run_context import OptimizerCapabilities +from ._run_context import RunStatus +from ._run_context import STAGE_CANDIDATE_GENERATION +from ._run_context import STAGE_REFLECTION +from ._run_context import TokenBudgetStatus +from ._run_context import UnsupportedOptimizationContextError +from ._run_context import UsageCoverage + +__all__ = [ + "ContextAlreadyAttachedError", + "ModelCallEvent", + "ModelCallState", + "OptimizationBudgetExceeded", + "OptimizationBudgets", + "OptimizationCancelledError", + "OptimizationFailedError", + "OptimizationProviderError", + "OptimizationRunContext", + "OptimizationRunContextError", + "OptimizationRunFinalizedError", + "OptimizationRunSnapshot", + "OptimizerCapabilities", + "RunStatus", + "STAGE_CANDIDATE_GENERATION", + "STAGE_REFLECTION", + "TokenBudgetStatus", + "UnsupportedOptimizationContextError", + "UsageCoverage", +] diff --git a/src/google/adk/optimization/run_context.py b/src/google/adk/optimization/_run_context.py similarity index 86% rename from src/google/adk/optimization/run_context.py rename to src/google/adk/optimization/_run_context.py index 4bcb3cd7cf6..334ba6d1954 100644 --- a/src/google/adk/optimization/run_context.py +++ b/src/google/adk/optimization/_run_context.py @@ -79,6 +79,10 @@ class ModelCallState(str, enum.Enum): COMPLETED = "completed" PROVIDER_ERROR = "provider_error" CANCELLED = "cancelled" + ABORTED = "aborted" + """The invocation was admitted but aborted locally (request construction, + scheduling, or other optimizer-side failure) before or during execution -- + distinct from a provider failure and from cancellation.""" class RunStatus(str, enum.Enum): @@ -276,18 +280,20 @@ class _CallRecord: "closed", ) - def __init__(self, sequence: int, stage: str, requested_model): - self.sequence = sequence - self.stage = stage - self.requested_model = requested_model - self.returned_model_version = None - self.start_time = time.time() - self.end_time = None + def __init__( + self, sequence: int, stage: str, requested_model: Optional[str] + ) -> None: + self.sequence: int = sequence + self.stage: str = stage + self.requested_model: Optional[str] = requested_model + self.returned_model_version: Optional[str] = None + self.start_time: float = time.time() + self.end_time: Optional[float] = None self.state: Optional[ModelCallState] = None self.usage: dict[str, Optional[int]] = {} self.error_code: Optional[str] = None self.error_type: Optional[str] = None - self.closed = False + self.closed: bool = False class _CallHandle: @@ -297,9 +303,11 @@ class _CallHandle: typed misuse error. """ - def __init__(self, context: "OptimizationRunContext", record: _CallRecord): - self._context = context - self._record = record + def __init__( + self, context: "OptimizationRunContext", record: _CallRecord + ) -> None: + self._context: "OptimizationRunContext" = context + self._record: _CallRecord = record _UNATTACHED = object() @@ -319,7 +327,7 @@ class OptimizationRunContext: any thread, and read :meth:`snapshot` at any time, including after failure. """ - def __init__(self, budgets: Optional[OptimizationBudgets] = None): + def __init__(self, budgets: Optional[OptimizationBudgets] = None) -> None: # OptimizationBudgets is frozen; keep a private reference and never # expose a mutable path to it. self._budgets = budgets or OptimizationBudgets() @@ -394,21 +402,40 @@ def raise_if_cancelled(self) -> None: # --- terminal transitions ------------------------------------------------- def finalize_success(self) -> None: - """Marks a normal optimizer return ``completed`` (first terminal wins). - - Defensive: a pending cancellation observed here commits ``cancelled`` - and raises the typed error -- a cancelled run can never be recorded as a - success, no matter how late the signal arrived. + """Commits ``completed`` -- and enforces that success is actually possible. + + This is an invariant boundary, not a status setter: + + - idempotent only when the run is already ``completed``; any other + existing terminal re-raises its typed outcome so a caller cannot + swallow a committed budget/provider/cancel terminal into success; + - a pending cancellation commits ``cancelled`` and raises; + - an open (unsettled) call makes success impossible: the run commits + ``failed`` (``OPEN_MODEL_CALLS``) and raises + ``OptimizationFailedError``. Already-admitted calls may still settle + later and will observe that terminal. """ with self._lock: if self._run_status is not None: - return - if self._cancel_requested: + if self._run_status == RunStatus.COMPLETED: + return + error = self._terminal_error_locked() + elif self._cancel_requested: self._transition_locked(RunStatus.CANCELLED) error = OptimizationCancelledError( f"Optimization cancelled: {self._cancel_reason}", self._snapshot_locked(), ) + elif any(not r.closed for r in self._records): + self._transition_locked( + RunStatus.FAILED, + error_code="OPEN_MODEL_CALLS", + error_type="OptimizationRunFinalizedError", + ) + error = OptimizationFailedError( + "Cannot finalize success with unsettled model calls.", + self._snapshot_locked(), + ) else: self._transition_locked(RunStatus.COMPLETED) return @@ -581,7 +608,9 @@ def end_model_call( was_terminal = self._run_status is not None record.closed = True record.end_time = time.time() - record.returned_model_version = returned_model_version + record.returned_model_version = _sanitize_optional_str( + returned_model_version, _MODEL_VERSION_MAX_LEN + ) record.usage = usage self._settled_calls += 1 total = usage.get("total_tokens") @@ -630,6 +659,51 @@ def end_model_call( if error is not None: raise error + def abort_model_call( + self, + handle: _CallHandle, + *, + error_code: Optional[str] = None, + error_type: Optional[str] = None, + ) -> None: + """Settles an admitted call that failed locally (non-provider). + + Closes the record as ``aborted``, transitions the run to ``failed`` with + sanitized metadata, and raises ``OptimizationFailedError`` -- truthfully + distinct from provider failures and from cancellation. Idempotent on an + already-settled handle; foreign handles are a typed misuse error. + """ + clean_code = _sanitize_error_meta(error_code) or "LOCAL_CALL_ABORT" + clean_type = _sanitize_error_meta(error_type) + record = handle._record + with self._lock: + if handle._context is not self: + raise OptimizationRunFinalizedError( + "Call handle belongs to a different OptimizationRunContext." + ) + error: Optional[Exception] = None + if not record.closed: + record.closed = True + record.end_time = time.time() + record.state = ModelCallState.ABORTED + record.error_code = clean_code + record.error_type = clean_type + self._settled_calls += 1 + if self._transition_locked( + RunStatus.FAILED, + sequence=record.sequence, + error_code=clean_code, + error_type=clean_type, + ): + error = OptimizationFailedError( + f"Local failure on call {record.sequence}: {clean_code}", + self._snapshot_locked(), + ) + if error is None and self._run_status is not None: + error = self._terminal_error_locked() + if error is not None: + raise error + # --- snapshots -------------------------------------------------------------- def snapshot(self) -> OptimizationRunSnapshot: @@ -701,6 +775,17 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: _ERROR_META_SAFE = re.compile(r"[^A-Za-z0-9_.:-]") +_MODEL_VERSION_MAX_LEN = 256 + + +def _sanitize_optional_str(value: Any, max_len: int) -> Optional[str]: + """Bounded optional-string normalization for ledger fields.""" + if value is None: + return None + text = str(value).strip() + return text[:max_len] if text else None + + def _sanitize_error_meta(value: Any) -> Optional[str]: """Normalizes error metadata to a bounded, single-token identifier. @@ -732,6 +817,10 @@ def _get(name: str) -> Optional[int]: return None if not math.isfinite(value) or value < 0: return None + if isinstance(value, float) and not value.is_integer(): + # Token counters are integral evidence; silent truncation is not + # truthful normalization -- fractional values are "not reported". + return None return int(value) return { diff --git a/src/google/adk/optimization/agent_optimizer.py b/src/google/adk/optimization/agent_optimizer.py index ff9b44cadca..2f72469b177 100644 --- a/src/google/adk/optimization/agent_optimizer.py +++ b/src/google/adk/optimization/agent_optimizer.py @@ -20,11 +20,11 @@ from typing import Optional from ..agents.llm_agent import Agent +from ._run_context import OptimizationRunContext +from ._run_context import OptimizerCapabilities from .data_types import AgentWithScoresT from .data_types import OptimizerResult from .data_types import SamplingResultT -from .run_context import OptimizationRunContext -from .run_context import OptimizerCapabilities from .sampler import Sampler diff --git a/src/google/adk/optimization/gepa_root_agent_optimizer.py b/src/google/adk/optimization/gepa_root_agent_optimizer.py index fc08ed5cb12..bb615bfddd2 100644 --- a/src/google/adk/optimization/gepa_root_agent_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_optimizer.py @@ -33,11 +33,11 @@ from ..tools.skill_toolset import SkillToolset from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._run_context import OptimizationRunContext from .agent_optimizer import AgentOptimizer from .data_types import AgentWithScores from .data_types import OptimizerResult from .data_types import UnstructuredSamplingResult -from .run_context import OptimizationRunContext from .sampler import Sampler logger = logging.getLogger("google_adk." + __name__) @@ -351,7 +351,7 @@ async def optimize( agent instance, its scores on the validation examples, and other metrics. """ if run_context is not None: - from .run_context import UnsupportedOptimizationContextError + from ._run_context import UnsupportedOptimizationContextError raise UnsupportedOptimizationContextError( "GEPARootAgentOptimizer does not support OptimizationRunContext in" diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py index de78e709fce..658e871d3a4 100644 --- a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -31,11 +31,11 @@ from ..models.registry import LLMRegistry from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._run_context import OptimizationRunContext from .agent_optimizer import AgentOptimizer from .data_types import AgentWithScores from .data_types import OptimizerResult from .data_types import UnstructuredSamplingResult -from .run_context import OptimizationRunContext from .sampler import Sampler _logger = logging.getLogger("google_adk." + __name__) @@ -224,7 +224,7 @@ async def optimize( agent instance, its scores on the validation examples, and other metrics. """ if run_context is not None: - from .run_context import UnsupportedOptimizationContextError + from ._run_context import UnsupportedOptimizationContextError raise UnsupportedOptimizationContextError( "GEPARootAgentPromptOptimizer does not support OptimizationRunContext" diff --git a/src/google/adk/optimization/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index 511f3e4b5cf..e98c844c3ff 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -24,11 +24,11 @@ from google.adk.evaluation._retry_options_utils import add_default_retry_options_if_not_present from google.adk.models.llm_request import LlmRequest from google.adk.models.registry import LLMRegistry +from google.adk.optimization._run_context import OptimizationRunContext from google.adk.optimization.agent_optimizer import AgentOptimizer from google.adk.optimization.data_types import AgentWithScores from google.adk.optimization.data_types import OptimizerResult from google.adk.optimization.data_types import UnstructuredSamplingResult -from google.adk.optimization.run_context import OptimizationRunContext from google.adk.optimization.sampler import Sampler from google.genai import types as genai_types from pydantic import BaseModel @@ -209,7 +209,7 @@ async def optimize( run_context: Optional[OptimizationRunContext] = None, ) -> OptimizerResult[AgentWithScores]: if run_context is not None: - from .run_context import UnsupportedOptimizationContextError + from ._run_context import UnsupportedOptimizationContextError raise UnsupportedOptimizationContextError( "SimplePromptOptimizer does not support OptimizationRunContext in" diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index 982d92e3d30..ee4f8dca0ff 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -17,20 +17,20 @@ import threading from types import SimpleNamespace -from google.adk.optimization.run_context import ContextAlreadyAttachedError -from google.adk.optimization.run_context import ModelCallState -from google.adk.optimization.run_context import OptimizationBudgetExceeded -from google.adk.optimization.run_context import OptimizationBudgets -from google.adk.optimization.run_context import OptimizationCancelledError -from google.adk.optimization.run_context import OptimizationProviderError -from google.adk.optimization.run_context import OptimizationRunContext -from google.adk.optimization.run_context import OptimizationRunFinalizedError -from google.adk.optimization.run_context import OptimizerCapabilities -from google.adk.optimization.run_context import RunStatus -from google.adk.optimization.run_context import STAGE_CANDIDATE_GENERATION -from google.adk.optimization.run_context import STAGE_REFLECTION -from google.adk.optimization.run_context import TokenBudgetStatus -from google.adk.optimization.run_context import UsageCoverage +from google.adk.optimization import ContextAlreadyAttachedError +from google.adk.optimization import ModelCallState +from google.adk.optimization import OptimizationBudgetExceeded +from google.adk.optimization import OptimizationBudgets +from google.adk.optimization import OptimizationCancelledError +from google.adk.optimization import OptimizationProviderError +from google.adk.optimization import OptimizationRunContext +from google.adk.optimization import OptimizationRunFinalizedError +from google.adk.optimization import OptimizerCapabilities +from google.adk.optimization import RunStatus +from google.adk.optimization import STAGE_CANDIDATE_GENERATION +from google.adk.optimization import STAGE_REFLECTION +from google.adk.optimization import TokenBudgetStatus +from google.adk.optimization import UsageCoverage import pydantic import pytest @@ -386,7 +386,7 @@ class TestBuiltInsRejectContextInContractsPr: async def test_simple_prompt_optimizer_rejects_context(self): from unittest import mock - from google.adk.optimization.run_context import UnsupportedOptimizationContextError + from google.adk.optimization import UnsupportedOptimizationContextError from google.adk.optimization.sampler import Sampler from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig @@ -498,7 +498,7 @@ def test_cancelled_settlement_preserves_usage_without_raising(self): class TestFailureTypes: def test_finalize_failed_is_generic_not_provider(self): - from google.adk.optimization.run_context import OptimizationFailedError + from google.adk.optimization import OptimizationFailedError ctx = OptimizationRunContext() ctx.finalize_failed(error_code="SAMPLER_CRASH", error_type="ValueError") @@ -526,3 +526,185 @@ def test_event_times_are_wall_clock(self): event = ctx.snapshot().events[0] after = _time.time() assert before <= event.start_time <= event.end_time <= after + + +class TestSuccessInvariantBoundary: + """Round-3: finalize_success is an invariant boundary, not a status setter.""" + + def test_success_with_open_call_commits_failed(self): + from google.adk.optimization import OptimizationFailedError + + ctx = OptimizationRunContext() + ctx.begin_model_call(STAGE_REFLECTION) # left open + with pytest.raises(OptimizationFailedError): + ctx.finalize_success() + snap = ctx.snapshot() + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "OPEN_MODEL_CALLS" + + def test_success_after_budget_terminal_reraises(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + with pytest.raises(OptimizationBudgetExceeded): + ctx.finalize_success() # cannot swallow the committed terminal + assert ctx.snapshot().run_status == RunStatus.BUDGET_EXCEEDED + + def test_success_after_provider_terminal_reraises(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError): + ctx.end_model_call(h, error_code="INTERNAL") + with pytest.raises(OptimizationProviderError): + ctx.finalize_success() + + def test_success_after_cancel_terminal_reraises(self): + ctx = OptimizationRunContext() + ctx.finalize_cancelled("x") + with pytest.raises(OptimizationCancelledError): + ctx.finalize_success() + + def test_success_idempotent_only_when_completed(self): + ctx = OptimizationRunContext() + ctx.finalize_success() + ctx.finalize_success() # idempotent + assert ctx.snapshot().run_status == RunStatus.COMPLETED + + +class TestLedgerFieldNormalization: + """Round-3: every event field normalized before mutation.""" + + def test_malformed_model_version_cannot_corrupt_snapshots(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, + usage_metadata=_usage(total_token_count=3), + returned_model_version=123, # non-string from a custom BaseLlm + ) + snap = ctx.snapshot() # must not raise + assert snap.events[0].returned_model_version == "123" + + def test_fractional_usage_is_unreported_not_truncated(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=1.9)) + event = ctx.snapshot().events[0] + assert event.total_tokens is None + assert event.usage_coverage == UsageCoverage.UNREPORTED + + def test_integral_float_zero_and_large_int_accepted(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, + usage_metadata=_usage( + total_token_count=2.0, + prompt_token_count=0, + candidates_token_count=10**12, + ), + ) + event = ctx.snapshot().events[0] + assert event.total_tokens == 2 + assert event.prompt_tokens == 0 + assert event.output_tokens == 10**12 + + +class TestLocalCallAbort: + + def test_abort_settles_as_failed_not_provider(self): + from google.adk.optimization import OptimizationFailedError + + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationFailedError): + ctx.abort_model_call( + h, error_code="SCHEDULING_FAILURE", error_type="RuntimeError" + ) + snap = ctx.snapshot() + assert snap.events[0].state == ModelCallState.ABORTED + assert snap.run_status == RunStatus.FAILED + assert snap.completed_calls == 1 # settled + + +class TestConcurrencyRaces: + """Round-3: the RFC-required concurrent attachment and admission races.""" + + def test_concurrent_attach_exactly_one_wins(self): + ctx = OptimizationRunContext() + results: list[bool] = [] + + def try_attach(): + try: + ctx.attach(owner=object()) + results.append(True) + except ContextAlreadyAttachedError: + results.append(False) + + threads = [threading.Thread(target=try_attach) for _ in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + assert results.count(True) == 1 + + def test_concurrent_admission_admits_exactly_k(self): + k = 3 + ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=k)) + admitted: list[int] = [] + rejected: list[Exception] = [] + + def try_admit(): + try: + h = ctx.begin_model_call(STAGE_REFLECTION) + admitted.append(h._record.sequence) + except OptimizationBudgetExceeded as e: + rejected.append(e) + + threads = [threading.Thread(target=try_admit) for _ in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + assert len(admitted) == k + assert sorted(admitted) == [1, 2, 3] # unique sequence ids + assert len(rejected) == 16 - k + assert all( + e.snapshot.run_status == RunStatus.BUDGET_EXCEEDED for e in rejected + ) + + def test_success_racing_last_settlement_never_completes_with_open_call( + self, + ): + from google.adk.optimization import OptimizationFailedError + + for _ in range(20): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + outcomes: list[str] = [] + + def settle(handle=h, context=ctx): + context.end_model_call(handle, usage_metadata=None) + + def finish(context=ctx, sink=outcomes): + try: + context.finalize_success() + sink.append("completed") + except OptimizationFailedError: + sink.append("failed_open") + + t1 = threading.Thread(target=settle) + t2 = threading.Thread(target=finish) + t2.start() + t1.start() + t1.join() + t2.join() + snap = ctx.snapshot() + if "completed" in outcomes: + # Success was only recordable because the call had already settled. + assert snap.events[0].state is not None + else: + assert snap.run_status == RunStatus.FAILED From b12baacfeb1203bc62b7ed19dffe8bb5acddfd4c Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 02:06:31 -0700 Subject: [PATCH 6/9] expose terminal provenance on the snapshot (terminal_from_provider_call) Lets consumers distinguish a provider-call FAILED terminal from a local/ optimizer FAILED terminal without private state -- required by the stacked enforcement PR's post-check mapping (OptimizationProviderError vs OptimizationFailedError). --- src/google/adk/optimization/_run_context.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/google/adk/optimization/_run_context.py b/src/google/adk/optimization/_run_context.py index 334ba6d1954..0bffca82016 100644 --- a/src/google/adk/optimization/_run_context.py +++ b/src/google/adk/optimization/_run_context.py @@ -213,6 +213,10 @@ class OptimizationRunSnapshot(BaseModel): terminal_error_code: Optional[str] = None terminal_error_type: Optional[str] = None + terminal_from_provider_call: bool = False + """True when a ``failed`` terminal came from a governed provider call; + False for local/optimizer failures recorded via ``finalize_failed`` or + ``abort_model_call``.""" class OptimizationRunContextError(Exception): @@ -768,6 +772,7 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: terminal_sequence=self._terminal_sequence, terminal_error_code=self._terminal_error_code, terminal_error_type=self._terminal_error_type, + terminal_from_provider_call=self._terminal_from_provider_call, ) From d0b10d5eeaf022aedea46ba54c9c231153b2c4a6 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 02:44:00 -0700 Subject: [PATCH 7/9] round-4 contracts: boundary-hostile inputs cannot strand or corrupt the ledger; duplicate terminal settlement re-raises; race test is warning-clean - _extract_usage is defensive around attribute ACCESS as well as conversion: a throwing usage property or hostile numeric degrades that counter to unreported and the call settles atomically (no more stranded open events). All str() normalization paths survive hostile __str__ (error metadata degrades to UNSTRINGABLE; optional fields to None). - Duplicate settlement of an already-closed call on a TERMINAL run re-raises the committed terminal (governance callers cannot continue past it); the no-op return remains only for a closed call on a nonterminal run. abort_model_call follows the same rule. - Admission and cancellation metadata are normalized at the boundary: stage must be a non-empty bounded string (ValueError otherwise, no slot consumed), requested_model and cancel reason are bounded normalized strings -- hostile values can no longer make every snapshot raise. - The remaining strict-mypy optional-.value read is narrowed via a local. - The success/settlement race test uses a Barrier, captures both threads' outcomes, asserts the two legal resolutions explicitly and completed_calls == started_calls -- clean under -W error::pytest.PytestUnhandledThreadExceptionWarning. - Seven new regression tests reproduce every round-4 contract finding. --- src/google/adk/optimization/_run_context.py | 77 +++++++++--- .../optimization/run_context_test.py | 110 ++++++++++++++++-- 2 files changed, 160 insertions(+), 27 deletions(-) diff --git a/src/google/adk/optimization/_run_context.py b/src/google/adk/optimization/_run_context.py index 0bffca82016..d1d02ba6410 100644 --- a/src/google/adk/optimization/_run_context.py +++ b/src/google/adk/optimization/_run_context.py @@ -377,7 +377,9 @@ def request_cancel(self, reason: str = "requested") -> None: with self._lock: if not self._cancel_requested: self._cancel_requested = True - self._cancel_reason = reason + self._cancel_reason = ( + _sanitize_optional_str(reason, _ERROR_META_MAX_LEN) or "requested" + ) @property def cancel_requested(self) -> bool: @@ -454,7 +456,9 @@ def finalize_cancelled(self, reason: str = "cancelled") -> None: with self._lock: if not self._cancel_requested: self._cancel_requested = True - self._cancel_reason = reason + self._cancel_reason = ( + _sanitize_optional_str(reason, _ERROR_META_MAX_LEN) or "requested" + ) for record in self._records: if not record.closed: record.closed = True @@ -519,8 +523,10 @@ def _terminal_error_locked(self) -> Exception: return OptimizationFailedError( f"Optimization failed: {self._terminal_error_code}", snapshot ) + status = self._run_status + status_name = status.value if status is not None else "unknown" return OptimizationRunFinalizedError( - f"OptimizationRunContext already finalized as {self._run_status.value}." + f"OptimizationRunContext already finalized as {status_name}." ) # --- the ledger ----------------------------------------------------------- @@ -555,8 +561,17 @@ def begin_model_call( self._snapshot_locked(), ) else: + clean_stage = _sanitize_optional_str(stage, _ERROR_META_MAX_LEN) + if not clean_stage: + raise ValueError( + "stage must be a non-empty string (see the STAGE_* constants)." + ) self._started_calls += 1 - record = _CallRecord(self._started_calls, stage, requested_model) + record = _CallRecord( + self._started_calls, + clean_stage, + _sanitize_optional_str(requested_model, _MODEL_VERSION_MAX_LEN), + ) self._records.append(record) return _CallHandle(self, record) raise error @@ -608,7 +623,14 @@ def end_model_call( "Call handle belongs to a different OptimizationRunContext." ) if record.closed: - return + # Ledger-idempotent, but a caller must not continue past a committed + # terminal: re-raise the existing terminal outcome. The no-op return + # is only for an already-closed call on a still-nonterminal run. + if self._run_status is not None: + error = self._terminal_error_locked() + else: + return + raise error was_terminal = self._run_status is not None record.closed = True record.end_time = time.time() @@ -686,7 +708,10 @@ def abort_model_call( "Call handle belongs to a different OptimizationRunContext." ) error: Optional[Exception] = None - if not record.closed: + if record.closed: + if self._run_status is not None: + error = self._terminal_error_locked() + else: record.closed = True record.end_time = time.time() record.state = ModelCallState.ABORTED @@ -784,10 +809,17 @@ def _snapshot_locked(self) -> OptimizationRunSnapshot: def _sanitize_optional_str(value: Any, max_len: int) -> Optional[str]: - """Bounded optional-string normalization for ledger fields.""" + """Bounded optional-string normalization for ledger fields. + + Defensive against hostile ``__str__``: a conversion failure degrades the + field to ``None`` rather than defeating settlement. + """ if value is None: return None - text = str(value).strip() + try: + text = str(value).strip() + except Exception: # pylint: disable=broad-except + return None return text[:max_len] if text else None @@ -799,7 +831,10 @@ def _sanitize_error_meta(value: Any) -> Optional[str]: """ if value is None: return None - text = str(value).strip() + try: + text = str(value).strip() + except Exception: # pylint: disable=broad-except + return "UNSTRINGABLE" if not text: return None text = _ERROR_META_SAFE.sub("_", text) @@ -817,16 +852,22 @@ def _extract_usage(usage_metadata: Any) -> dict[str, Optional[int]]: return {} def _get(name: str) -> Optional[int]: - value = getattr(usage_metadata, name, None) - if isinstance(value, bool) or not isinstance(value, (int, float)): - return None - if not math.isfinite(value) or value < 0: - return None - if isinstance(value, float) and not value.is_integer(): - # Token counters are integral evidence; silent truncation is not - # truthful normalization -- fractional values are "not reported". + # Defensive around BOTH attribute access and conversion: a throwing + # property or hostile numeric type from a custom provider degrades that + # counter to unreported -- it must never strand an admitted call open. + try: + value = getattr(usage_metadata, name, None) + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + if not math.isfinite(value) or value < 0: + return None + if isinstance(value, float) and not value.is_integer(): + # Token counters are integral evidence; silent truncation is not + # truthful normalization -- fractional values are "not reported". + return None + return int(value) + except Exception: # pylint: disable=broad-except return None - return int(value) return { "prompt_tokens": _get("prompt_token_count"), diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index ee4f8dca0ff..5baa32feb1d 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -684,27 +684,119 @@ def test_success_racing_last_settlement_never_completes_with_open_call( for _ in range(20): ctx = OptimizationRunContext() h = ctx.begin_model_call(STAGE_REFLECTION) - outcomes: list[str] = [] + barrier = threading.Barrier(2) + outcomes: dict[str, object] = {} - def settle(handle=h, context=ctx): - context.end_model_call(handle, usage_metadata=None) + def settle(handle=h, context=ctx, sync=barrier, sink=outcomes): + sync.wait() + try: + context.end_model_call(handle, usage_metadata=None) + sink["settle"] = "ok" + except OptimizationFailedError: + # Success won the race, committed FAILED/OPEN_MODEL_CALLS, and + # this settling caller observes it after committing its event. + sink["settle"] = "observed_terminal" - def finish(context=ctx, sink=outcomes): + def finish(context=ctx, sync=barrier, sink=outcomes): + sync.wait() try: context.finalize_success() - sink.append("completed") + sink["success"] = "completed" except OptimizationFailedError: - sink.append("failed_open") + sink["success"] = "failed_open" t1 = threading.Thread(target=settle) t2 = threading.Thread(target=finish) - t2.start() t1.start() + t2.start() t1.join() t2.join() snap = ctx.snapshot() - if "completed" in outcomes: - # Success was only recordable because the call had already settled. + # Exactly two legal resolutions -- and never COMPLETED with open calls. + assert snap.completed_calls == snap.started_calls == 1 + if outcomes["success"] == "completed": + assert outcomes["settle"] == "ok" + assert snap.run_status == RunStatus.COMPLETED assert snap.events[0].state is not None else: + assert outcomes["success"] == "failed_open" + assert outcomes["settle"] == "observed_terminal" assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "OPEN_MODEL_CALLS" + + +class TestRoundFourFindings: + + def test_throwing_usage_accessor_degrades_to_unreported(self): + class ExplodingUsage: + + @property + def total_token_count(self): + raise RuntimeError("bad provider metadata") + + prompt_token_count = 5 + + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=ExplodingUsage()) + snap = ctx.snapshot() + # The call settles atomically; the throwing counter is unreported and + # the well-behaved counter survives. + assert snap.completed_calls == 1 + event = snap.events[0] + assert event.state == ModelCallState.COMPLETED + assert event.total_tokens is None + assert event.prompt_tokens == 5 + + def test_hostile_str_cannot_defeat_settlement(self): + class HostileCode: + + def __str__(self): + raise RuntimeError("no string for you") + + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError): + ctx.end_model_call(h, error_code=HostileCode()) + snap = ctx.snapshot() + assert snap.completed_calls == 1 + assert snap.events[0].error_code == "UNSTRINGABLE" + + def test_duplicate_terminal_settlement_reraises_terminal(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + # A duplicate close on a terminal run must not return normally. + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=50)) + assert ctx.snapshot().completed_calls == 1 # committed exactly once + + def test_duplicate_close_on_nonterminal_run_is_noop(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=None) + ctx.end_model_call(h, usage_metadata=None) # no raise, no double count + assert ctx.snapshot().completed_calls == 1 + + def test_admission_metadata_normalized(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(123, requested_model=456) # hostile callers + ctx.end_model_call(h, usage_metadata=None) + snap = ctx.snapshot() # must not raise + assert snap.events[0].stage == "123" + assert snap.events[0].requested_model == "456" + + def test_empty_stage_rejected(self): + ctx = OptimizationRunContext() + with pytest.raises(ValueError): + ctx.begin_model_call(" ") + assert ctx.snapshot().started_calls == 0 + + def test_cancel_reason_normalized(self): + ctx = OptimizationRunContext() + ctx.request_cancel(123) + snap = ctx.snapshot() # must not raise + assert snap.cancel_reason == "123" From e9483709da4226b0769fbf97fdde0b4ad946b509 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 02:06:03 -0700 Subject: [PATCH 8/9] round-3 enforcement: pre-admission request build, exception-complete settlement, pre-cancel precedence, typed internals, GEPA behavior pin - GEPA reflection builds and validates LlmRequest BEFORE reserving the model-call slot (a construction failure is a GEPA-native no-proposal with a clean ledger, never an open call); from admission onward every operation -- including run_coroutine_threadsafe scheduling -- lives inside one settlement boundary; local scheduling failures settle via the new abort_model_call as call-status ABORTED / run-status FAILED (OptimizationFailedError, truthfully distinct from provider failures and cancellation) and convert to the sentinel so GEPA cannot swallow them into success. - Snapshot terminal provenance (terminal_from_provider_call) distinguishes provider FAILED from local FAILED in the post-check and sentinel-escape mappings. - Pre-cancellation wins over discovery: both optimizers observe raise_if_cancelled immediately after attach, before imports, model/ adapter construction, and any sampler method; cancellation also wins over a competing discovery exception (tested with sampler-method assert_not_called for both optimizers). - finalize_success is gated to genuinely non-terminal runs on the governed GEPA partial path (the authoritative BUDGET_EXCEEDED terminal survives). - Typed _ReflectionCapture replaces the heterogeneous holder dict; adapter run_context parameter typed; gepa.optimize is two fully typed branches -- the legacy call shape preserved exactly for absent-context parity, the governed branch pinning raise_on_exception=True so generic evaluator/ sampler failures stay fail-closed under the unbounded gepa>=0.1 range; sentinel docstring attribution corrected (reflective_mutation.py converts reflection errors to no-proposal independent of the engine flag). - New real-GEPA durable coverage: in-band provider error, post-sampler cancellation, reflection-request failure. Full suite: 106/106. --- .../gepa_root_agent_prompt_optimizer.py | 467 +++++++++-- .../optimization/simple_prompt_optimizer.py | 215 ++++- .../run_context_enforcement_test.py | 791 ++++++++++++++++++ .../optimization/run_context_test.py | 151 +--- 4 files changed, 1352 insertions(+), 272 deletions(-) create mode 100644 tests/unittests/optimization/run_context_enforcement_test.py diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py index 658e871d3a4..8befca35e3d 100644 --- a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -31,7 +31,14 @@ from ..models.registry import LLMRegistry from ..utils.context_utils import Aclosing from ..utils.feature_decorator import experimental +from ._run_context import OptimizationBudgetExceeded +from ._run_context import OptimizationFailedError +from ._run_context import OptimizationProviderError from ._run_context import OptimizationRunContext +from ._run_context import OptimizationRunContextError +from ._run_context import OptimizerCapabilities +from ._run_context import RunStatus +from ._run_context import STAGE_REFLECTION from .agent_optimizer import AgentOptimizer from .data_types import AgentWithScores from .data_types import OptimizerResult @@ -43,6 +50,33 @@ _AGENT_PROMPT_NAME = "agent_prompt" +class _ReflectionCapture: + """Typed carrier for reflection stream metadata, visible to error paths.""" + + __slots__ = ("usage", "model_version", "error_code") + + def __init__(self) -> None: + self.usage: Any = None + self.model_version: Optional[str] = None + self.error_code: Optional[str] = None + + +class _GovernedRunAbort(Exception): + """Private iteration-abort sentinel for governed GEPA runs. + + GEPA's reflective-mutation proposer converts reflection/proposal + exceptions into "no proposal" (``reflective_mutation.py``), independent of + the engine's ``raise_on_exception`` flag -- so a typed budget, provider, or + cancellation error raised from the reflection callable would be swallowed + and the run could still return success. (Other iteration failures default + to ``raise_on_exception=True`` in the admitted GEPA releases; governed runs + pin that explicitly.) The governed terminal state is + already committed on the run context when this sentinel is raised; the stop + callback observes it at the next loop boundary and the post-check after + ``gepa.optimize`` returns maps it to the typed outcome. + """ + + class GEPARootAgentPromptOptimizerConfig(BaseModel): """Contains configuration options required by the GEPARootAgentPromptOptimizer.""" @@ -103,10 +137,12 @@ def __init__( initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], main_loop: asyncio.AbstractEventLoop, + run_context: Optional[OptimizationRunContext] = None, ): self._initial_agent = initial_agent self._sampler = sampler self._main_loop = main_loop + self._run_context = run_context self._train_example_ids = set(sampler.get_train_example_ids()) self._validation_example_ids = set(sampler.get_validation_example_ids()) @@ -131,6 +167,15 @@ def evaluate( else: raise ValueError(f"Invalid batch composition: {batch}") + if self._run_context is not None: + # Boundary: before a sampler evaluation is scheduled onto the main + # loop. Raising here is converted to the private sentinel below so + # GEPA cannot swallow it into a later success. + try: + self._run_context.raise_if_cancelled() + except OptimizationRunContextError as e: + raise _GovernedRunAbort(str(e)) from e + # Run the evaluation in the main loop future = asyncio.run_coroutine_threadsafe( self._sampler.sample_and_score( @@ -143,6 +188,13 @@ def evaluate( ) result: UnstructuredSamplingResult = future.result() + if self._run_context is not None: + # Boundary: immediately after the sampler evaluation settles. + try: + self._run_context.raise_if_cancelled() + except OptimizationRunContextError as e: + raise _GovernedRunAbort(str(e)) from e + scores = [] outputs = [] trajectories = [] @@ -204,6 +256,17 @@ def __init__( llm_registry = LLMRegistry() self._llm_class = llm_registry.resolve(self._config.optimizer_model) + @property + def capabilities(self) -> OptimizerCapabilities: + return OptimizerCapabilities( + accepts_run_context=True, + model_calls_observable=True, + logical_call_limits_enforceable=True, + reported_token_limits_enforceable=True, + cooperative_cancellation=True, + sampler_usage_included=False, + ) + async def optimize( self, initial_agent: Agent, @@ -224,114 +287,338 @@ async def optimize( agent instance, its scores on the validation examples, and other metrics. """ if run_context is not None: - from ._run_context import UnsupportedOptimizationContextError - - raise UnsupportedOptimizationContextError( - "GEPARootAgentPromptOptimizer does not support OptimizationRunContext" - " in this release; check optimizer.capabilities before passing a" - " context." - ) - - if initial_agent.sub_agents: - _logger.warning( - "The GEPARootAgentPromptOptimizer will not optimize prompts for" - " sub-agents." - ) - - _logger.info("Setting up the GEPA optimizer...") + run_context.attach(owner=self) try: - import gepa # lazy import as gepa is not in core ADK package + if run_context is not None: + # A pending cancellation wins over everything -- observed before + # imports, model construction, adapter construction, and sampler + # discovery. + run_context.raise_if_cancelled() + + if initial_agent.sub_agents: + _logger.warning( + "The GEPARootAgentPromptOptimizer will not optimize prompts for" + " sub-agents." + ) - _AgentGEPAAdapter = _create_agent_gepa_adapter_class() - except ImportError as e: - raise ImportError(MISSING_EVAL_DEPENDENCIES_MESSAGE) from e + _logger.info("Setting up the GEPA optimizer...") - loop = asyncio.get_running_loop() + try: + import gepa # lazy import as gepa is not in core ADK package - adapter = _AgentGEPAAdapter( - initial_agent=initial_agent, - sampler=sampler, - main_loop=loop, - ) + _AgentGEPAAdapter = _create_agent_gepa_adapter_class() + except ImportError as e: + raise ImportError(MISSING_EVAL_DEPENDENCIES_MESSAGE) from e - llm = self._llm_class(model=self._config.optimizer_model) + loop = asyncio.get_running_loop() - def reflection_lm(prompt: str) -> str: - llm_request = LlmRequest( - model=self._config.optimizer_model, - config=self._config.model_configuration, - contents=[ - genai_types.Content( - parts=[genai_types.Part(text=prompt)], - role="user", - ) - ], + adapter = _AgentGEPAAdapter( + initial_agent=initial_agent, + sampler=sampler, + main_loop=loop, + run_context=run_context, ) - async def _generate(): - response_text = "" - async with Aclosing(llm.generate_content_async(llm_request)) as agen: - async for llm_response in agen: - llm_response: LlmResponse - generated_content: genai_types.Content = llm_response.content - if not generated_content.parts: - continue - response_text = "".join( - part.text - for part in generated_content.parts - if part.text and not part.thought + llm = self._llm_class(model=self._config.optimizer_model) + + def reflection_lm(prompt: str) -> str: + # Build and validate the request BEFORE reserving the model-call slot: a + # construction failure is a failed proposal (GEPA-native no-proposal + # path) with a clean ledger -- no admitted call is left open. + llm_request = LlmRequest( + model=self._config.optimizer_model, + config=self._config.model_configuration, + contents=[ + genai_types.Content( + parts=[genai_types.Part(text=prompt)], + role="user", + ) + ], + ) + + capture = _ReflectionCapture() + + async def _generate() -> str: + response_text = "" + async with Aclosing(llm.generate_content_async(llm_request)) as agen: + async for llm_response in agen: + llm_response: LlmResponse + if getattr(llm_response, "usage_metadata", None) is not None: + capture.usage = llm_response.usage_metadata + if getattr(llm_response, "model_version", None): + capture.model_version = llm_response.model_version + error_code = getattr(llm_response, "error_code", None) + if error_code and run_context is not None: + # Governed-only: terminate on the in-band provider error. With + # no context, legacy stream semantics continue iterating. + capture.error_code = str(error_code) + return response_text + generated_content: genai_types.Content = llm_response.content + if not generated_content.parts: + continue + response_text = "".join( + part.text + for part in generated_content.parts + if part.text and not part.thought + ) + return response_text + + if run_context is None: + # Absent-context parity: the legacy scheduling/return path exactly. + future = asyncio.run_coroutine_threadsafe(_generate(), loop) + return future.result() + + # One logical optimizer-owned reflection invocation. Atomic admission + # observes cancellation and the call budget; from admission onward, + # every operation lives inside one settlement boundary so no exception + # can leave the call open. + try: + handle = run_context.begin_model_call( + STAGE_REFLECTION, + requested_model=self._config.optimizer_model, + ) + except OptimizationRunContextError as e: + raise _GovernedRunAbort(str(e)) from e + + try: + future = asyncio.run_coroutine_threadsafe(_generate(), loop) + except Exception as e: + # Local scheduling failure: settle as an aborted call -> run FAILED + # (OptimizationFailedError), truthfully distinct from a provider + # failure and from cancellation. + try: + run_context.abort_model_call( + handle, + error_code="SCHEDULING_FAILURE", + error_type=type(e).__name__, ) + except OptimizationRunContextError as ctx_err: + raise _GovernedRunAbort(str(ctx_err)) from ctx_err + raise _GovernedRunAbort(str(e)) from e + + try: + response_text = future.result() + except Exception as e: + try: + run_context.end_model_call( + handle, + usage_metadata=capture.usage, + error_code=str( + getattr(e, "error_code", None) or "PROVIDER_EXCEPTION" + ), + error_type=type(e).__name__, + ) + except OptimizationRunContextError as ctx_err: + raise _GovernedRunAbort(str(ctx_err)) from ctx_err + raise _GovernedRunAbort(str(e)) from e + + try: + if capture.error_code: + # In-band provider error: commit sanitized failure atomically + # (run status FAILED, provider-error-first precedence). + run_context.end_model_call( + handle, + usage_metadata=capture.usage, + error_code=capture.error_code, + error_type="LlmResponseError", + ) + else: + # Commit-then-terminate: an over-budget final reflection is + # persisted first (call status completed; run status + # budget_exceeded). + run_context.end_model_call( + handle, + usage_metadata=capture.usage, + returned_model_version=capture.model_version, + ) + except OptimizationRunContextError as e: + raise _GovernedRunAbort(str(e)) from e return response_text - future = asyncio.run_coroutine_threadsafe(_generate(), loop) - return future.result() + train_ids = sampler.get_train_example_ids() + val_ids = sampler.get_validation_example_ids() - train_ids = sampler.get_train_example_ids() - val_ids = sampler.get_validation_example_ids() + if set(train_ids).intersection(val_ids): + _logger.warning( + "The training and validation example UIDs overlap. This WILL cause" + " aliasing issues unless each common UID refers to the same example" + " in both sets." + ) - if set(train_ids).intersection(val_ids): - _logger.warning( - "The training and validation example UIDs overlap. This WILL cause" - " aliasing issues unless each common UID refers to the same example" - " in both sets." - ) + def run_gepa(): + if run_context is None: + # Legacy call shape preserved exactly for absent-context parity. + return gepa.optimize( + seed_candidate={_AGENT_PROMPT_NAME: initial_agent.instruction}, + trainset=train_ids, + valset=val_ids, + adapter=adapter, + max_metric_calls=self._config.max_metric_calls, + reflection_lm=reflection_lm, + reflection_minibatch_size=self._config.reflection_minibatch_size, + run_dir=self._config.run_dir, + ) - def run_gepa(): - return gepa.optimize( - seed_candidate={_AGENT_PROMPT_NAME: initial_agent.instruction}, - trainset=train_ids, - valset=val_ids, - adapter=adapter, - max_metric_calls=self._config.max_metric_calls, - reflection_lm=reflection_lm, - reflection_minibatch_size=self._config.reflection_minibatch_size, - run_dir=self._config.run_dir, - ) + # Governed: the stopper is observed at every loop boundary and must + # reflect ANY terminal governed state (budget, provider failure, + # cancellation), because GEPA's reflective-mutation proposer converts an + # aborted reflection into "no proposal" and keeps looping. + # raise_on_exception=True is pinned so generic evaluator/sampler failures + # stay fail-closed even if the upstream default changes under the + # unbounded gepa>=0.1 range. GEPA types stay private to ADK. + def _governed_stop(*_args: Any, **_kwargs: Any) -> bool: + return ( + run_context.cancel_requested + or run_context.snapshot().run_status is not None + ) - _logger.info("Running the GEPA optimizer...") + return gepa.optimize( + seed_candidate={_AGENT_PROMPT_NAME: initial_agent.instruction}, + trainset=train_ids, + valset=val_ids, + adapter=adapter, + max_metric_calls=self._config.max_metric_calls, + reflection_lm=reflection_lm, + reflection_minibatch_size=self._config.reflection_minibatch_size, + run_dir=self._config.run_dir, + stop_callbacks=[_governed_stop], + raise_on_exception=True, + ) - ctx = contextvars.copy_context() - gepa_results = await loop.run_in_executor(None, lambda: ctx.run(run_gepa)) + _logger.info("Running the GEPA optimizer...") - _logger.info("GEPA optimization finished. Preparing final results...") + ctx = contextvars.copy_context() + if run_context is None: + # Absent-context parity: the legacy direct await is preserved exactly, + # including immediate task-cancellation propagation semantics. + gepa_results = await loop.run_in_executor( + None, lambda: ctx.run(run_gepa) + ) + else: + executor_future = loop.run_in_executor(None, lambda: ctx.run(run_gepa)) + try: + gepa_results = await asyncio.shield(executor_future) + except asyncio.CancelledError: + # Governed native cancellation: request the cooperative stop, drain + # the executor worker once the active operation reaches a boundary, + # finalize the snapshot, then re-raise promptly. The caller's + # watchdog remains the bounded hard backstop. + run_context.request_cancel("task_cancelled") + try: + await asyncio.shield(executor_future) + except BaseException: # pylint: disable=broad-except + pass + run_context.finalize_cancelled("task_cancelled") + raise + except _GovernedRunAbort: + # The sentinel escaped gepa.optimize (e.g. raised during the seed + # evaluation, outside GEPA's swallowing loop). The terminal state is + # already committed; map it to the typed outcome -- the raw sentinel + # never reaches the caller. + snap = run_context.snapshot() + if snap.run_status == RunStatus.FAILED: + if snap.terminal_from_provider_call: + raise OptimizationProviderError( + f"Provider failure: {snap.terminal_error_code}", snap + ) from None + raise OptimizationFailedError( + f"Optimization failed: {snap.terminal_error_code}", snap + ) from None + if snap.run_status == RunStatus.BUDGET_EXCEEDED: + if run_context.budgets.on_budget_exceeded == "return_partial": + _logger.warning( + "Optimization budget exhausted; returning the seed as the" + " partial result." + ) + return GEPARootAgentPromptOptimizerResult( + optimized_agents=[ + AgentWithScores(optimized_agent=initial_agent) + ], + ) + raise OptimizationBudgetExceeded( + "Optimization budget exhausted.", snap + ) from None + run_context.raise_if_cancelled() + raise # pragma: no cover -- unknown sentinel state + except Exception as e: + run_context.finalize_failed( + error_code=getattr(e, "error_code", None) or "OPTIMIZER_FAILURE", + error_type=type(e).__name__, + ) + raise + + # Post-check: GEPA returns normally when the stopper ends the loop, so + # a committed governed terminal state must map to its typed outcome + # here -- never a false success. + snap = run_context.snapshot() + if snap.run_status == RunStatus.FAILED: + if snap.terminal_from_provider_call: + raise OptimizationProviderError( + f"Provider failure: {snap.terminal_error_code}", snap + ) + raise OptimizationFailedError( + f"Optimization failed: {snap.terminal_error_code}", snap + ) + if snap.cancel_requested and snap.run_status in ( + None, + RunStatus.CANCELLED, + ): + run_context.raise_if_cancelled() + if snap.run_status == RunStatus.BUDGET_EXCEEDED: + if run_context.budgets.on_budget_exceeded == "return_partial": + # GEPA stopped at a loop boundary, so its best-so-far candidates + # are available. The snapshot (run_status=budget_exceeded) is the + # authoritative record; the result schema stays unchanged and a + # downstream gate treats the attempt as never promotion-eligible. + _logger.warning( + "Optimization budget exhausted; returning partial result." + ) + else: + raise OptimizationBudgetExceeded( + "Optimization budget exhausted.", snap + ) - optimized_prompts = [ - candidate[_AGENT_PROMPT_NAME] for candidate in gepa_results.candidates - ] - scores = gepa_results.val_aggregate_scores + _logger.info("GEPA optimization finished. Preparing final results...") - optimized_agents = [ - AgentWithScores( - optimized_agent=initial_agent.clone( - update={"instruction": optimized_prompt}, - ), - overall_score=score, - ) - for optimized_prompt, score in zip(optimized_prompts, scores) - ] + optimized_prompts = [ + candidate[_AGENT_PROMPT_NAME] for candidate in gepa_results.candidates + ] + scores = gepa_results.val_aggregate_scores - return GEPARootAgentPromptOptimizerResult( - optimized_agents=optimized_agents, - gepa_result=gepa_results.to_dict(), - ) + optimized_agents = [ + AgentWithScores( + optimized_agent=initial_agent.clone( + update={"instruction": optimized_prompt}, + ), + overall_score=score, + ) + for optimized_prompt, score in zip(optimized_prompts, scores) + ] + + if run_context is not None and run_context.snapshot().run_status is None: + # Only a genuinely non-terminal run may finalize as success; a + # governed budget partial keeps its authoritative BUDGET_EXCEEDED + # terminal on the snapshot. + run_context.finalize_success() + return GEPARootAgentPromptOptimizerResult( + optimized_agents=optimized_agents, + gepa_result=gepa_results.to_dict(), + ) + except asyncio.CancelledError: + if run_context is not None: + run_context.finalize_cancelled("task_cancelled") + raise + except OptimizationRunContextError: + # Typed run-context outcomes already committed their terminal state. + raise + except Exception as e: + if run_context is not None: + # Any failure after successful attachment -- including sampler ID + # setup before the executor boundary -- finalizes the authoritative + # snapshot (first terminal wins). + run_context.finalize_failed( + error_code=getattr(e, "error_code", None) or "OPTIMIZER_FAILURE", + error_type=type(e).__name__, + ) + raise diff --git a/src/google/adk/optimization/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index e98c844c3ff..9b1475ec697 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -16,6 +16,7 @@ from __future__ import annotations +import asyncio import logging import random from typing import Optional @@ -24,7 +25,12 @@ from google.adk.evaluation._retry_options_utils import add_default_retry_options_if_not_present from google.adk.models.llm_request import LlmRequest from google.adk.models.registry import LLMRegistry +from google.adk.optimization._run_context import OptimizationBudgetExceeded +from google.adk.optimization._run_context import OptimizationProviderError from google.adk.optimization._run_context import OptimizationRunContext +from google.adk.optimization._run_context import OptimizerCapabilities +from google.adk.optimization._run_context import RunStatus +from google.adk.optimization._run_context import STAGE_CANDIDATE_GENERATION from google.adk.optimization.agent_optimizer import AgentOptimizer from google.adk.optimization.data_types import AgentWithScores from google.adk.optimization.data_types import OptimizerResult @@ -96,7 +102,10 @@ def __init__(self, config: SimplePromptOptimizerConfig): self._llm = llm_registry.new_llm(self._config.optimizer_model) async def _generate_candidate_prompt( - self, best_agent: Agent, best_score: float + self, + best_agent: Agent, + best_score: float, + run_context: Optional[OptimizationRunContext] = None, ) -> str: """Generates a new prompt candidate using the optimizer LLM.""" prompt_for_optimizer = _OPTIMIZER_PROMPT_TEMPLATE.format( @@ -115,13 +124,71 @@ async def _generate_candidate_prompt( ) add_default_retry_options_if_not_present(llm_request) + # One logical optimizer-owned model invocation: atomic admission checks + # cancellation and the call budget before the call starts; the terminal + # outcome (usage, provider error, or cancellation) is committed after. + handle = None + if run_context is not None: + handle = run_context.begin_model_call( + STAGE_CANDIDATE_GENERATION, + requested_model=self._config.optimizer_model, + ) + response_text = "" - async for llm_response in self._llm.generate_content_async(llm_request): - if not (llm_response.content and llm_response.content.parts): - continue - for part in llm_response.content.parts: - if part.text and not part.thought: - response_text += part.text + last_usage = None + model_version = None + try: + async for llm_response in self._llm.generate_content_async(llm_request): + if getattr(llm_response, "usage_metadata", None) is not None: + last_usage = llm_response.usage_metadata + if getattr(llm_response, "model_version", None): + model_version = llm_response.model_version + error_code = getattr(llm_response, "error_code", None) + if error_code and run_context is not None and handle is not None: + # Governed runs preserve usage-so-far and terminate on an in-band + # provider error instead of silently succeeding. end_model_call + # transitions the run to FAILED and raises the typed error. + run_context.end_model_call( + handle, + usage_metadata=last_usage, + error_code=str(error_code), + error_type=type(llm_response).__name__, + ) + if not (llm_response.content and llm_response.content.parts): + continue + for part in llm_response.content.parts: + if part.text and not part.thought: + response_text += part.text + except asyncio.CancelledError: + # Native task cancellation: settle the open call WITH the usage + # evidence accumulated so far, finalize the governed run as cancelled, + # and re-raise promptly. + if run_context is not None: + if handle is not None: + run_context.end_model_call( + handle, usage_metadata=last_usage, cancelled=True + ) + run_context.finalize_cancelled("task_cancelled") + raise + except OptimizationProviderError: + raise + except Exception as e: + if run_context is not None and handle is not None: + run_context.end_model_call( + handle, + usage_metadata=last_usage, + error_code=str( + getattr(e, "error_code", None) or "PROVIDER_EXCEPTION" + ), + error_type=type(e).__name__, + ) + raise + if run_context is not None and handle is not None: + run_context.end_model_call( + handle, + usage_metadata=last_usage, + returned_model_version=model_version, + ) return response_text async def _score_agent_on_batch( @@ -144,9 +211,14 @@ async def _run_optimization_iterations( initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], train_example_ids: list[str], + run_context: Optional[OptimizationRunContext] = None, ) -> tuple[Agent, float]: """Runs the optimization loop and returns the best agent and score.""" best_agent = initial_agent + if run_context is not None: + # Boundary: before the baseline sampler evaluation. A pre-cancelled + # run performs zero sampler work. + run_context.raise_if_cancelled() logger.info("Evaluating initial agent to get baseline score...") best_score = await self._score_agent_on_batch( best_agent, sampler, train_example_ids @@ -154,14 +226,35 @@ async def _run_optimization_iterations( logger.info("Initial agent baseline score: %f", best_score) for i in range(self._config.num_iterations): + if run_context is not None: + # Boundary: after the previous sampler operation / before this + # iteration's model call. + run_context.raise_if_cancelled() logger.info( "--- Starting optimization iteration %d/%d ---", i + 1, self._config.num_iterations, ) - new_prompt_text = await self._generate_candidate_prompt( - best_agent, best_score - ) + try: + new_prompt_text = await self._generate_candidate_prompt( + best_agent, best_score, run_context + ) + except OptimizationBudgetExceeded: + if ( + run_context is not None + and run_context.budgets.on_budget_exceeded == "return_partial" + ): + # Stop scheduling immediately after the final in-flight call + # settled; the best-so-far agent survives as the partial result. + logger.warning( + "Optimization budget exhausted at iteration %d; stopping.", + i + 1, + ) + return best_agent, best_score + raise + if run_context is not None: + # Boundary: after candidate generation / before candidate scoring. + run_context.raise_if_cancelled() candidate_agent = best_agent.clone( update={"instruction": new_prompt_text} ) @@ -169,6 +262,9 @@ async def _run_optimization_iterations( candidate_score = await self._score_agent_on_batch( candidate_agent, sampler, train_example_ids ) + if run_context is not None: + # Boundary: after candidate scoring. + run_context.raise_if_cancelled() logger.info( "Candidate score: %f (vs. best score: %f)", candidate_score, @@ -201,6 +297,17 @@ async def _run_final_validation( validation_results.scores ) + @property + def capabilities(self) -> OptimizerCapabilities: + return OptimizerCapabilities( + accepts_run_context=True, + model_calls_observable=True, + logical_call_limits_enforceable=True, + reported_token_limits_enforceable=True, + cooperative_cancellation=True, + sampler_usage_included=False, + ) + async def optimize( self, initial_agent: Agent, @@ -209,36 +316,72 @@ async def optimize( run_context: Optional[OptimizationRunContext] = None, ) -> OptimizerResult[AgentWithScores]: if run_context is not None: - from ._run_context import UnsupportedOptimizationContextError + run_context.attach(owner=self) - raise UnsupportedOptimizationContextError( - "SimplePromptOptimizer does not support OptimizationRunContext in" - " this release; check optimizer.capabilities before passing a" - " context." - ) + try: + if run_context is not None: + # A pending cancellation wins over sampler discovery -- observed + # before any sampler method is touched. + run_context.raise_if_cancelled() + train_example_ids = sampler.get_train_example_ids() - train_example_ids = sampler.get_train_example_ids() + if self._config.batch_size > len(train_example_ids): + logger.warning( + "Batch size (%d) is larger than the number of training examples" + " (%d). Using all training examples for each evaluation.", + self._config.batch_size, + len(train_example_ids), + ) + self._config.batch_size = len(train_example_ids) - if self._config.batch_size > len(train_example_ids): - logger.warning( - "Batch size (%d) is larger than the number of training examples" - " (%d). Using all training examples for each evaluation.", - self._config.batch_size, - len(train_example_ids), + best_agent, _ = await self._run_optimization_iterations( + initial_agent, sampler, train_example_ids, run_context ) - self._config.batch_size = len(train_example_ids) - best_agent, _ = await self._run_optimization_iterations( - initial_agent, sampler, train_example_ids - ) + if ( + run_context is not None + and run_context.snapshot().run_status == RunStatus.BUDGET_EXCEEDED + ): + # return_partial stop: no further optimizer-owned or sampler work is + # scheduled -- final validation included, so overall_score stays + # None. The caller-owned snapshot (run_status=budget_exceeded) is + # the authoritative record; the public result schema is unchanged. + return OptimizerResult( + optimized_agents=[ + AgentWithScores(optimized_agent=best_agent, overall_score=None) + ], + ) - final_score = await self._run_final_validation(best_agent, sampler) - logger.info("Final validation score: %f", final_score) + if run_context is not None: + # Boundary: before final validation. + run_context.raise_if_cancelled() + final_score = await self._run_final_validation(best_agent, sampler) + logger.info("Final validation score: %f", final_score) - return OptimizerResult( - optimized_agents=[ - AgentWithScores( - optimized_agent=best_agent, overall_score=final_score - ) - ] - ) + if run_context is not None: + # Boundary: after final validation -- cancellation requested during + # validation must not end as COMPLETED (finalize_success is also + # defensive, but the explicit boundary keeps evidence and ordering + # clear). + run_context.raise_if_cancelled() + run_context.finalize_success() + return OptimizerResult( + optimized_agents=[ + AgentWithScores( + optimized_agent=best_agent, overall_score=final_score + ) + ] + ) + except asyncio.CancelledError: + if run_context is not None: + run_context.finalize_cancelled("task_cancelled") + raise + except Exception as e: + if run_context is not None: + # Non-provider failures (sampler crashes, etc.) still finalize the + # authoritative snapshot; typed run-context errors already did. + run_context.finalize_failed( + error_code=getattr(e, "error_code", None) or "OPTIMIZER_FAILURE", + error_type=type(e).__name__, + ) + raise diff --git a/tests/unittests/optimization/run_context_enforcement_test.py b/tests/unittests/optimization/run_context_enforcement_test.py new file mode 100644 index 00000000000..2c42abdf1be --- /dev/null +++ b/tests/unittests/optimization/run_context_enforcement_test.py @@ -0,0 +1,791 @@ +# 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 +# +# http://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. + +"""Enforcement tests: run-context instrumentation in the P0a optimizers.""" + +from __future__ import annotations + +import asyncio +import sys +from types import SimpleNamespace +from unittest import mock + +from google.adk.agents.llm_agent import Agent +from google.adk.optimization import ContextAlreadyAttachedError +from google.adk.optimization import OptimizationBudgetExceeded +from google.adk.optimization import OptimizationBudgets +from google.adk.optimization import OptimizationCancelledError +from google.adk.optimization import OptimizationProviderError +from google.adk.optimization import OptimizationRunContext +from google.adk.optimization import RunStatus +from google.adk.optimization import STAGE_CANDIDATE_GENERATION +from google.adk.optimization import UsageCoverage +from google.adk.optimization.data_types import UnstructuredSamplingResult +from google.adk.optimization.sampler import Sampler +from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer +from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig +from google.genai import types as genai_types +import pytest + + +def _mock_sampler() -> mock.MagicMock: + sampler = mock.MagicMock(spec=Sampler) + sampler.get_train_example_ids.return_value = ["e1", "e2"] + sampler.get_validation_example_ids.return_value = ["v1"] + + async def mock_sample_and_score( + agent, example_set="validation", batch=None, capture_full_eval_data=False + ): + ids = batch or (["v1"] if example_set == "validation" else ["e1", "e2"]) + return UnstructuredSamplingResult(scores={i: 0.5 for i in ids}) + + sampler.sample_and_score.side_effect = mock_sample_and_score + return sampler + + +@pytest.fixture +def mock_sampler() -> mock.MagicMock: + return _mock_sampler() + + +def _mock_llm(usage=None, error_code=None, raise_exc=None): + mock_llm = mock.MagicMock() + + async def mock_generate_content_async(*args, **kwargs): + if raise_exc is not None: + raise raise_exc + yield SimpleNamespace( + content=genai_types.Content( + parts=[genai_types.Part(text="improved prompt")], role="model" + ), + usage_metadata=usage, + model_version="test-model-001", + error_code=error_code, + partial=False, + ) + + mock_llm.generate_content_async.side_effect = mock_generate_content_async + return mock.MagicMock(return_value=mock_llm) + + +def _optimizer(mock_llm_class, iterations=2): + with mock.patch( + "google.adk.optimization.simple_prompt_optimizer.LLMRegistry.resolve", + return_value=mock_llm_class, + ): + return SimplePromptOptimizer( + SimplePromptOptimizerConfig(num_iterations=iterations, batch_size=1) + ) + + +def _agent() -> Agent: + return Agent(name="test_agent", model="gemini-2.5-flash", instruction="v0") + + +class TestSimplePromptOptimizerEnforcement: + + def test_capabilities_opt_in(self): + caps = _optimizer(_mock_llm()).capabilities + assert caps.accepts_run_context + assert caps.model_calls_observable + assert caps.logical_call_limits_enforceable + assert caps.reported_token_limits_enforceable + assert caps.cooperative_cancellation + assert not caps.sampler_usage_included + + @pytest.mark.asyncio + async def test_normal_success_finalizes_completed(self, mock_sampler): + usage = SimpleNamespace(total_token_count=15) + optimizer = _optimizer(_mock_llm(usage), iterations=3) + ctx = OptimizationRunContext() + result = await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + # One candidate-generation call per iteration; nothing more. + assert snap.started_calls == 3 + assert snap.completed_calls == 3 + assert all(e.stage == STAGE_CANDIDATE_GENERATION for e in snap.events) + assert snap.cumulative_total_tokens == 45 + assert snap.usage_coverage == UsageCoverage.VERIFIED + # Normal success is finalized COMPLETED -- the snapshot is authoritative. + assert snap.run_status == RunStatus.COMPLETED + assert result.optimized_agents[0].overall_score is not None + + @pytest.mark.asyncio + async def test_missing_usage_is_unreported_not_zero(self, mock_sampler): + optimizer = _optimizer(_mock_llm(usage=None), iterations=1) + ctx = OptimizationRunContext() + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + event = ctx.snapshot().events[0] + assert event.total_tokens is None + assert event.usage_coverage == UsageCoverage.UNREPORTED + + @pytest.mark.asyncio + async def test_call_budget_raise_mode(self, mock_sampler): + optimizer = _optimizer(_mock_llm(), iterations=5) + ctx = OptimizationRunContext(OptimizationBudgets(max_model_calls=2)) + with pytest.raises(OptimizationBudgetExceeded): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + # Exactly N calls started; the rejected third reservation is not an event. + assert ctx.snapshot().started_calls == 2 + assert len(ctx.snapshot().events) == 2 + + @pytest.mark.asyncio + async def test_token_budget_return_partial(self, mock_sampler): + usage = SimpleNamespace(total_token_count=100) + optimizer = _optimizer(_mock_llm(usage), iterations=5) + ctx = OptimizationRunContext( + OptimizationBudgets( + max_provider_reported_tokens=150, + on_budget_exceeded="return_partial", + ) + ) + sampler_calls_before = mock_sampler.sample_and_score.await_count + result = await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + assert snap.cumulative_total_tokens == 200 + # Validation never ran, so overall_score is None; the result schema is + # unchanged and the snapshot is the authoritative record. + assert result.optimized_agents[0].overall_score is None + # Baseline + iteration-1 candidate scoring only: no scoring after stop. + assert mock_sampler.sample_and_score.await_count - sampler_calls_before == 2 + + @pytest.mark.asyncio + async def test_in_band_provider_error_terminates_governed_run( + self, mock_sampler + ): + usage = SimpleNamespace(total_token_count=30) + optimizer = _optimizer( + _mock_llm(usage, error_code="RESOURCE_EXHAUSTED"), iterations=2 + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + event = snap.events[0] + assert event.total_tokens == 30 # usage preserved + assert event.error_code == "RESOURCE_EXHAUSTED" + assert snap.run_status == RunStatus.FAILED + + @pytest.mark.asyncio + async def test_raised_provider_error_finalizes_failed(self, mock_sampler): + optimizer = _optimizer( + _mock_llm(raise_exc=RuntimeError("boom")), iterations=2 + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + assert snap.run_status == RunStatus.FAILED + assert snap.events[0].error_type == "RuntimeError" + # Sanitized: the raw exception text is not in the ledger. + assert snap.events[0].error_code == "PROVIDER_EXCEPTION" + + @pytest.mark.asyncio + async def test_pre_cancelled_run_does_zero_work(self, mock_sampler): + optimizer = _optimizer(_mock_llm(), iterations=5) + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + # Zero model calls AND zero sampler calls. + assert ctx.snapshot().started_calls == 0 + assert mock_sampler.sample_and_score.await_count == 0 + + @pytest.mark.asyncio + async def test_cancel_during_model_work_prevents_next_sampler_call( + self, mock_sampler + ): + ctx = OptimizationRunContext() + mock_llm = mock.MagicMock() + + async def gen_and_cancel(*args, **kwargs): + ctx.request_cancel("mid-model") + yield SimpleNamespace( + content=genai_types.Content( + parts=[genai_types.Part(text="p")], role="model" + ), + usage_metadata=None, + model_version=None, + error_code=None, + ) + + mock_llm.generate_content_async.side_effect = gen_and_cancel + optimizer = _optimizer(mock.MagicMock(return_value=mock_llm), iterations=3) + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + # Only the baseline scoring ran; the post-generation boundary stopped + # candidate scoring. + assert mock_sampler.sample_and_score.await_count == 1 + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + @pytest.mark.asyncio + async def test_native_cancellation_finalizes_and_reraises(self, mock_sampler): + ctx = OptimizationRunContext() + mock_llm = mock.MagicMock() + + async def gen_cancelled(*args, **kwargs): + raise asyncio.CancelledError() + yield # pylint: disable=unreachable + + mock_llm.generate_content_async.side_effect = gen_cancelled + optimizer = _optimizer(mock.MagicMock(return_value=mock_llm), iterations=2) + with pytest.raises(asyncio.CancelledError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + # The open call event is closed and the run finalized CANCELLED. + assert snap.run_status == RunStatus.CANCELLED + assert all(e.end_time is not None for e in snap.events) + + @pytest.mark.asyncio + async def test_no_context_path_unchanged(self, mock_sampler): + optimizer = _optimizer(_mock_llm(), iterations=2) + result = await optimizer.optimize(_agent(), mock_sampler) + assert result.optimized_agents + + @pytest.mark.asyncio + async def test_context_reuse_rejected_across_runs(self, mock_sampler): + optimizer = _optimizer(_mock_llm(), iterations=1) + ctx = OptimizationRunContext() + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + with pytest.raises(ContextAlreadyAttachedError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + + +class TestRootSkillOptimizerRejectsContext: + + @pytest.mark.asyncio + async def test_unsupported_context_rejected_before_work(self): + pytest.importorskip("gepa") + from google.adk.optimization import UnsupportedOptimizationContextError + from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizer + from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizerConfig + + optimizer = GEPARootAgentOptimizer(GEPARootAgentOptimizerConfig()) + sampler = mock.MagicMock(spec=Sampler) + with pytest.raises(UnsupportedOptimizationContextError): + await optimizer.optimize( + _agent(), sampler, run_context=OptimizationRunContext() + ) + # Rejected before any sampler or model work. + sampler.get_train_example_ids.assert_not_called() + sampler.sample_and_score.assert_not_called() + + def test_conservative_capabilities(self): + pytest.importorskip("gepa") + from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizer + from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizerConfig + + caps = GEPARootAgentOptimizer(GEPARootAgentOptimizerConfig()).capabilities + assert not caps.accepts_run_context + assert not caps.model_calls_observable + + +def _gepa_stub(mocker, evaluate_counter): + """A gepa stub with faithful engine semantics: per-iteration exceptions are + swallowed ("no proposal") and stop callbacks are checked at loop + boundaries.""" + + class MockEvaluationBatch: + + def __init__(self, outputs, scores, trajectories): + self.outputs, self.scores, self.trajectories = ( + outputs, + scores, + trajectories, + ) + + class MockGEPAAdapter: + + def __class_getitem__(cls, item): + return cls + + def fake_optimize( + *, + seed_candidate, + trainset, + valset, + adapter, + max_metric_calls, + reflection_lm, + reflection_minibatch_size, + run_dir, + stop_callbacks=None, + **_kwargs, + ): + def should_stop(): + return any(cb(None) for cb in stop_callbacks or []) + + for _ in range(5): + if should_stop(): + break + try: + reflection_lm("reflect") + except Exception: # engine.py:588 -- swallowed, loop continues + continue + if should_stop(): + break + evaluate_counter["count"] += 1 + adapter.evaluate(list(trainset), dict(seed_candidate), False) + return SimpleNamespace( + candidates=[dict(seed_candidate)], + val_aggregate_scores=[0.0], + to_dict=lambda: {}, + ) + + gepa_module = mocker.MagicMock() + gepa_module.optimize = fake_optimize + gepa_module.core.adapter.EvaluationBatch = MockEvaluationBatch + gepa_module.core.adapter.GEPAAdapter = MockGEPAAdapter + mocker.patch.dict( + sys.modules, + { + "gepa": gepa_module, + "gepa.core": gepa_module.core, + "gepa.core.adapter": gepa_module.core.adapter, + }, + ) + return gepa_module + + +def _gepa_optimizer(llm_class): + from google.adk.optimization.gepa_root_agent_prompt_optimizer import GEPARootAgentPromptOptimizer + from google.adk.optimization.gepa_root_agent_prompt_optimizer import GEPARootAgentPromptOptimizerConfig + + with mock.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer.LLMRegistry.resolve", + return_value=llm_class, + ): + return GEPARootAgentPromptOptimizer(GEPARootAgentPromptOptimizerConfig()) + + +class TestGepaSentinelBridge: + + @pytest.mark.asyncio + async def test_no_sampler_call_after_overshoot_commits(self, mocker): + counter = {"count": 0} + _gepa_stub(mocker, counter) + usage = SimpleNamespace(total_token_count=500) + optimizer = _gepa_optimizer(_mock_llm(usage)) + sampler = _mock_sampler() + ctx = OptimizationRunContext( + OptimizationBudgets( + max_provider_reported_tokens=100, + on_budget_exceeded="return_partial", + ) + ) + result = await optimizer.optimize(_agent(), sampler, run_context=ctx) + # The first reflection commits 500 tokens (> 100): the sentinel is + # swallowed by the engine, the stopper ends the loop at the next + # boundary, and NO sampler evaluation is ever scheduled. + assert counter["count"] == 0 + assert sampler.sample_and_score.call_count == 0 + snap = ctx.snapshot() + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + assert snap.completed_calls == 1 # committed; call status completed + assert result.optimized_agents # best-so-far partial, schema unchanged + + @pytest.mark.asyncio + async def test_raise_mode_maps_post_check_to_typed_error(self, mocker): + _gepa_stub(mocker, {"count": 0}) + usage = SimpleNamespace(total_token_count=500) + optimizer = _gepa_optimizer(_mock_llm(usage)) + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=100) + ) + with pytest.raises(OptimizationBudgetExceeded): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + + @pytest.mark.asyncio + async def test_in_band_provider_error_cannot_become_success(self, mocker): + _gepa_stub(mocker, {"count": 0}) + usage = SimpleNamespace(total_token_count=10) + optimizer = _gepa_optimizer( + _mock_llm(usage, error_code="RESOURCE_EXHAUSTED") + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + snap = ctx.snapshot() + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "RESOURCE_EXHAUSTED" + assert snap.events[0].total_tokens == 10 # usage preserved + + @pytest.mark.asyncio + async def test_raised_provider_error_cannot_become_success(self, mocker): + _gepa_stub(mocker, {"count": 0}) + optimizer = _gepa_optimizer(_mock_llm(raise_exc=RuntimeError("boom"))) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + assert ctx.snapshot().run_status == RunStatus.FAILED + + @pytest.mark.asyncio + async def test_normal_gepa_success_finalizes_completed(self, mocker): + counter = {"count": 0} + _gepa_stub(mocker, counter) + usage = SimpleNamespace(total_token_count=5) + optimizer = _gepa_optimizer(_mock_llm(usage)) + ctx = OptimizationRunContext() + result = await optimizer.optimize( + _agent(), _mock_sampler(), run_context=ctx + ) + assert result.optimized_agents + assert ctx.snapshot().run_status == RunStatus.COMPLETED + assert counter["count"] == 5 # all iterations ran + + @pytest.mark.asyncio + async def test_pre_cancelled_gepa_run_does_zero_work(self, mocker): + counter = {"count": 0} + _gepa_stub(mocker, counter) + optimizer = _gepa_optimizer(_mock_llm()) + sampler = _mock_sampler() + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + assert ctx.snapshot().started_calls == 0 + assert sampler.sample_and_score.call_count == 0 + + +class TestRealGepaIntegration: + """Budget/cancel/provider paths against the real installed gepa engine. + + These prove the sentinel/stopper/post-check bridge against actual GEPA + loop semantics, not the stub. A min-plus-latest gepa version matrix is a + CI concern tracked in the RFC. + """ + + @pytest.mark.asyncio + async def test_real_gepa_token_overshoot_raise_mode(self): + pytest.importorskip("gepa") + usage = SimpleNamespace(total_token_count=10_000) + optimizer = _gepa_optimizer(_mock_llm(usage)) + sampler = _mock_sampler() + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=100) + ) + with pytest.raises(OptimizationBudgetExceeded): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + snap = ctx.snapshot() + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + # The triggering reflection is committed with call status completed. + reflections = [e for e in snap.events if e.stage == "reflection"] + assert reflections and reflections[-1].state is not None + + @pytest.mark.asyncio + async def test_real_gepa_pre_cancelled_stops_promptly(self): + pytest.importorskip("gepa") + optimizer = _gepa_optimizer(_mock_llm()) + sampler = _mock_sampler() + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + assert ctx.snapshot().started_calls == 0 + + +class TestRoundTwoEnforcementFindings: + """Regression tests for the round-2 adversarial findings.""" + + @pytest.mark.asyncio + async def test_no_context_gepa_stream_keeps_legacy_semantics(self, mocker): + # An error-bearing chunk followed by a valid final chunk: the legacy + # (no-context) path must keep iterating and reflect the FINAL text. + reflected = {} + + def fake_optimize(*, reflection_lm, stop_callbacks=None, **_kwargs): + reflected["text"] = reflection_lm("reflect") + return SimpleNamespace( + candidates=[{"agent_prompt": "p1"}], + val_aggregate_scores=[0.0], + to_dict=lambda: {}, + ) + + gepa_module = mocker.MagicMock() + gepa_module.optimize = fake_optimize + + class MockEvaluationBatch: + + def __init__(self, outputs, scores, trajectories): + pass + + class MockGEPAAdapter: + + def __class_getitem__(cls, item): + return cls + + gepa_module.core.adapter.EvaluationBatch = MockEvaluationBatch + gepa_module.core.adapter.GEPAAdapter = MockGEPAAdapter + mocker.patch.dict( + sys.modules, + { + "gepa": gepa_module, + "gepa.core": gepa_module.core, + "gepa.core.adapter": gepa_module.core.adapter, + }, + ) + + mock_llm = mock.MagicMock() + + async def two_chunk_stream(*args, **kwargs): + yield SimpleNamespace( + content=genai_types.Content(parts=[], role="model"), + usage_metadata=None, + model_version=None, + error_code="TRANSIENT", + ) + yield SimpleNamespace( + content=genai_types.Content( + parts=[genai_types.Part(text="legacy final")], role="model" + ), + usage_metadata=None, + model_version=None, + error_code=None, + ) + + mock_llm.generate_content_async.side_effect = two_chunk_stream + optimizer = _gepa_optimizer(mock.MagicMock(return_value=mock_llm)) + await optimizer.optimize(_agent(), _mock_sampler()) # no run_context + assert reflected["text"] == "legacy final" + + @pytest.mark.asyncio + async def test_gepa_sampler_setup_failure_finalizes_failed(self, mocker): + _gepa_stub(mocker, {"count": 0}) + optimizer = _gepa_optimizer(_mock_llm()) + sampler = _mock_sampler() + sampler.get_train_example_ids.side_effect = ValueError("setup crash") + ctx = OptimizationRunContext() + with pytest.raises(ValueError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + # The attached lifecycle terminalizes even before the executor boundary. + assert ctx.snapshot().run_status == RunStatus.FAILED + assert ctx.snapshot().terminal_error_type == "ValueError" + + @pytest.mark.asyncio + async def test_simple_native_cancel_preserves_usage_evidence( + self, mock_sampler + ): + from google.adk.optimization import ModelCallState + + ctx = OptimizationRunContext() + mock_llm = mock.MagicMock() + + async def gen_usage_then_cancel(*args, **kwargs): + yield SimpleNamespace( + content=genai_types.Content( + parts=[genai_types.Part(text="p")], role="model" + ), + usage_metadata=SimpleNamespace(total_token_count=7), + model_version=None, + error_code=None, + ) + raise asyncio.CancelledError() + + mock_llm.generate_content_async.side_effect = gen_usage_then_cancel + optimizer = _optimizer(mock.MagicMock(return_value=mock_llm), iterations=2) + with pytest.raises(asyncio.CancelledError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + snap = ctx.snapshot() + # The seven reported tokens survive as evidence on the cancelled call. + assert snap.events[0].state == ModelCallState.CANCELLED + assert snap.events[0].total_tokens == 7 + assert snap.completed_calls == 1 # settled + assert snap.run_status == RunStatus.CANCELLED + + @pytest.mark.asyncio + async def test_cancel_during_final_validation_never_completes(self): + ctx = OptimizationRunContext() + sampler = mock.MagicMock(spec=Sampler) + sampler.get_train_example_ids.return_value = ["e1"] + sampler.get_validation_example_ids.return_value = ["v1"] + + async def sample_and_score( + agent, + example_set="validation", + batch=None, + capture_full_eval_data=False, + ): + if example_set == "validation": + # Cancellation lands while final validation is in flight. + ctx.request_cancel("during-validation") + ids = batch or (["v1"] if example_set == "validation" else ["e1"]) + return UnstructuredSamplingResult(scores={i: 0.5 for i in ids}) + + sampler.sample_and_score.side_effect = sample_and_score + optimizer = _optimizer(_mock_llm(), iterations=1) + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + @pytest.mark.asyncio + async def test_numeric_error_code_normalized_through_raised_path( + self, mock_sampler + ): + exc = RuntimeError("boom") + exc.error_code = 429 # numeric provider code + optimizer = _optimizer(_mock_llm(raise_exc=exc), iterations=1) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + # Normalized to a bounded string at both the adapter and the context. + assert ctx.snapshot().events[0].error_code == "429" + + +class TestRoundThreeEnforcementFindings: + """Regression tests for the round-3 adversarial findings.""" + + @pytest.mark.asyncio + async def test_reflection_request_failure_leaves_clean_ledger(self, mocker): + # Request construction happens BEFORE slot admission: a failure is a + # GEPA-native no-proposal with a clean ledger, never an open call, and + # finalize_success (second line of defense) can commit truthfully. + counter = {"count": 0} + _gepa_stub(mocker, counter) + mocker.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer.LlmRequest", + side_effect=RuntimeError("bad request config"), + ) + optimizer = _gepa_optimizer(_mock_llm()) + ctx = OptimizationRunContext() + result = await optimizer.optimize( + _agent(), _mock_sampler(), run_context=ctx + ) + snap = ctx.snapshot() + assert snap.started_calls == 0 # no admitted call was left open + assert snap.run_status == RunStatus.COMPLETED + assert result.optimized_agents + + @pytest.mark.asyncio + async def test_reflection_scheduling_failure_becomes_failed_not_success( + self, mocker + ): + from google.adk.optimization import OptimizationFailedError + + counter = {"count": 0} + _gepa_stub(mocker, counter) + optimizer = _gepa_optimizer(_mock_llm()) + mocker.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer.asyncio" + ".run_coroutine_threadsafe", + side_effect=RuntimeError("loop closed"), + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationFailedError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + snap = ctx.snapshot() + # The admitted call is settled (aborted), never open; the failure is + # generic, not a provider failure, and can never become success. + assert snap.run_status == RunStatus.FAILED + assert not snap.terminal_from_provider_call + assert snap.completed_calls == snap.started_calls == 1 + assert counter["count"] == 0 # no sampler evaluation after the abort + + @pytest.mark.asyncio + async def test_pre_cancel_wins_before_simple_discovery(self): + optimizer = _optimizer(_mock_llm(), iterations=2) + sampler = mock.MagicMock(spec=Sampler) + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + # Every sampler method remains untouched. + sampler.get_train_example_ids.assert_not_called() + sampler.get_validation_example_ids.assert_not_called() + sampler.sample_and_score.assert_not_called() + + @pytest.mark.asyncio + async def test_pre_cancel_wins_before_gepa_discovery(self, mocker): + _gepa_stub(mocker, {"count": 0}) + optimizer = _gepa_optimizer(_mock_llm()) + sampler = mock.MagicMock(spec=Sampler) + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + sampler.get_train_example_ids.assert_not_called() + sampler.get_validation_example_ids.assert_not_called() + + @pytest.mark.asyncio + async def test_pre_cancel_wins_over_discovery_exception(self): + # Cancellation pending + discovery that would raise: cancellation wins. + optimizer = _optimizer(_mock_llm(), iterations=2) + sampler = mock.MagicMock(spec=Sampler) + sampler.get_train_example_ids.side_effect = ValueError("discovery crash") + ctx = OptimizationRunContext() + ctx.request_cancel("pre") + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + +class TestRealGepaRoundThree: + """Round-3 durable real-GEPA coverage additions.""" + + @pytest.mark.asyncio + async def test_real_gepa_in_band_provider_error_is_typed(self): + pytest.importorskip("gepa") + usage = SimpleNamespace(total_token_count=10) + optimizer = _gepa_optimizer( + _mock_llm(usage, error_code="RESOURCE_EXHAUSTED") + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + snap = ctx.snapshot() + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_from_provider_call + + @pytest.mark.asyncio + async def test_real_gepa_post_sampler_cancellation_is_typed(self): + pytest.importorskip("gepa") + optimizer = _gepa_optimizer(_mock_llm()) + ctx = OptimizationRunContext() + sampler = mock.MagicMock(spec=Sampler) + sampler.get_train_example_ids.return_value = ["e1", "e2"] + sampler.get_validation_example_ids.return_value = ["v1"] + + async def sample_then_cancel( + agent, + example_set="validation", + batch=None, + capture_full_eval_data=False, + ): + ids = batch or (["v1"] if example_set == "validation" else ["e1", "e2"]) + ctx.request_cancel("post-sampler") + return UnstructuredSamplingResult(scores={i: 0.5 for i in ids}) + + sampler.sample_and_score.side_effect = sample_then_cancel + with pytest.raises(OptimizationCancelledError): + await optimizer.optimize(_agent(), sampler, run_context=ctx) + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + @pytest.mark.asyncio + async def test_real_gepa_reflection_request_failure_completes_cleanly( + self, mocker + ): + pytest.importorskip("gepa") + mocker.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer.LlmRequest", + side_effect=RuntimeError("bad request config"), + ) + optimizer = _gepa_optimizer(_mock_llm()) + ctx = OptimizationRunContext() + result = await optimizer.optimize( + _agent(), _mock_sampler(), run_context=ctx + ) + snap = ctx.snapshot() + assert snap.started_calls == 0 + assert snap.run_status == RunStatus.COMPLETED + assert result.optimized_agents diff --git a/tests/unittests/optimization/run_context_test.py b/tests/unittests/optimization/run_context_test.py index 5baa32feb1d..c1a57191711 100644 --- a/tests/unittests/optimization/run_context_test.py +++ b/tests/unittests/optimization/run_context_test.py @@ -380,152 +380,11 @@ def test_optimizer_result_has_no_run_context_fields(self): assert "terminal_status" not in OptimizerResult.model_fields -class TestBuiltInsRejectContextInContractsPr: - - @pytest.mark.asyncio - async def test_simple_prompt_optimizer_rejects_context(self): - from unittest import mock - - from google.adk.optimization import UnsupportedOptimizationContextError - from google.adk.optimization.sampler import Sampler - from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer - from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig - - with mock.patch( - "google.adk.optimization.simple_prompt_optimizer.LLMRegistry.resolve" - ): - optimizer = SimplePromptOptimizer(SimplePromptOptimizerConfig()) - sampler = mock.MagicMock(spec=Sampler) - with pytest.raises(UnsupportedOptimizationContextError): - await optimizer.optimize( - mock.MagicMock(), sampler, run_context=OptimizationRunContext() - ) - # Rejected before any sampler work. - sampler.get_train_example_ids.assert_not_called() - - -class TestAtomicSettlement: - """Round-2 adversarial findings: validate-before-commit.""" - - def test_nan_usage_does_not_corrupt_the_record(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - ctx.end_model_call(h, usage_metadata=_usage(total_token_count=float("nan"))) - event = ctx.snapshot().events[0] - # NaN/inf/negative counters are "not reported", never a crash or a - # half-closed event. - assert event.state == ModelCallState.COMPLETED - assert event.total_tokens is None - assert event.usage_coverage == UsageCoverage.UNREPORTED - - def test_negative_and_inf_usage_treated_as_unreported(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - ctx.end_model_call( - h, - usage_metadata=_usage( - total_token_count=float("inf"), prompt_token_count=-3 - ), - ) - event = ctx.snapshot().events[0] - assert event.total_tokens is None - assert event.prompt_tokens is None - - def test_numeric_error_code_is_normalized_to_string(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - with pytest.raises(OptimizationProviderError) as exc: - ctx.end_model_call(h, usage_metadata=None, error_code=429) - snap = exc.value.snapshot - assert snap.events[0].error_code == "429" - assert snap.terminal_error_code == "429" - - def test_error_metadata_is_sanitized_and_bounded(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - with pytest.raises(OptimizationProviderError) as exc: - ctx.end_model_call( - h, - error_code="multi\nline /etc/passwd " + "x" * 500, - error_type="/usr/lib/module.py", - ) - event = exc.value.snapshot.events[0] - assert "\n" not in event.error_code - assert "/" not in event.error_code - assert len(event.error_code) <= 128 - assert "/" not in event.error_type - - -class TestTerminalCompleteness: - """Round-2 adversarial findings: every caller observes the terminal.""" - - def test_finalize_success_observes_pending_cancellation(self): - ctx = OptimizationRunContext() - ctx.request_cancel("deadline") - with pytest.raises(OptimizationCancelledError): - ctx.finalize_success() - assert ctx.snapshot().run_status == RunStatus.CANCELLED - - def test_inflight_settlement_after_terminal_raises_existing_terminal(self): - ctx = OptimizationRunContext( - OptimizationBudgets(max_provider_reported_tokens=10) - ) - h1 = ctx.begin_model_call(STAGE_REFLECTION) - h2 = ctx.begin_model_call(STAGE_REFLECTION) # admitted before terminal - with pytest.raises(OptimizationBudgetExceeded): - ctx.end_model_call(h1, usage_metadata=_usage(total_token_count=50)) - # The second in-flight call settles for truthful accounting, but its - # caller receives the existing terminal and cannot continue. - with pytest.raises(OptimizationBudgetExceeded): - ctx.end_model_call(h2, usage_metadata=_usage(total_token_count=5)) - snap = ctx.snapshot() - assert snap.completed_calls == 2 # both settled - assert snap.cumulative_total_tokens == 55 # evidence preserved - assert snap.run_status == RunStatus.BUDGET_EXCEEDED - - def test_cancelled_settlement_preserves_usage_without_raising(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - ctx.end_model_call( - h, usage_metadata=_usage(total_token_count=7), cancelled=True - ) - snap = ctx.snapshot() - assert snap.events[0].state == ModelCallState.CANCELLED - assert snap.events[0].total_tokens == 7 - assert snap.completed_calls == 1 # settled calls, any call status - - -class TestFailureTypes: - - def test_finalize_failed_is_generic_not_provider(self): - from google.adk.optimization import OptimizationFailedError - - ctx = OptimizationRunContext() - ctx.finalize_failed(error_code="SAMPLER_CRASH", error_type="ValueError") - with pytest.raises(OptimizationFailedError): - ctx.begin_model_call(STAGE_REFLECTION) - - def test_provider_call_failure_stays_provider_typed(self): - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - with pytest.raises(OptimizationProviderError): - ctx.end_model_call(h, error_code="INTERNAL") - with pytest.raises(OptimizationProviderError): - ctx.begin_model_call(STAGE_REFLECTION) - - -class TestDurableTimestamps: - - def test_event_times_are_wall_clock(self): - import time as _time - - before = _time.time() - ctx = OptimizationRunContext() - h = ctx.begin_model_call(STAGE_REFLECTION) - ctx.end_model_call(h, usage_metadata=None) - event = ctx.snapshot().events[0] - after = _time.time() - assert before <= event.start_time <= event.end_time <= after +# The contracts-PR behavior (built-ins reject a supplied context) is +# superseded in this stacked enforcement PR: SimplePromptOptimizer and +# GEPARootAgentPromptOptimizer now support run contexts (see +# run_context_enforcement_test.py); GEPARootAgentOptimizer still rejects, +# covered there as well. class TestSuccessInvariantBoundary: From 8a19693d9f68b259c9ed6b1d37c0f9edc74d2b2b Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 02:50:05 -0700 Subject: [PATCH 9/9] round-4 enforcement: governed request-construction failures are FAILED, coroutine hygiene, terminal evidence parity - A governed LlmRequest construction failure keeps the call ledger correctly empty (no provider call began) but finalizes the run FAILED (REQUEST_CONSTRUCTION_FAILURE) and aborts via the sentinel -- a malformed optimizer request can never be recorded as a successful governed optimization. Legacy no-context behavior preserved. The two round-3 tests pinning the superseded clean-ledger-success shape are updated to the corrected semantics (stub + real GEPA). - Scheduling failures close the never-awaited coroutine before settling the abort (reflection and the adapter's sampler coroutine) -- the regression now runs with RuntimeWarning promoted to error. - Terminal settlements carry the same evidence as successes: in-band and raised provider errors and Simple native cancellation all preserve returned_model_version alongside usage (regressions per optimizer). - Full suite: 116/116. --- .../gepa_root_agent_prompt_optimizer.py | 65 ++++++---- .../optimization/simple_prompt_optimizer.py | 6 +- .../gepa_root_agent_optimizer_test.py | 1 - .../run_context_enforcement_test.py | 112 +++++++++++++++--- 4 files changed, 142 insertions(+), 42 deletions(-) diff --git a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py index 8befca35e3d..c4343e66725 100644 --- a/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py +++ b/src/google/adk/optimization/gepa_root_agent_prompt_optimizer.py @@ -177,15 +177,19 @@ def evaluate( raise _GovernedRunAbort(str(e)) from e # Run the evaluation in the main loop - future = asyncio.run_coroutine_threadsafe( - self._sampler.sample_and_score( - new_agent, - example_set=example_set, - batch=batch, - capture_full_eval_data=capture_traces, - ), - self._main_loop, + sample_coroutine = self._sampler.sample_and_score( + new_agent, + example_set=example_set, + batch=batch, + capture_full_eval_data=capture_traces, ) + try: + future = asyncio.run_coroutine_threadsafe( + sample_coroutine, self._main_loop + ) + except Exception: + sample_coroutine.close() + raise result: UnstructuredSamplingResult = future.result() if self._run_context is not None: @@ -323,19 +327,31 @@ async def optimize( llm = self._llm_class(model=self._config.optimizer_model) def reflection_lm(prompt: str) -> str: - # Build and validate the request BEFORE reserving the model-call slot: a - # construction failure is a failed proposal (GEPA-native no-proposal - # path) with a clean ledger -- no admitted call is left open. - llm_request = LlmRequest( - model=self._config.optimizer_model, - config=self._config.model_configuration, - contents=[ - genai_types.Content( - parts=[genai_types.Part(text=prompt)], - role="user", - ) - ], - ) + # Build and validate the request BEFORE reserving the model-call + # slot: the call ledger stays correctly empty (no provider call + # began). On the governed path a construction failure is still an + # optimizer-internal failure, NOT a successful optimization: + # finalize FAILED and abort so GEPA cannot convert it to a silent + # no-proposal success. Legacy no-context behavior is preserved. + try: + llm_request = LlmRequest( + model=self._config.optimizer_model, + config=self._config.model_configuration, + contents=[ + genai_types.Content( + parts=[genai_types.Part(text=prompt)], + role="user", + ) + ], + ) + except Exception as e: + if run_context is not None: + run_context.finalize_failed( + error_code="REQUEST_CONSTRUCTION_FAILURE", + error_type=type(e).__name__, + ) + raise _GovernedRunAbort(str(e)) from e + raise capture = _ReflectionCapture() @@ -381,9 +397,12 @@ async def _generate() -> str: except OptimizationRunContextError as e: raise _GovernedRunAbort(str(e)) from e + generate_coroutine = _generate() try: - future = asyncio.run_coroutine_threadsafe(_generate(), loop) + future = asyncio.run_coroutine_threadsafe(generate_coroutine, loop) except Exception as e: + # Close the never-awaited coroutine before settling the abort. + generate_coroutine.close() # Local scheduling failure: settle as an aborted call -> run FAILED # (OptimizationFailedError), truthfully distinct from a provider # failure and from cancellation. @@ -404,6 +423,7 @@ async def _generate() -> str: run_context.end_model_call( handle, usage_metadata=capture.usage, + returned_model_version=capture.model_version, error_code=str( getattr(e, "error_code", None) or "PROVIDER_EXCEPTION" ), @@ -420,6 +440,7 @@ async def _generate() -> str: run_context.end_model_call( handle, usage_metadata=capture.usage, + returned_model_version=capture.model_version, error_code=capture.error_code, error_type="LlmResponseError", ) diff --git a/src/google/adk/optimization/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index 9b1475ec697..be21b121152 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -151,6 +151,7 @@ async def _generate_candidate_prompt( run_context.end_model_call( handle, usage_metadata=last_usage, + returned_model_version=model_version, error_code=str(error_code), error_type=type(llm_response).__name__, ) @@ -166,7 +167,10 @@ async def _generate_candidate_prompt( if run_context is not None: if handle is not None: run_context.end_model_call( - handle, usage_metadata=last_usage, cancelled=True + handle, + usage_metadata=last_usage, + returned_model_version=model_version, + cancelled=True, ) run_context.finalize_cancelled("task_cancelled") raise diff --git a/tests/unittests/optimization/gepa_root_agent_optimizer_test.py b/tests/unittests/optimization/gepa_root_agent_optimizer_test.py index 70c1cdea17f..2958b1645f8 100644 --- a/tests/unittests/optimization/gepa_root_agent_optimizer_test.py +++ b/tests/unittests/optimization/gepa_root_agent_optimizer_test.py @@ -27,7 +27,6 @@ from google.adk.optimization.gepa_root_agent_optimizer import _update_skill_toolset from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizer from google.adk.optimization.gepa_root_agent_optimizer import GEPARootAgentOptimizerConfig -from google.adk.optimization.sampler import Sampler from google.adk.skills import models from google.adk.tools.skill_toolset import SkillToolset import pytest diff --git a/tests/unittests/optimization/run_context_enforcement_test.py b/tests/unittests/optimization/run_context_enforcement_test.py index 2c42abdf1be..92ebe7e2c81 100644 --- a/tests/unittests/optimization/run_context_enforcement_test.py +++ b/tests/unittests/optimization/run_context_enforcement_test.py @@ -647,10 +647,12 @@ class TestRoundThreeEnforcementFindings: """Regression tests for the round-3 adversarial findings.""" @pytest.mark.asyncio - async def test_reflection_request_failure_leaves_clean_ledger(self, mocker): - # Request construction happens BEFORE slot admission: a failure is a - # GEPA-native no-proposal with a clean ledger, never an open call, and - # finalize_success (second line of defense) can commit truthfully. + async def test_reflection_request_failure_is_failed_not_success(self, mocker): + # Request construction happens BEFORE slot admission, so the call ledger + # stays empty -- but a malformed request is an optimizer-internal + # failure, never a successful governed optimization (round-4). + from google.adk.optimization import OptimizationFailedError + counter = {"count": 0} _gepa_stub(mocker, counter) mocker.patch( @@ -659,13 +661,13 @@ async def test_reflection_request_failure_leaves_clean_ledger(self, mocker): ) optimizer = _gepa_optimizer(_mock_llm()) ctx = OptimizationRunContext() - result = await optimizer.optimize( - _agent(), _mock_sampler(), run_context=ctx - ) + with pytest.raises(OptimizationFailedError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) snap = ctx.snapshot() - assert snap.started_calls == 0 # no admitted call was left open - assert snap.run_status == RunStatus.COMPLETED - assert result.optimized_agents + assert snap.started_calls == 0 # clean ledger: no provider call began + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "REQUEST_CONSTRUCTION_FAILURE" + assert not snap.terminal_from_provider_call @pytest.mark.asyncio async def test_reflection_scheduling_failure_becomes_failed_not_success( @@ -772,20 +774,94 @@ async def sample_then_cancel( assert ctx.snapshot().run_status == RunStatus.CANCELLED @pytest.mark.asyncio - async def test_real_gepa_reflection_request_failure_completes_cleanly( - self, mocker - ): + async def test_real_gepa_reflection_request_failure_is_failed(self, mocker): pytest.importorskip("gepa") + from google.adk.optimization import OptimizationFailedError + mocker.patch( "google.adk.optimization.gepa_root_agent_prompt_optimizer.LlmRequest", side_effect=RuntimeError("bad request config"), ) optimizer = _gepa_optimizer(_mock_llm()) ctx = OptimizationRunContext() - result = await optimizer.optimize( - _agent(), _mock_sampler(), run_context=ctx - ) + with pytest.raises(OptimizationFailedError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) snap = ctx.snapshot() assert snap.started_calls == 0 - assert snap.run_status == RunStatus.COMPLETED - assert result.optimized_agents + assert snap.run_status == RunStatus.FAILED + assert snap.terminal_error_code == "REQUEST_CONSTRUCTION_FAILURE" + + +class TestRoundFourEnforcementFindings: + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("error::RuntimeWarning") + async def test_scheduling_failure_closes_coroutine(self, mocker): + # Promotes 'coroutine was never awaited' to an error: the abort path + # must close the never-scheduled coroutine. + from google.adk.optimization import OptimizationFailedError + + counter = {"count": 0} + _gepa_stub(mocker, counter) + optimizer = _gepa_optimizer(_mock_llm()) + mocker.patch( + "google.adk.optimization.gepa_root_agent_prompt_optimizer.asyncio" + ".run_coroutine_threadsafe", + side_effect=RuntimeError("loop closed"), + ) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationFailedError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + assert ctx.snapshot().run_status == RunStatus.FAILED + + @pytest.mark.asyncio + async def test_gepa_in_band_terminal_keeps_model_version(self): + pytest.importorskip("gepa") + usage = SimpleNamespace(total_token_count=7) + mock_llm = mock.MagicMock() + + async def stream(*args, **kwargs): + yield SimpleNamespace( + content=genai_types.Content(parts=[], role="model"), + usage_metadata=usage, + model_version="provider-model-v9", + error_code="RESOURCE_EXHAUSTED", + ) + + mock_llm.generate_content_async.side_effect = stream + optimizer = _gepa_optimizer(mock.MagicMock(return_value=mock_llm)) + ctx = OptimizationRunContext() + with pytest.raises(OptimizationProviderError): + await optimizer.optimize(_agent(), _mock_sampler(), run_context=ctx) + event = ctx.snapshot().events[0] + # Terminal attempt evidence is as strong as success evidence. + assert event.total_tokens == 7 + assert event.error_code == "RESOURCE_EXHAUSTED" + assert event.returned_model_version == "provider-model-v9" + + @pytest.mark.asyncio + async def test_simple_cancel_terminal_keeps_model_version(self, mock_sampler): + from google.adk.optimization import ModelCallState + + ctx = OptimizationRunContext() + mock_llm = mock.MagicMock() + + async def gen(*args, **kwargs): + yield SimpleNamespace( + content=genai_types.Content( + parts=[genai_types.Part(text="p")], role="model" + ), + usage_metadata=SimpleNamespace(total_token_count=7), + model_version="provider-model-v9", + error_code=None, + ) + raise asyncio.CancelledError() + + mock_llm.generate_content_async.side_effect = gen + optimizer = _optimizer(mock.MagicMock(return_value=mock_llm), iterations=1) + with pytest.raises(asyncio.CancelledError): + await optimizer.optimize(_agent(), mock_sampler, run_context=ctx) + event = ctx.snapshot().events[0] + assert event.state == ModelCallState.CANCELLED + assert event.total_tokens == 7 + assert event.returned_model_version == "provider-model-v9"