Skip to content

Commit 538ccce

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 538ccce

2 files changed

Lines changed: 232 additions & 8 deletions

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 62 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3241,6 +3241,53 @@ 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`` — the originating paused call is present in this
3258+
session's in-memory history, so this completion is definitely
3259+
not an orphan.
3260+
* ``None`` — unknown / not yet settled: there is no id to pair on,
3261+
no session history is available, or the pause was not found in
3262+
the (possibly trimmed) in-memory history. Callers omit the
3263+
``pause_orphan`` attribute in this case so consumers read it as
3264+
SQL NULL ("not yet determined").
3265+
3266+
Never returns ``True``. Declaring a true orphan requires the settled
3267+
fallback (an off-hot-path session-service / BigQuery reconciliation),
3268+
which is intentionally out of scope here so the resume path stays
3269+
free of added reads. A persistent ``None`` is what that future
3270+
fallback upgrades to ``True``.
3271+
"""
3272+
if not function_call_id:
3273+
return None
3274+
session = getattr(invocation_context, "session", None)
3275+
events = getattr(session, "events", None) if session is not None else None
3276+
if not events:
3277+
return None
3278+
for event in events:
3279+
ids = getattr(event, "long_running_tool_ids", None)
3280+
if not ids or function_call_id not in ids:
3281+
continue
3282+
content = getattr(event, "content", None)
3283+
parts = getattr(content, "parts", None) if content is not None else None
3284+
for part in parts or ():
3285+
function_call = getattr(part, "function_call", None)
3286+
if function_call is not None and function_call.id == function_call_id:
3287+
# The originating paused call is in this session's history.
3288+
return False
3289+
return None
3290+
32443291
@_safe_callback
32453292
async def on_user_message_callback(
32463293
self,
@@ -3300,26 +3347,33 @@ async def on_user_message_callback(
33003347
# tool calls complete inside the agent run via
33013348
# after_tool_callback, so a function_response inside a user
33023349
# message is the resume side of a previously-paused tool.
3303-
# Stamp the pair keys; pause_orphan / registry semantics
3304-
# are intentionally deferred.
33053350
if not part.function_response.id:
33063351
logger.debug(
33073352
"User-message function_response for tool %s has no id;"
33083353
" the resulting TOOL_COMPLETED row cannot pair with a"
33093354
" TOOL_PAUSED row.",
33103355
part.function_response.name,
33113356
)
3357+
adk_extras = {
3358+
"pause_kind": "tool",
3359+
"function_call_id": part.function_response.id,
3360+
}
3361+
# pause_orphan: stamped False only when this session's history
3362+
# proves the completion pairs with a real pause. Omitted (->
3363+
# SQL NULL = "unknown / not yet settled") otherwise. True is
3364+
# reserved for the off-hot-path settled fallback so the resume
3365+
# path stays free of session-service / BigQuery reads.
3366+
pause_orphan = self._resolve_same_session_pause_orphan(
3367+
invocation_context, part.function_response.id
3368+
)
3369+
if pause_orphan is not None:
3370+
adk_extras["pause_orphan"] = pause_orphan
33123371
await self._log_event(
33133372
"TOOL_COMPLETED",
33143373
callback_ctx,
33153374
raw_content=content_dict,
33163375
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-
),
3376+
event_data=EventData(adk_extras=adk_extras),
33233377
)
33243378

33253379
@_safe_callback

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

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

0 commit comments

Comments
 (0)