From e726c58e2d1b96e1ed3b6e88b8e092ca57b4d548 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Fri, 10 Jul 2026 02:46:33 -0700 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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"