feat(optimization): run-scoped usage, budget, and cancellation contracts (RFC #6357, PR 1/2)

[caohy1988]

RFC: #6357 — please review the RFC before this code. PR 1 of the two-PR structure: contracts only, zero optimizer behavior change. PR 2 (#6359, stacked) instruments SimplePromptOptimizer and GEPARootAgentPromptOptimizer.

Hardened across two adversarial review rounds (terminal state machine, frozen ledger, atomic settlement, truthful failure types).

What this adds

• optimization/run_context.py (new, experimental):
  • OptimizationRunContext — optional, one-shot, caller-owned, concurrency-safe ledger of logical optimizer-owned model invocations (control-plane metadata only; no prompt/response content).
  • Explicit terminal state machine, first terminal wins: every governed run finalizes exactly once as completed | budget_exceeded | cancelled | failed via ledger operations or finalize_success / finalize_cancelled / finalize_failed. begin_model_call rejects terminal contexts; a late cancellation cannot overwrite an earlier terminal; finalize_success observes pending cancellation and can never record a cancelled run as success.
  • Atomic settlement, validate-before-commit: usage counters must be finite non-negative numbers (NaN/inf/negative → not-reported, never a half-closed record); error metadata is normalized to bounded single-token identifiers (numeric provider codes become strings; multiline/path text scrubbed); event, counters, and terminal transition commit under one lock acquisition. An already-admitted call settles for truthful accounting and then receives the existing terminal error, so no caller continues past a committed terminal.
  • Call status ≠ run status: a call is completed | provider_error | cancelled; the run is completed | budget_exceeded | cancelled | failed. The triggering call of a token overshoot stays completed while the run becomes budget_exceeded. end_model_call(cancelled=True) settles a call with its usage evidence preserved.
  • Truthful usage coverage: verified / partial / unreported per call and per run; token_budget_status: within_limit | exceeded | indeterminate — missing totals are never read as compliance. Only logical-call limits are hard (atomic slot admission); max_provider_reported_tokens terminates reactively, commit-then-terminate.
  • Truthful failure types: OptimizationProviderError is reserved for governed provider-call terminals (provider-error-first precedence over simultaneous token overshoot); OptimizationFailedError covers sampler/adapter/optimizer failures via finalize_failed.
  • on_budget_exceeded: "raise" | "return_partial" (default raise). The public OptimizerResult schema is unchanged — the caller-owned snapshot is authoritative; snapshots/events are frozen models (tuple events) carrying budgets, terminal sequence, and terminal error metadata; event timestamps are wall-clock for durable attempt records.
  • Ledger integrity: context-bound handles (foreign commit is a typed error), lock-atomic close (concurrent double-close commits once), private unattached sentinel (attach(None) is one-shot), frozen ge=0-validated budgets.
• OptimizerCapabilities — granular, conservative defaults: accepts_run_context, model_calls_observable, logical_call_limits_enforceable, reported_token_limits_enforceable, cooperative_cancellation, sampler_usage_included (always False). Callers omit the keyword when preflight says unsupported.
• AgentOptimizer.optimize(...) gains keyword-only run_context: Optional[OptimizationRunContext] = None; the three built-ins accept the typed keyword and reject a supplied context with a typed error in this PR, so the abstract-signature change is independently type-clean; real support lands in the stacked PR.
• Stage is an extensible string with documented constants.

Verification

tests/unittests/optimization/ isolated on this head: 70/70 — including both adversarial rounds' probes as first-class tests: foreign handles, 8-thread concurrent double-close, post-terminal admission (budget and completed), frozen budgets/snapshot mutation attempts, attach-None reuse, terminal precedence (late cancel vs budget), NaN/inf/negative usage, numeric error codes, sanitized/bounded error metadata, in-flight settlement after terminal, cancelled settlement with usage evidence, generic-vs-provider failure typing, wall-clock timestamps. ruff --select F401,F811,F841 clean; isort + pyink applied. Stacked with #6359: 96/96.

Downstream consumer: GoogleCloudPlatform/BigQuery-Agent-Analytics-SDK#318.

[caohy1988]

Verification update: after installing the missing local test dependencies ([eval] extra + pytest-mock), the entire tests/unittests/optimization/ suite passes — 53/53 on the stacked enforcement branch (#6359), which includes this PR's commits. The "25 pre-existing collection errors" mentioned in the description were all environmental (missing pandas/pytest-mock in my venv), not code issues — confirmed identical on upstream/main before any changes.

[caohy1988]

Code Review Results

Scope: google/adk-python PR #6358, 9d306f5..e726c58 (4 files, +646/-0)
Intent: Add the contracts half of RFC #6357 without changing existing optimizer behavior or serialization.
Mode: report-only external PR review
Review coverage: correctness, testing, maintainability, project standards, API contract, reliability, concurrency, security/privacy, and adversarial state-machine probes. No subagents or peer-model passes were used, per the user's instruction.

Triage Groups

Group	Findings	Context	Preferred resolution	Why
Sync PR 1 to the current RFC	#1, #2, #3, #5, #7	The PR commit predates the issue's precision pass and still implements the superseded result/status/capability shape	Rebase the public data model on the current RFC first, then adjust adapters/tests	These findings share one cause; fixing them piecemeal would preserve contradictions
Make the ledger trustworthy under concurrency	#4, #6, #8	Handle completion, evidence objects, configuration, and attachment all participate in the context's concurrency/one-shot guarantees	Define the lifecycle invariants, enforce them under one lock, then add deterministic race tests	A governance ledger is useful only if its counts and snapshots cannot be corrupted

P1 -- High

#	File	Issue	Confidence
1	src/google/adk/optimization/data_types.py:78	terminal_status changes every existing OptimizerResult schema and contradicts the current RFC	10/10
2	src/google/adk/optimization/run_context.py:100	The snapshot cannot represent final run state or truthful reported-token compliance	10/10
3	src/google/adk/optimization/agent_optimizer.py:45	The new base signature is incompatible with all three built-in overrides and capabilities cannot express keyword acceptance/token enforcement	10/10
4	src/google/adk/optimization/run_context.py:300	One call handle can be committed twice, and a foreign-context handle corrupts both ledgers	10/10
5	src/google/adk/optimization/run_context.py:308	Provider errors are not a typed terminal outcome and can be masked by a token-budget exception	10/10

• How to use api_transport "REST"? in Agent #1: On main, OptimizerResult(optimized_agents=[]).model_dump() is {'optimized_agents': []}; this PR changes it to {'optimized_agents': [], 'terminal_status': None}. That violates the RFC's explicit result-schema preservation and the PR's no-context equivalence claim. Remove the field, put run_status on the caller-owned snapshot, and add a before/after serialization regression test for base and GEPA result subclasses.
• Google cloud requirement #2: terminal_control_state=None means both running and normally completed, while the snapshot omits configured limits, typed run_status, within_limit | exceeded | indeterminate reported-token status, triggering invocation ID, and terminal error metadata. max_total_tokens also restores the misleading name the RFC changed to max_provider_reported_tokens. Add typed run/budget statuses and explicit finalization methods so finally: persist(context.snapshot()) yields an unambiguous terminal record.
• we need typescript version #3: Strict mypy reports three new [override] failures at simple_prompt_optimizer.py:202, gepa_root_agent_prompt_optimizer.py:206, and gepa_root_agent_optimizer.py:332; all still expose the old two-argument signature. The capability model also omits accepts_run_context and a distinct reported-token-enforcement flag, so a wrapper cannot safely decide whether to omit the keyword. Update all built-in signatures in PR 1, reject non-null contexts with UnsupportedOptimizationContextError until the relevant PR 2 instrumentation is present, and implement the RFC's six granular flags.
• fix: Update README evaluate link #4: _closed is checked and written before acquiring the context lock. A deterministic two-thread interleaving produced one event with completed_calls=2 and cumulative_total_tokens=22 from one 11-token handle. Passing a handle from context A to context B also leaves A with a completed event but zero completed calls and B with one completed call but zero events. Bind handles to their creating context and move ownership validation plus the open-to-closed transition under the same lock as ledger mutation.
• Evaluation link broken #5: end_model_call(..., error_message=...) records provider_error, but if the same response crosses the token threshold, line 324 raises OptimizationBudgetExceeded before the caller can preserve the provider failure as primary. The RFC instead requires sanitized error code/exception type, run_status='failed', and OptimizationProviderError(snapshot, error_code). Add that typed exception and define precedence so usage and budget status are recorded without converting a provider failure into a budget failure. Do not persist arbitrary provider exception messages in the control-plane ledger.

P2 -- Moderate

#	File	Issue	Confidence
6	src/google/adk/optimization/run_context.py:100	Snapshots and the live budget configuration are mutable despite the immutable/concurrency-safe contract	10/10
7	src/google/adk/optimization/run_context.py:60	The call schema closes the RFC's extensible stage field and contains an invalid run-level call state	9/10
8	src/google/adk/optimization/run_context.py:205	attach(None) can succeed repeatedly, bypassing the one-shot guard	10/10

• Elevated Permission Error on Windows #6: OptimizationRunSnapshot and nested ModelCallEvent models are mutable, events is a mutable list, and budgets exposes the context's live mutable object. A caller can change max_model_calls after construction and alter returned evidence in place. Freeze/copy the budget config, validate limits as non-negative, and return deeply frozen point-in-time snapshots (for example frozen Pydantic models plus a tuple of frozen events).
• docs: update ToolboxTool docstring #7: stage: ModelCallStage rejects any third-party/custom stage although the RFC specifies an extensible string with two initial values. ModelCallState.BUDGET_EXCEEDED is a run outcome, not a call outcome, and is never assigned; CANCELLED is declared but no completion API can set it. Use an extensible string/constant surface and keep call status to completed | provider_error | cancelled, with budget termination represented only at run level.
• Using Google's prompt templates with ADK #8: _attached_owner is not None uses None as both a legal runtime argument and the unattached sentinel. Track attachment with a separate boolean/private sentinel and add concurrent-attachment plus None/invalid-owner tests.

P3 -- Low

#	File	Issue	Confidence
9	src/google/adk/optimization/__init__.py:1	The new public contract is not explicitly exported as required by ADK's API principles	8/10

• Malformed function call #9: Export the intended public names through google.adk.optimization.__init__ and __all__, or explicitly document that this experimental module is intentionally module-qualified and reconcile the repository's API-principles rule.

Verification

• PR tests: 15 passed; focused coverage reports 97% for run_context.py.
• Full PR feat(optimization): run-scoped usage, budget, and cancellation contracts (RFC #6357, PR 1/2) #6358 optimization suite with the eval/test extras: 43 passed. The author's 53/53 update is consistent with stacked PR feat(optimization): enforce run-context usage, budgets, and cancellation in the prompt optimizers (RFC #6357, PR 2/2) #6359 adding ten enforcement tests; it is not the PR feat(optimization): run-scoped usage, budget, and cancellation contracts (RFC #6357, PR 1/2) #6358 test count.
• Adversarial probes: 6 passed, each proving a current failure mode (double commit, foreign handle, mutable snapshot/config, attach(None) reuse, provider-error masking).
• Strict mypy comparison: PR has 20 errors versus 17 on main; the three new errors are exactly the built-in optimize override incompatibilities.
• Ruff: failed on the new unused AgentWithScores import in run_context_test.py:189.
• Pyink: failed; run_context.py would be reformatted. Isort and git diff --check: passed.
• Python 3.10 and 3.14 focused runs were inconclusive because this local macOS toolchain could not build cryptography for the downloaded 3.10 x86_64 interpreter and the downloaded 3.14 alpha interpreter segfaulted in pydantic_core; these are environment failures, not attributed to the PR.
• Live GitHub checks at review time: CLA and header/check-changes passed; the substantive test/lint matrix had not run.

Testing gaps

• No before/after serialization-equivalence test.
• No concurrent call admission or concurrent handle completion test.
• No foreign/unknown handle test.
• No snapshot/config immutability test.
• No final completed/failed/indeterminate snapshot test.
• No typed unsupported preflight against built-in and old third-party overrides.
• No extensible stage or cancelled-call finalization test.
• No provider-error-plus-budget-precedence test.

---

Verdict: Not ready.

Reasoning: The direction is right and the happy-path ledger tests are clean, but PR 1 implements the pre-precision RFC shape, introduces three strict type failures, changes existing result serialization, and cannot yet produce the authoritative immutable final record the downstream governance use case needs.

Fix order: Sync the public contract (#1, #2, #3, #5, #7) -> repair handle/lifecycle concurrency (#4, #6, #8) -> add the missing semantic tests -> clear Ruff/Pyink -> rerun the supported-version matrix.

Actionable Findings

1. [P1] data_types.py:78 - remove terminal_status; make snapshot run_status authoritative.
2. [P1] run_context.py:100 - add complete typed final snapshot/lifecycle and reported-token budget status.
3. [P1] agent_optimizer.py:45 - make built-ins signature-compatible and restore granular capability/preflight behavior.
4. [P1] run_context.py:300 - make handle ownership and close/commit atomic under the context lock.
5. [P1] run_context.py:308 - add typed sanitized provider failure and correct error precedence.
6. [P2] run_context.py:100 - deeply freeze snapshots and budget configuration.
7. [P2] run_context.py:60 - make stages extensible and separate call/run statuses.
8. [P2] run_context.py:205 - replace the None attachment sentinel and test concurrent reuse.
9. [P3] optimization/__init__.py:1 - export the intended public experimental surface or document the exception.

[caohy1988]

Fresh review of the contracts PR

Reviewed at f345eee1e2bdb171abd58ca0fde8c5a8359b0956, independently against the current #6357 contract rather than relying on the previous disposition.

Verdict: keep this PR in draft / request changes. The revised direction is right, but the run context is not yet an authoritative, immutable terminal record. Several earlier load-bearing findings remain present.

What is now right

• OptimizerResult remains unchanged; the caller-owned snapshot is correctly intended to be authoritative.
• Reported-token compliance distinguishes within_limit, exceeded, and indeterminate; missing totals are not treated as proof of compliance.
• Call status and run status are separate.
• Capabilities are granular and default conservatively.

[P1] Terminal lifecycle is incomplete

run_context.py:283-410

The context has terminal status values but no complete terminal state machine:

• There is no success finalizer, so successful governed optimizations finish with run_status=None.
• Provider failures do not set run_status=FAILED.
• Native cancellation can leave run_status=None and an open call event.
• begin_model_call() does not reject an already-terminal context, so another call can start after a committed token overshoot.
• raise_if_cancelled() unconditionally writes CANCELLED, allowing a late cancellation to overwrite an earlier budget terminal state.
• The snapshot omits configured limits, the triggering invocation ID, and terminal error metadata required by the RFC.

The fix should be an explicit, lock-protected terminal transition API: completed / budget exceeded / cancelled / failed, exactly once, with first-terminal-wins semantics. begin_model_call() must reject every terminal context. Every optimizer exit path in the enforcement PR must invoke the corresponding transition.

[P1] Provider failure is only a call annotation, not a governed terminal outcome

run_context.py:322-367

end_model_call(error_message=...) marks the event provider_error, but it does not:

• set the run to FAILED;
• return or raise OptimizationProviderError;
• carry a structured/sanitized error_code and exception type;
• define provider-error versus token-overshoot precedence.

Today, usage that crosses the token limit causes OptimizationBudgetExceeded to replace the provider failure. The RFC needs the provider failure to remain primary while preserving both the committed usage and token-compliance evidence. Persisting arbitrary str(exception) also weakens the control-plane-only/privacy boundary.

Please make terminal provider commit atomic and structured. The public typed error should expose the sanitized code and final snapshot; raw provider/request text should not be copied into the ledger.

[P1] The concurrency-safe, one-shot ledger can be corrupted or mutated

run_context.py:214-257 and run_context.py:322-410

• A handle created by context A can be committed into context B, incrementing B's counters with A's event.
• _CallHandle._closed is checked and written outside the context lock, so concurrent close is not atomic.
• attach(None) can be called repeatedly because None is both a valid argument and the unattached sentinel.
• The context retains and returns the caller's mutable OptimizationBudgets object. Limits can therefore change during a run.
• OptimizationRunSnapshot, its events list, and its event models are mutable despite the documented immutable contract.

Bind every handle to a private context identity, validate ownership under the lock, make close atomic, replace the attachment sentinel with an explicit boolean/private sentinel, copy/freeze budgets at construction, and return frozen snapshots with tuples/frozen event models.

[P2] Contract precision and independent landability

• ModelCallStage is still a closed enum, while RFC: run-scoped usage, budget, and cancellation controls for AgentOptimizer (OptimizationRunContext + OptimizerCapabilities) #6357 specifies an extensible stage string. Use str plus documented constants if arbitrary optimizer stages must survive without an ADK release.
• Negative model-call and token limits are accepted. Add ge=0 validation.
• The budget documentation at lines 170-177 still refers to the removed terminal_status result field; the end_model_call documentation still calls the reported-token threshold a hard ceiling.
• PR 1 changes the abstract optimize() signature without updating the built-in implementations in the same independently landable PR. In my local comparison this adds three mypy override errors over the reviewed main baseline; the signatures only converge after stacking feat(optimization): enforce run-context usage, budgets, and cancellation in the prompt optimizers (RFC #6357, PR 2/2) #6359.

Verification

• Isolated tests/unittests/optimization: 44 passed.
• Fresh contract probes for attachment reuse, foreign handles, snapshot/budget mutability, extensible stages, and post-terminal call admission: 5 failed out of 5.

The passing suite validates the positive accounting path, but it does not currently exercise the adversarial invariants above. Please add deterministic tests for every terminal transition, terminal precedence, cross-context handles, concurrent double-close, post-terminal admission, immutable budgets/snapshots, and native cancellation before treating the contract PR as complete.

This is a good API direction, but the ledger is the governance boundary. Its lifecycle and integrity guarantees need to be true in code, not only in the RFC prose.

[caohy1988]

Review disposition — all three P1s and the P2s implemented (head bd5b94d)

[P1] Terminal lifecycle → implemented as an explicit lock-protected, first-terminal-wins state machine. finalize_success() / finalize_cancelled() / finalize_failed() added; begin_model_call() rejects every terminal context with the corresponding typed error; raise_if_cancelled() re-raises an earlier terminal outcome instead of overwriting it (test: late cancel after budget overshoot stays BUDGET_EXCEEDED); finalize_cancelled() closes open call events; the snapshot now carries the configured budgets, the terminal-triggering sequence, and terminal error code/type.

[P1] Provider failure as governed terminal → end_model_call(error_code=, error_type=) commits usage, transitions the run to FAILED, and raises OptimizationProviderError atomically under the lock. Provider-error-first precedence over a simultaneous token overshoot is implemented and tested (usage and token_budget_status=EXCEEDED evidence both preserved while the run terminates failed). Error metadata is sanitized identifiers only — raw exception text never enters the ledger.

[P1] Ledger integrity → handles are bound to their context (foreign commit raises, victim ledger untouched — tested); close is atomic under the context lock (8-thread concurrent double-close commits exactly once — tested); the unattached sentinel is a private object so attach(None) is one-shot (tested); OptimizationBudgets is a frozen model with ge=0 validation; snapshots and events are frozen models with tuple events (mutation attempts raise — tested).

[P2s] → stage is an extensible string with documented constants; stale terminal_status/hard-ceiling doc text removed in the rewrite; independent type-cleanliness: this PR now also updates the three built-in optimizers to accept the keyword and reject a supplied context with a typed error, so the abstract-signature change introduces no override mismatch before the stacked PR lands.

Verification on this head: the isolated optimization suite is 60/60, including the review's adversarial probes as first-class tests (foreign handles, concurrent double-close, post-terminal admission for both budget and completed states, frozen budgets/snapshots, attach-None reuse, terminal precedence). ruff --select F401,F811,F841 clean; isort + pyink applied.

[caohy1988]

Fresh adversarial review of the hardened contracts head

Reviewed at bd5b94d3f123602bc524bd66c48b376b72d39cf5, including the complete diff, the prior disposition, official tests, type/lint gates, concurrent-terminal behavior, malformed boundary values, and current-main compatibility.

Verdict: keep draft / request changes. The rewrite genuinely fixes the previous attachment, handle ownership, immutability, extensible-stage, terminal-admission, and provider-precedence findings. The governance boundary still has two load-bearing correctness gaps plus several smaller contract/gate issues.

Verified improvements

• First terminal status is retained against a later observed cancellation.
• A terminal context rejects new call admission.
• Foreign-context handles are rejected without changing the victim ledger.
• Handle close is atomic; the eight-thread double-close test commits once.
• attach(None) is now one-shot.
• Budgets, snapshots, events, and the event collection are immutable.
• Negative budgets are rejected and stages are extensible strings.
• Provider error wins over simultaneous token overshoot while usage and token-compliance evidence remain visible.
• Built-in signatures are updated in this PR, eliminating the earlier override mismatch.
• The isolated optimization suite is 60/60 and merge-tree analysis shows no conflict with current main.

[P1] The metadata boundary is not normalized or exception-safe

run_context.py:490-563 documents sanitized identifiers, but persists whatever string the caller supplies. There is no length, character, or structure enforcement. A raw multiline/provider payload can enter both the event and terminal snapshot unchanged.

The commit order also makes invalid usage non-atomic: the record is marked closed and receives end_time before _extract_usage() finishes. A malformed value such as NaN, or an accessor that raises, leaves a half-closed event with end_time != None, state=None, and no completed-call increment.

This becomes a live failure in #6359: a raised provider exception with numeric error_code=429 reaches this boundary, snapshot construction raises Pydantic ValidationError, the caller never receives OptimizationProviderError, and subsequent snapshots remain broken.

Fix: normalize/cap error identifiers at this boundary, validate/extract usage into local values before mutating the record, and perform the complete commit only after conversion succeeds. Invalid provider metadata should terminate with a defined sanitized failure or explicitly degrade coverage to indeterminate, never partially mutate the ledger.

[P1] First-terminal-wins is incomplete for pending cancellation and already-in-flight calls

finalize_success() blindly transitions to COMPLETED even when cancel_requested=True. I reproduced a pending cancellation being finalized as success.

There is a second concurrency hole in end_model_call(): admit calls A and B, let A cross the token limit, then settle already-in-flight B. B commits and returns normally because _transition_locked() loses to the existing terminal and no existing terminal error is propagated. B's caller can schedule sampler work after the run already terminated.

Fix: success finalization must atomically observe pending cancellation and require no open calls. An already-in-flight call may settle and commit after another call wins the terminal race, but its caller must receive the first terminal outcome before returning to downstream work.

[P2] Failed-run typing conflates provider and non-provider failures

_terminal_error_locked() maps every RunStatus.FAILED to OptimizationProviderError. A context finalized for ADAPTER_CRASH or a sampler/setup failure is therefore recast as a provider failure when observed again.

Either add a generic typed run-failure/finalized outcome or retain terminal error kind/source so only provider failures become OptimizationProviderError.

[P2] Cancelled-call counters are internally inconsistent

finalize_cancelled() closes open events as CANCELLED but leaves completed_calls unchanged. The resulting snapshot can contain one terminal event with started_calls=1, completed_calls=0.

If this counter means terminal/settled calls, increment it when cancellation closes an event. If it means successful/provider-settled calls only, rename and document it so downstream reconciliation does not infer missing calls.

[P2] Persisted start_time / end_time are process-relative

The event uses time.monotonic(). That is correct for duration measurement but is not a timestamp that can be compared across processes or interpreted after persistence.

Rename these fields to start_monotonic / end_monotonic, or expose UTC timestamps plus a monotonic duration. The current names invite downstream consumers to store process-uptime values as event timestamps.

Merge gates and claim drift

• Ruff fails on the unused typing.Optional import in simple_prompt_optimizer.py.
• Strict Mypy reports five new errors in run_context.py: one untyped parameter, three attributes inferred as None, and one optional .value access. The same changed-file check on current main is clean.
• Pyink and Isort are clean.
• GitHub currently runs only lightweight header, CLA, and workflow-scan checks on this fork PR, so the missing substantive gates are not visible in the check list.
• The PR body still describes 16 contract tests / 57 stack tests. The reviewed head has 60 isolated optimization tests and materially different lifecycle semantics.

Tests still needed on this PR

• pending cancellation racing success finalization;
• two already-admitted calls settling after the first terminal wins;
• numeric, oversized, multiline, and otherwise malformed error identifiers;
• malformed/non-finite usage without partial ledger mutation;
• cancelled-event counter semantics;
• non-provider FAILED observation;
• concurrent call-slot admission and concurrent attachment, which the RFC names but the suite currently tests only sequentially.

The rewrite moved the contract a long way forward. The next pass should make the public boundary defensive against the values adapters actually produce and make terminal propagation correct for all admitted concurrency, not only admission itself.

[caohy1988]

Fresh review round — contracts PR (bd5b94d)

Verdict: keep draft / changes required. I re-reviewed the current remote head against current ADK main (57e1ba6), reran the official optimization suite, and exercised the contract with independent concurrency, cancellation, metadata, and ledger-integrity probes. The advertised core improvements are real, but the state machine is not yet a reliable governance boundary.

P1 — ledger settlement is not atomic

end_model_call() marks a record closed and writes end_time before usage and error metadata have been validated or converted:

adk-python/src/google/adk/optimization/run_context.py

Lines 513 to 563 in bd5b94d

	record = handle._record
	with self._lock:
	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
	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

Two concrete reproductions:

1. total_token_count=float("nan") raises during int(value) after closed=True and end_time are committed. The resulting event has end_time != None but state=None, and the handle cannot be retried because it is already closed.
2. A provider exception with numeric error_code=429 stores an integer where the snapshot expects Optional[str]. Snapshot construction then raises a Pydantic validation error instead of OptimizationProviderError, leaving the governed record corrupted.

The public boundary also accepts arbitrary multiline error_code strings and path-like error_type values despite documenting these fields as sanitized identifiers.

Required fix: normalize and validate every usage/model/error field before mutating the record. Provider counters should be finite, non-negative integers; error codes/types should be bounded normalized strings. Then commit the full event, counters, and terminal transition atomically under the lock. No exception path should expose a half-closed event.

P1 — “first terminal wins” is incomplete

finalize_success() commits COMPLETED without observing a pending cancellation:

adk-python/src/google/adk/optimization/run_context.py

Lines 378 to 382 in bd5b94d

	def finalize_success(self) -> None:
	"""Marks a normal optimizer return ``completed`` (first terminal wins)."""
	with self._lock:
	self._transition_locked(RunStatus.COMPLETED)

Reproduction: request_cancel("deadline") followed by finalize_success() produces a completed run instead of a typed cancellation.

There is also an already-inflight concurrency gap. If two calls have been admitted and the first crosses the token ceiling, the second can settle afterward and return normally. _transition_locked() declines to overwrite the existing terminal, but end_model_call() does not propagate that existing terminal outcome. Its caller can therefore schedule downstream work despite the committed budget terminal.

Required fix:

• finalize_success() must observe and commit/raise pending cancellation before success.
• An already-admitted call may settle for truthful accounting, but after settlement its caller must receive the existing terminal error when another call has already terminated the run.

P1 — terminal error typing collapses unrelated failures

Every RunStatus.FAILED is reconstructed as OptimizationProviderError:

adk-python/src/google/adk/optimization/run_context.py

Lines 433 to 450 in bd5b94d

	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}."
	)

That includes sampler, adapter, validation, and optimizer failures recorded by finalize_failed(). The typed API therefore misclassifies non-provider failures. Please add/retain a generic optimization failure outcome and reserve OptimizationProviderError for governed provider-call failures.

Contract clarifications

• finalize_cancelled() closes open events as CANCELLED, but does not increment completed_calls. Either define that field as successful completions only or rename/count it as settled/closed calls. The current event/counter pair is ambiguous.
• Audit-facing start_time/end_time are time.monotonic() values. These are process-relative and cannot be persisted or compared across workers. Keep monotonic duration internally, but expose wall-clock timestamps for durable attempt records.

Verification

• Official tests/unittests/optimization: 60 passed.
• Independent cross-stack probes: 11 failed / 4 passed; six failures directly exercise this contract.
• git diff --check: pass.
• Mypy on the affected optimization surface: current main has 24 existing errors; this PR has 32, so 8 new errors are introduced (five in run_context.py and three newly untyped optimizer signatures).
• Ruff: this PR adds one new failure over current main (typing.Optional unused in simple_prompt_optimizer.py).
• The PR body still says 16 tests / 57 passing; the current suite count is 60.

The earlier fixes—context-bound handles, lock-protected double-close behavior, frozen snapshots/budgets, provider-error precedence, and explicit terminal finalizers—are valuable. The remaining issues are narrower, but they are load-bearing: BQAA cannot treat the snapshot as authoritative until settlement is atomic and every caller observes the committed terminal.

[caohy1988]

Round-2 disposition — all three P1s and both clarifications implemented (head 9ed15ac)

[P1] Atomic settlement → all usage/error metadata is now normalized and validated before the record is mutated: counters must be finite non-negative numbers (NaN/inf/negative → not-reported — reproduced in tests, no more half-closed events), error codes/types pass _sanitize_error_meta (numeric 429 → "429"; multiline/path text scrubbed; bounded to 128 chars). Event, counters, and terminal transition commit under one lock acquisition; no exception path can expose a half-closed event.

[P1] First-terminal completeness → finalize_success() observes pending cancellation and commits/raises CANCELLED (a cancelled run can never be recorded as success — tested). The in-flight gap is closed: an already-admitted call settles for truthful accounting (usage evidence preserved, counted), and its caller then receives the existing terminal error — the two-in-flight reproduction from the review is a passing regression test (both calls settled, 55 tokens accounted, run stays BUDGET_EXCEEDED, second caller raises).

[P1] Failure-type collapse → OptimizationFailedError added for non-provider failures (terminal provenance tracked); OptimizationProviderError is reserved for governed provider-call terminals. Both directions tested.

[Clarifications] → completed_calls is documented and counted as settled calls regardless of call status (finalize_cancelled counts its closures; cancelled settlements count); event timestamps are now wall-clock epoch seconds for durable, cross-worker attempt records (tested with before/after bounds).

Gates: the three untyped built-in optimize() signatures are now fully typed (Optional[OptimizationRunContext]), addressing that named mypy delta; the unused-Optional ruff finding no longer applies (the import is used by the typed signature); ruff --select F401,F811,F841 clean; isort/pyink applied. I could not reproduce your exact mypy count locally (config-dependent) — if any of the five run_context.py errors persist under the repo config, happy to chase them with your invocation. PR body refreshed with current counts.

Verification: isolated optimization suite 70/70, with all round-2 reproductions as first-class tests.

[caohy1988]

Round-3 fresh review — contracts head 9ed15ac

Verdict: keep draft / changes required. Round 2 genuinely closes the previously reported atomic-settlement, terminal-provenance, counter, and timestamp findings. The new implementation passes 70/70 repository tests. Fresh contract probes found two remaining correctness gaps plus repository type/public-API gates.

P1 — finalize_success() permits an impossible final snapshot

adk-python/src/google/adk/optimization/run_context.py

Lines 396 to 415 in 9ed15ac

	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.
	"""
	with self._lock:
	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

Two reproduced cases:

1. A call is admitted and left open, then finalize_success() records run_status=COMPLETED with started_calls=1, completed_calls=0, and an event with no terminal call state.
2. A raise-mode budget terminal is committed, then finalize_success() silently returns because _run_status is not None. The required OptimizationBudgetExceeded is no longer observable by that caller.

This is the contract half of a real #6359 false-success path: GEPA can swallow a reflection setup exception after admission and later invoke this finalizer.

Suggested fix:

• Return idempotently only when the existing status is COMPLETED.
• For every other existing terminal, raise _terminal_error_locked().
• Before committing success, reject any record where closed is False.
• Prefer committing a generic failed terminal such as OPEN_MODEL_CALLS and raising OptimizationFailedError(snapshot) so the run cannot escape nonterminal. Already-admitted calls may still settle later and will observe the existing failure.
• Add direct regressions for success-with-open-call and success-after-budget/provider/cancel terminals.

P1 — returned provider metadata can still corrupt every snapshot

adk-python/src/google/adk/optimization/run_context.py

Lines 565 to 585 in 9ed15ac

	# 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:
	return
	was_terminal = self._run_status is not None
	record.closed = True
	record.end_time = time.time()
	record.returned_model_version = returned_model_version
	record.usage = usage

Usage and error metadata are normalized before mutation, but returned_model_version is stored unchanged. end_model_call(..., returned_model_version=123) succeeds; every later snapshot() raises Pydantic ValidationError because ModelCallEvent.returned_model_version requires a string.

Official LlmResponse validates this field, but custom ADK LLM implementations are a supported extension boundary. The run context is supposed to be the final ledger-normalization boundary.

Suggested fix: normalize or validate all event fields before mutation, including returned model version. For example, use a bounded optional-string helper (None stays None; accepted values become a bounded string) and build/validate the proposed immutable event before copying it into the mutable record. Add a malformed model-version regression that proves snapshots remain readable.

P2 — fractional usage is silently marked verified after truncation

adk-python/src/google/adk/optimization/run_context.py

Lines 719 to 744 in 9ed15ac

	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]:
	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
	return int(value)
	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"),
	}

1.9 passes the finite/non-negative checks, becomes int(1.9) == 1, and is classified as authoritative verified usage. Token counters are integral evidence; silent truncation is not truthful normalization.

Suggested fix: accept integer types excluding bool; accept floats only when value.is_integer(). Treat fractional values as unreported, consistent with NaN/inf/negative behavior. Add zero, integral-float, fractional-float, and very-large-integer tests.

Merge gate — five new Mypy errors are reproducible

Using the repository's own [tool.mypy] configuration, with no custom strict flags:

python -m mypy --no-error-summary --no-pretty src/google/adk/optimization

• current main: 47 errors;
• this PR: 52 errors;
• delta: 5 new errors, all in run_context.py.

Exact new errors are at lines 279, 430, 492, 583, and 584: an untyped constructor argument, attributes inferred as None, and unsafe optional .value access.

Suggested fix: explicitly type every _CallRecord attribute and constructor parameter, and narrow/assert _run_status before reading .value. This should produce zero normalized delta against main under the same command CI uses.

Public API packaging violates ADK's private-by-default rule

src/google/adk/optimization/run_context.py is the only newly added source module. ADK's repository style requires new modules under src/google/adk/ to begin with _, with public names exported through package __init__.py and __all__. optimization/__init__.py is currently empty.

Reference: https://github.com/google/adk-python/blob/57e1ba66ee5f762e22cf5aa727f10fea3485a572/.agents/skills/adk-style/references/visibility.md

Suggested fix: rename to _run_context.py, update internal direct imports, and export only the supported experimental surface from google.adk.optimization with an explicit __all__. This is not one of RFC #6357's four open API questions, so it should be resolved before maintainer review.

Missing contract tests

The RFC still calls for concurrent attachment and concurrent call-slot admission. Current tests cover sequential double attachment and concurrent double close, but not those two races.

Suggested tests:

• N threads call attach() simultaneously; exactly one succeeds.
• N threads race begin_model_call() under max_model_calls=k; exactly k handles are admitted, sequence IDs are unique, and the remainder receive the same committed budget terminal.
• A success finalizer racing the last settlement cannot produce COMPLETED while any call remains open.

Verification

• Official optimization suite: 70 passed.
• Prior round findings: verified fixed.
• Fresh applicable contract probes: open-call success, swallowed existing terminal, malformed model version, and fractional usage all reproduce.
• Pyink, Isort, and git diff --check: pass.
• Ruff: no new delta relative to current main.

The contract is close. The load-bearing change is making finalize_success() a real invariant boundary, not just a status setter; after that, complete provider-field normalization and clear the repository type/module gates.