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 new file mode 100644 index 00000000000..d1d02ba6410 --- /dev/null +++ b/src/google/adk/optimization/_run_context.py @@ -0,0 +1,922 @@ +# 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, +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 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 + +import enum +import math +import re +import threading +import time +from typing import Any +from typing import Literal +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.""" + + 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 ModelCallState(str, enum.Enum): + """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" + 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): + """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): + """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. 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: str + 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 + 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_code: Optional[str] = None + error_type: Optional[str] = None + + +class OptimizationBudgets(BaseModel): + """Configured ceilings for optimizer-owned logical model invocations. + + Immutable by construction, so limits cannot change mid-run. + """ + + 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 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; 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," + " actual-token compliance is indeterminate." + ), + ) + 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 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 + """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).""" + + 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 + 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): + """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 OptimizationProviderError(OptimizationRunContextError): + """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 + committed usage and the token-compliance evidence are preserved, but the + run terminates ``failed`` and this error is raised. + """ + + +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.""" + + +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 _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: 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: bool = False + + +class _CallHandle: + """Opaque handle for one in-flight logical model invocation. + + Bound to exactly one context; committing it into a different context is a + typed misuse error. + """ + + def __init__( + self, context: "OptimizationRunContext", record: _CallRecord + ) -> None: + self._context: "OptimizationRunContext" = context + self._record: _CallRecord = record + + +_UNATTACHED = object() + + +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 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) -> 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._records: list[_CallRecord] = [] + self._started_calls = 0 + self._settled_calls = 0 + self._cumulative_total_tokens = 0 + self._cancel_requested = False + self._cancel_reason: Optional[str] = None + self._run_status: Optional[RunStatus] = 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 + 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 _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. + + 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 + self._cancel_reason = ( + _sanitize_optional_str(reason, _ERROR_META_MAX_LEN) or "requested" + ) + + @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. + + 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 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 + raise error + + # --- terminal transitions ------------------------------------------------- + + def finalize_success(self) -> None: + """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: + 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 + raise error + + 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 = ( + _sanitize_optional_str(reason, _ERROR_META_MAX_LEN) or "requested" + ) + for record in self._records: + if not record.closed: + record.closed = True + record.end_time = time.time() + record.state = ModelCallState.CANCELLED + self._settled_calls += 1 + 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=_sanitize_error_meta(error_code), + error_type=_sanitize_error_meta(error_type), + ) + + def _transition_locked( + self, + status: RunStatus, + *, + 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: + return False + self._run_status = status + 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: + """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: + 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 + ) + status = self._run_status + status_name = status.value if status is not None else "unknown" + return OptimizationRunFinalizedError( + f"OptimizationRunContext already finalized as {status_name}." + ) + + # --- the ledger ----------------------------------------------------------- + + def begin_model_call( + self, + stage: str, + requested_model: Optional[str] = None, + ) -> _CallHandle: + """Atomically admits and records the start of one logical invocation. + + 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._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: + 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: + 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, + clean_stage, + _sanitize_optional_str(requested_model, _MODEL_VERSION_MAX_LEN), + ) + self._records.append(record) + return _CallHandle(self, record) + raise error + + def end_model_call( + self, + handle: _CallHandle, + *, + usage_metadata: Any = None, + 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. + + 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. ``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: + raise OptimizationRunFinalizedError( + "Call handle belongs to a different OptimizationRunContext." + ) + if record.closed: + # 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() + 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") + if total is not None: + self._cumulative_total_tokens += total + + error: Optional[Exception] = None + if cancelled: + record.state = ModelCallState.CANCELLED + elif is_provider_error: + record.state = ModelCallState.PROVIDER_ERROR + 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=clean_code, + error_type=clean_type, + provider=True, + ): + error = OptimizationProviderError( + f"Provider failure on call {record.sequence}: {clean_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 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 + + 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 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 + 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: + with self._lock: + return self._snapshot_locked() + + def _snapshot_locked(self) -> OptimizationRunSnapshot: + 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( + c in (UsageCoverage.VERIFIED, UsageCoverage.PARTIAL) for c in coverages + ): + 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 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=events, + budgets=self._budgets, + started_calls=self._started_calls, + completed_calls=self._settled_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, + run_status=self._run_status, + 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, + ) + + +_ERROR_META_MAX_LEN = 128 +_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. + + Defensive against hostile ``__str__``: a conversion failure degrades the + field to ``None`` rather than defeating settlement. + """ + if value is None: + return None + try: + text = str(value).strip() + except Exception: # pylint: disable=broad-except + return None + 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. + + 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 + try: + text = str(value).strip() + except Exception: # pylint: disable=broad-except + return "UNSTRINGABLE" + 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). + + 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]: + # 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 { + "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): + """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_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).""" + + model_calls_observable: bool = False + """Optimizer-owned logical model invocations are recorded on the context.""" + + 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.""" + + 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/src/google/adk/optimization/agent_optimizer.py b/src/google/adk/optimization/agent_optimizer.py index 7a6da422b09..2f72469b177 100644 --- a/src/google/adk/optimization/agent_optimizer.py +++ b/src/google/adk/optimization/agent_optimizer.py @@ -17,8 +17,11 @@ from abc import ABC from abc import abstractmethod from typing import Generic +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 @@ -28,11 +31,23 @@ 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/gepa_root_agent_optimizer.py b/src/google/adk/optimization/gepa_root_agent_optimizer.py index 01d93aae029..bb615bfddd2 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 @@ -32,6 +33,7 @@ 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 @@ -333,6 +335,8 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context: Optional[OptimizationRunContext] = None, ) -> GEPARootAgentOptimizerResult: """Runs the GEPARootAgentOptimizer. @@ -346,6 +350,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..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,6 +31,7 @@ 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 @@ -207,6 +208,8 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context: Optional[OptimizationRunContext] = None, ) -> GEPARootAgentPromptOptimizerResult: """Runs the GEPARootAgentPromptOptimizer. @@ -220,6 +223,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/simple_prompt_optimizer.py b/src/google/adk/optimization/simple_prompt_optimizer.py index 6a1c7398179..e98c844c3ff 100644 --- a/src/google/adk/optimization/simple_prompt_optimizer.py +++ b/src/google/adk/optimization/simple_prompt_optimizer.py @@ -18,11 +18,13 @@ 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 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 @@ -203,7 +205,18 @@ async def optimize( self, initial_agent: Agent, sampler: Sampler[UnstructuredSamplingResult], + *, + run_context: Optional[OptimizationRunContext] = 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 new file mode 100644 index 00000000000..5baa32feb1d --- /dev/null +++ b/tests/unittests/optimization/run_context_test.py @@ -0,0 +1,802 @@ +# 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 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 + + +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_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(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(STAGE_CANDIDATE_GENERATION) + ctx.end_model_call(h, usage_metadata=None) + with pytest.raises(OptimizationBudgetExceeded) as exc: + 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 + assert len(snap.events) == 2 + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + + +class TestTokenBudget: + + def test_overshoot_commits_then_raises(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=100) + ) + 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 + # 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.run_status == RunStatus.BUDGET_EXCEEDED + assert snap.token_budget_status == TokenBudgetStatus.EXCEEDED + + def test_unreported_usage_is_indeterminate_not_compliant(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + 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 + # 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(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=40)) + assert ctx.snapshot().token_budget_status == TokenBudgetStatus.WITHIN_LIMIT + + +class TestUsageClassification: + + def test_verified_partial_unreported(self): + ctx = OptimizationRunContext() + 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(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(prompt_token_count=5)) + 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 + 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(STAGE_CANDIDATE_GENERATION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=3)) + assert ctx.snapshot().usage_coverage == UsageCoverage.VERIFIED + + +class TestProviderFailure: + + def test_provider_error_is_a_governed_terminal(self): + ctx = OptimizationRunContext() + 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 # 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: + + 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(STAGE_REFLECTION) + assert "deadline" in str(exc.value) + assert exc.value.snapshot.cancel_reason == "deadline" + assert exc.value.snapshot.run_status == RunStatus.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(STAGE_CANDIDATE_GENERATION) + # A governance caller can persist the attempt from a finally block. + 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.logical_call_limits_enforceable + assert not caps.reported_token_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 TestResultShapeUnchanged: + + 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 + + assert "terminal_status" not in OptimizerResult.model_fields + + +class TestBuiltInsRejectContextInContractsPr: + + @pytest.mark.asyncio + async def test_simple_prompt_optimizer_rejects_context(self): + from unittest import mock + + from google.adk.optimization import UnsupportedOptimizationContextError + from google.adk.optimization.sampler import Sampler + from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizer + from google.adk.optimization.simple_prompt_optimizer import SimplePromptOptimizerConfig + + with mock.patch( + "google.adk.optimization.simple_prompt_optimizer.LLMRegistry.resolve" + ): + optimizer = SimplePromptOptimizer(SimplePromptOptimizerConfig()) + sampler = mock.MagicMock(spec=Sampler) + with pytest.raises(UnsupportedOptimizationContextError): + await optimizer.optimize( + mock.MagicMock(), sampler, run_context=OptimizationRunContext() + ) + # Rejected before any sampler work. + sampler.get_train_example_ids.assert_not_called() + + +class TestAtomicSettlement: + """Round-2 adversarial findings: validate-before-commit.""" + + def test_nan_usage_does_not_corrupt_the_record(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=_usage(total_token_count=float("nan"))) + event = ctx.snapshot().events[0] + # NaN/inf/negative counters are "not reported", never a crash or a + # half-closed event. + assert event.state == ModelCallState.COMPLETED + assert event.total_tokens is None + assert event.usage_coverage == UsageCoverage.UNREPORTED + + def test_negative_and_inf_usage_treated_as_unreported(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, + usage_metadata=_usage( + total_token_count=float("inf"), prompt_token_count=-3 + ), + ) + event = ctx.snapshot().events[0] + assert event.total_tokens is None + assert event.prompt_tokens is None + + def test_numeric_error_code_is_normalized_to_string(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call(h, usage_metadata=None, error_code=429) + snap = exc.value.snapshot + assert snap.events[0].error_code == "429" + assert snap.terminal_error_code == "429" + + def test_error_metadata_is_sanitized_and_bounded(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError) as exc: + ctx.end_model_call( + h, + error_code="multi\nline /etc/passwd " + "x" * 500, + error_type="/usr/lib/module.py", + ) + event = exc.value.snapshot.events[0] + assert "\n" not in event.error_code + assert "/" not in event.error_code + assert len(event.error_code) <= 128 + assert "/" not in event.error_type + + +class TestTerminalCompleteness: + """Round-2 adversarial findings: every caller observes the terminal.""" + + def test_finalize_success_observes_pending_cancellation(self): + ctx = OptimizationRunContext() + ctx.request_cancel("deadline") + with pytest.raises(OptimizationCancelledError): + ctx.finalize_success() + assert ctx.snapshot().run_status == RunStatus.CANCELLED + + def test_inflight_settlement_after_terminal_raises_existing_terminal(self): + ctx = OptimizationRunContext( + OptimizationBudgets(max_provider_reported_tokens=10) + ) + h1 = ctx.begin_model_call(STAGE_REFLECTION) + h2 = ctx.begin_model_call(STAGE_REFLECTION) # admitted before terminal + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h1, usage_metadata=_usage(total_token_count=50)) + # The second in-flight call settles for truthful accounting, but its + # caller receives the existing terminal and cannot continue. + with pytest.raises(OptimizationBudgetExceeded): + ctx.end_model_call(h2, usage_metadata=_usage(total_token_count=5)) + snap = ctx.snapshot() + assert snap.completed_calls == 2 # both settled + assert snap.cumulative_total_tokens == 55 # evidence preserved + assert snap.run_status == RunStatus.BUDGET_EXCEEDED + + def test_cancelled_settlement_preserves_usage_without_raising(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call( + h, usage_metadata=_usage(total_token_count=7), cancelled=True + ) + snap = ctx.snapshot() + assert snap.events[0].state == ModelCallState.CANCELLED + assert snap.events[0].total_tokens == 7 + assert snap.completed_calls == 1 # settled calls, any call status + + +class TestFailureTypes: + + def test_finalize_failed_is_generic_not_provider(self): + from google.adk.optimization import OptimizationFailedError + + ctx = OptimizationRunContext() + ctx.finalize_failed(error_code="SAMPLER_CRASH", error_type="ValueError") + with pytest.raises(OptimizationFailedError): + ctx.begin_model_call(STAGE_REFLECTION) + + def test_provider_call_failure_stays_provider_typed(self): + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + with pytest.raises(OptimizationProviderError): + ctx.end_model_call(h, error_code="INTERNAL") + with pytest.raises(OptimizationProviderError): + ctx.begin_model_call(STAGE_REFLECTION) + + +class TestDurableTimestamps: + + def test_event_times_are_wall_clock(self): + import time as _time + + before = _time.time() + ctx = OptimizationRunContext() + h = ctx.begin_model_call(STAGE_REFLECTION) + ctx.end_model_call(h, usage_metadata=None) + event = ctx.snapshot().events[0] + after = _time.time() + assert before <= event.start_time <= event.end_time <= after + + +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) + barrier = threading.Barrier(2) + outcomes: dict[str, object] = {} + + 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, sync=barrier, sink=outcomes): + sync.wait() + try: + context.finalize_success() + sink["success"] = "completed" + except OptimizationFailedError: + sink["success"] = "failed_open" + + t1 = threading.Thread(target=settle) + t2 = threading.Thread(target=finish) + t1.start() + t2.start() + t1.join() + t2.join() + snap = ctx.snapshot() + # 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"