Skip to content

Commit 8dc8621

Browse files
committed
feat(bigquery-analytics): hot-path same-session pause_orphan resolution
When a long-running tool resumes (a function_response in a user message), resolve pause_orphan on the TOOL_COMPLETED row from the in-memory session history only — a zero-I/O check that adds no latency to the resume path and introduces no session-service or BigQuery read. _resolve_same_session_pause_orphan returns: * False — the originating paused call is present in this session's in-memory history (definitely not an orphan); * None — unknown / not yet settled: no id to pair on, no history, or the pause was not found in the (possibly trimmed) in-memory history. The emit path stamps attributes.adk.pause_orphan only when the resolver returns a real bool; on None the key is omitted so consumers read it as SQL NULL. The resolver never returns True: declaring a true orphan needs the settled fallback (an off-hot-path session-service / BigQuery reconciliation), deliberately out of scope here to keep the resume path free of added reads. A persistent None is what that future fallback upgrades to True. Tests: resolver returns False only on a found pause, None for not-found / empty / missing-id / no-session / id-without-matching-part, never True, and does no session_service read; emit path stamps pause_orphan=False when the pause is in history and omits it otherwise.
1 parent c66dc1d commit 8dc8621

2 files changed

Lines changed: 237 additions & 8 deletions

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 64 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3241,6 +3241,55 @@ async def _log_event(
32413241

32423242
# --- UPDATED CALLBACKS FOR V1 PARITY ---
32433243

3244+
def _resolve_same_session_pause_orphan(
3245+
self,
3246+
invocation_context: "InvocationContext",
3247+
function_call_id: Optional[str],
3248+
) -> Optional[bool]:
3249+
"""Same-session resolution of ``pause_orphan`` for a long-running resume.
3250+
3251+
Scans the in-memory session history for the originating paused
3252+
function call. This is a hot-path-safe, zero-I/O check: it never
3253+
issues a session-service or BigQuery read, so it adds no latency to
3254+
the resume path.
3255+
3256+
Returns:
3257+
* ``False`` — an event in this session's in-memory history carries
3258+
``function_call_id`` in its ``long_running_tool_ids``, so the
3259+
producer emitted a pairable ``TOOL_PAUSED`` for it and this
3260+
completion is definitely not an orphan. Membership alone is
3261+
sufficient: the producer emits a ``TOOL_PAUSED`` for *every* id
3262+
in ``long_running_tool_ids`` — both the per-part path and the
3263+
fallback path that fires when the originating ``function_call``
3264+
part was rewritten away (e.g. by an ``after_model_callback``).
3265+
Requiring a surviving ``function_call`` part here would miss
3266+
that fallback case and diverge from what the producer wrote.
3267+
* ``None`` — unknown / not yet settled: there is no id to pair on,
3268+
no session history is available, or the id was not found in the
3269+
(possibly trimmed) in-memory history. Callers omit the
3270+
``pause_orphan`` attribute in this case so consumers read it as
3271+
SQL NULL ("not yet determined").
3272+
3273+
Never returns ``True``. Declaring a true orphan requires the settled
3274+
fallback (an off-hot-path session-service / BigQuery reconciliation),
3275+
which is intentionally out of scope here so the resume path stays
3276+
free of added reads. A persistent ``None`` is what that future
3277+
fallback upgrades to ``True``.
3278+
"""
3279+
if not function_call_id:
3280+
return None
3281+
session = getattr(invocation_context, "session", None)
3282+
events = getattr(session, "events", None) if session is not None else None
3283+
if not events:
3284+
return None
3285+
for event in events:
3286+
ids = getattr(event, "long_running_tool_ids", None)
3287+
if ids and function_call_id in ids:
3288+
# The producer emitted a pairable TOOL_PAUSED for this id (per
3289+
# part or via the rewritten-away fallback), so it is not orphan.
3290+
return False
3291+
return None
3292+
32443293
@_safe_callback
32453294
async def on_user_message_callback(
32463295
self,
@@ -3300,26 +3349,33 @@ async def on_user_message_callback(
33003349
# tool calls complete inside the agent run via
33013350
# after_tool_callback, so a function_response inside a user
33023351
# message is the resume side of a previously-paused tool.
3303-
# Stamp the pair keys; pause_orphan / registry semantics
3304-
# are intentionally deferred.
33053352
if not part.function_response.id:
33063353
logger.debug(
33073354
"User-message function_response for tool %s has no id;"
33083355
" the resulting TOOL_COMPLETED row cannot pair with a"
33093356
" TOOL_PAUSED row.",
33103357
part.function_response.name,
33113358
)
3359+
adk_extras = {
3360+
"pause_kind": "tool",
3361+
"function_call_id": part.function_response.id,
3362+
}
3363+
# pause_orphan: stamped False only when this session's history
3364+
# proves the completion pairs with a real pause. Omitted (->
3365+
# SQL NULL = "unknown / not yet settled") otherwise. True is
3366+
# reserved for the off-hot-path settled fallback so the resume
3367+
# path stays free of session-service / BigQuery reads.
3368+
pause_orphan = self._resolve_same_session_pause_orphan(
3369+
invocation_context, part.function_response.id
3370+
)
3371+
if pause_orphan is not None:
3372+
adk_extras["pause_orphan"] = pause_orphan
33123373
await self._log_event(
33133374
"TOOL_COMPLETED",
33143375
callback_ctx,
33153376
raw_content=content_dict,
33163377
is_truncated=is_truncated,
3317-
event_data=EventData(
3318-
adk_extras={
3319-
"pause_kind": "tool",
3320-
"function_call_id": part.function_response.id,
3321-
},
3322-
),
3378+
event_data=EventData(adk_extras=adk_extras),
33233379
)
33243380

33253381
@_safe_callback

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import dataclasses
1919
import json
2020
import os
21+
from types import SimpleNamespace
2122
from unittest import mock
2223

2324
from google.adk.agents import base_agent
@@ -8815,3 +8816,175 @@ async def test_matched_id_not_double_emitted_by_fallback(
88158816
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
88168817
pauses = [r for r in rows if r["event_type"] == "TOOL_PAUSED"]
88178818
assert len(pauses) == 1
8819+
8820+
8821+
class TestPauseOrphanSameSession:
8822+
"""#206 PR 1 — hot-path same-session pause resolution.
8823+
8824+
``_resolve_same_session_pause_orphan`` returns False only when the
8825+
originating pause is present in the in-memory session history, None
8826+
otherwise (unknown / not yet settled), and never True. The emit path
8827+
stamps ``pause_orphan`` only when the resolver returns a real bool.
8828+
"""
8829+
8830+
def _plugin(self):
8831+
return bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin(
8832+
project_id=PROJECT_ID, dataset_id=DATASET_ID, table_id=TABLE_ID
8833+
)
8834+
8835+
def _pause_event(self, fc_id):
8836+
fc = types.FunctionCall(id=fc_id, name="long_search", args={})
8837+
return event_lib.Event(
8838+
author="agent",
8839+
content=types.Content(
8840+
role="model", parts=[types.Part(function_call=fc)]
8841+
),
8842+
long_running_tool_ids={fc_id},
8843+
actions=event_actions_lib.EventActions(),
8844+
)
8845+
8846+
def _ic(self, events):
8847+
return SimpleNamespace(session=SimpleNamespace(events=events))
8848+
8849+
# -- resolver: the only path that returns False --
8850+
def test_found_in_session_history_is_not_orphan(self):
8851+
plugin = self._plugin()
8852+
ic = self._ic([self._pause_event("call-1")])
8853+
assert plugin._resolve_same_session_pause_orphan(ic, "call-1") is False
8854+
8855+
# -- everything else is None ("unknown"), never True --
8856+
def test_not_found_is_unknown(self):
8857+
plugin = self._plugin()
8858+
ic = self._ic([self._pause_event("other-call")])
8859+
assert plugin._resolve_same_session_pause_orphan(ic, "call-1") is None
8860+
8861+
def test_empty_history_is_unknown(self):
8862+
plugin = self._plugin()
8863+
assert (
8864+
plugin._resolve_same_session_pause_orphan(self._ic([]), "call-1")
8865+
is None
8866+
)
8867+
8868+
def test_missing_function_call_id_is_unknown(self):
8869+
plugin = self._plugin()
8870+
ic = self._ic([self._pause_event("call-1")])
8871+
assert plugin._resolve_same_session_pause_orphan(ic, None) is None
8872+
assert plugin._resolve_same_session_pause_orphan(ic, "") is None
8873+
8874+
def test_id_in_long_running_without_part_is_not_orphan(self):
8875+
# long_running_tool_ids carries the id but the originating
8876+
# function_call part was rewritten away (e.g. after_model_callback).
8877+
# The producer still emits a pairable TOOL_PAUSED via its fallback,
8878+
# so id-membership alone proves the pause: not orphan.
8879+
plugin = self._plugin()
8880+
ev = event_lib.Event(
8881+
author="agent",
8882+
content=types.Content(role="model", parts=[types.Part(text="hi")]),
8883+
long_running_tool_ids={"call-1"},
8884+
actions=event_actions_lib.EventActions(),
8885+
)
8886+
assert (
8887+
plugin._resolve_same_session_pause_orphan(self._ic([ev]), "call-1")
8888+
is False
8889+
)
8890+
8891+
def test_no_session_is_unknown(self):
8892+
plugin = self._plugin()
8893+
ic = SimpleNamespace(session=None)
8894+
assert plugin._resolve_same_session_pause_orphan(ic, "call-1") is None
8895+
8896+
def test_never_returns_true(self):
8897+
# The resolver must never declare a true orphan in PR 1 — that is
8898+
# reserved for the off-hot-path settled fallback.
8899+
plugin = self._plugin()
8900+
for ic in (
8901+
self._ic([]),
8902+
self._ic([self._pause_event("x")]),
8903+
SimpleNamespace(session=None),
8904+
):
8905+
assert plugin._resolve_same_session_pause_orphan(ic, "call-1") is not True
8906+
8907+
def test_resolver_does_no_session_service_read(self):
8908+
# Hot-path guarantee: the resolver must not call session_service /
8909+
# get_session. A session_service whose every attr raises proves the
8910+
# resolver only touches the in-memory session.events.
8911+
plugin = self._plugin()
8912+
8913+
class _Exploding:
8914+
8915+
def __getattr__(self, name):
8916+
raise AssertionError(f"resolver touched session_service.{name}")
8917+
8918+
ic = SimpleNamespace(
8919+
session=SimpleNamespace(events=[self._pause_event("call-1")]),
8920+
session_service=_Exploding(),
8921+
)
8922+
assert plugin._resolve_same_session_pause_orphan(ic, "call-1") is False
8923+
8924+
# -- emit path: pause_orphan lands on the TOOL_COMPLETED row --
8925+
@pytest.mark.asyncio
8926+
async def test_emit_stamps_false_when_pause_in_history(
8927+
self,
8928+
bq_plugin_inst,
8929+
mock_write_client,
8930+
invocation_context,
8931+
dummy_arrow_schema,
8932+
):
8933+
type(invocation_context.session).events = mock.PropertyMock(
8934+
return_value=[self._pause_event("call-1")]
8935+
)
8936+
user_message = types.Content(
8937+
role="user",
8938+
parts=[
8939+
types.Part(
8940+
function_response=types.FunctionResponse(
8941+
id="call-1", name="long_search", response={"ok": True}
8942+
)
8943+
)
8944+
],
8945+
)
8946+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
8947+
await bq_plugin_inst.on_user_message_callback(
8948+
invocation_context=invocation_context, user_message=user_message
8949+
)
8950+
await bq_plugin_inst.flush()
8951+
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
8952+
completed = [r for r in rows if r["event_type"] == "TOOL_COMPLETED"]
8953+
assert len(completed) == 1
8954+
adk = json.loads(completed[0]["attributes"])["adk"]
8955+
assert adk["function_call_id"] == "call-1"
8956+
assert adk["pause_kind"] == "tool"
8957+
assert adk["pause_orphan"] is False
8958+
8959+
@pytest.mark.asyncio
8960+
async def test_emit_omits_pause_orphan_when_unknown(
8961+
self,
8962+
bq_plugin_inst,
8963+
mock_write_client,
8964+
invocation_context,
8965+
dummy_arrow_schema,
8966+
):
8967+
# No matching pause in history -> pause_orphan omitted (SQL NULL),
8968+
# never stamped true.
8969+
type(invocation_context.session).events = mock.PropertyMock(return_value=[])
8970+
user_message = types.Content(
8971+
role="user",
8972+
parts=[
8973+
types.Part(
8974+
function_response=types.FunctionResponse(
8975+
id="call-1", name="long_search", response={"ok": True}
8976+
)
8977+
)
8978+
],
8979+
)
8980+
bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context)
8981+
await bq_plugin_inst.on_user_message_callback(
8982+
invocation_context=invocation_context, user_message=user_message
8983+
)
8984+
await bq_plugin_inst.flush()
8985+
rows = await _get_captured_rows_async(mock_write_client, dummy_arrow_schema)
8986+
completed = [r for r in rows if r["event_type"] == "TOOL_COMPLETED"]
8987+
assert len(completed) == 1
8988+
adk = json.loads(completed[0]["attributes"])["adk"]
8989+
assert adk["function_call_id"] == "call-1"
8990+
assert "pause_orphan" not in adk

0 commit comments

Comments
 (0)