fix: scope tool confirmation replay claims by session

[AnvitDevadiga]

Summary

Scope in-process confirmation claims by app, user, session, branch, and tool-call ID so independent sessions cannot consume each other’s approvals. Release a claim after its tool response is persisted.

Add public-Runner coverage for same-ID approvals across separate sessions, users, and apps, plus controls for different sessions and different call IDs. The tests assert one execution per approval and no duplicate confirmation prompt on replay.

Validation

• 107 related confirmation, runner, and invocation-context tests passed.
• Repository pre-commit checks passed.

Scope

The claim registry protects concurrent invocations within one process; it does not provide cross-process or multi-replica exactly-once guarantees. The original denial gate is already present on main; this PR focuses on replay protection discussed in #7148.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[babyblueviper1]

Good fix — moving denial enforcement to the framework boundary instead of trusting each BaseTool to reimplement it is the right shape (matches the general lesson: a verdict object is only as strong as the code path that's forced to consume it, not the code path that happens to check it).

One adjacent question worth asking before merge, not asserted as a bug since I haven't reproduced it against this exact tree: does the resume path guard against the same function_call_id/ToolConfirmation being submitted more than once? If a client can call the resume endpoint twice with an identical confirmed=True payload for a function_call_id that's already been executed, does _resolve_confirmation_targets re-execute it, or is there state elsewhere marking that confirmation as consumed?

Reason I ask: we hit the identical shape in our own verdict-issuing service — a pure function computing a decision reference had no state, so an identical still-valid verdict could be replayed to authorize a second dispatch (caught via a negative-control test explicitly checking for replay, not by inspection). The fix there was a small addition on top of the same "enforce at the boundary" pattern this PR already uses: a primary-keyed consumed-set with atomic check-and-mark, so a race between two identical resume calls has exactly one winner instead of two executions. If tools_to_resume_with_confirmation here is keyed by function_call_id and nothing marks an id as spent after this code path runs, the same gap likely exists — cheap to close now while the boundary-enforcement code is already being touched, much easier than adding it after the fact once callers start depending on today's semantics.

[Tardfyou]

Following @surajksharma07's request in #7148, I validated the proposed shared-caller workaround locally on google-adk==2.9.1 using the original deterministic fake-LLM/custom-BaseTool harness. The only effect is a temporary marker file; no external model or service is contacted.

Results:

Caller	User verdict	Tool executions after verdict	Marker written
Unmodified 2.9.1	confirmed=False	1	yes
Caller denies confirmed=False before run_async	confirmed=False	0	no
Same workaround	confirmed=True	1	yes
Unmodified 2.9.1	confirmed=True	1	yes

All four assertions passed. This validates the suggested caller-level workaround for the reported custom-tool path.

One distinction matters before treating this as validation of this PR: the current head 3557a41 filters denied calls in tools/_confirmation.py, whereas the workaround I tested guards _call_tool_async (in _tool_caller.py in release 2.9.1). I have not run this PR head, so this comment is not an end-to-end verification of #7151. Please retain a runner-level regression with a custom BaseTool that does not inspect the verdict, plus the approved control. The replay question above is also outside the scope of my four-case check.

[babyblueviper1]

Verified this against the actual PR head (6627ed4), not just the comment thread — the fix is the right shape: _consumed_tool_confirmation_ids is a primary-keyed set on InvocationContext, claimed under _tool_confirmation_consume_lock before dispatch, exactly the atomic check-and-mark pattern that closes the race from my original question (a confirmed=True payload for the same function_call_id submitted twice, one winner).

I went one step further before calling it closed: does this survive a crash/resume into a new process? _consumed_tool_confirmation_ids is a Pydantic PrivateAttr on InvocationContext, and private attrs don't round-trip through model serialization — so a fresh process resuming a persisted session gets an empty consumed-set. Read through _confirmation.py's full pipeline to check whether that reopens anything, and it doesn't: Step 2 (event-history dedup, responded_fc_ids from the persisted event log) runs before Step 3's in-memory claim and drops any confirmation that already has a recorded function_response — that check is durable across process boundaries because the event log is what gets persisted, not the in-memory set. So the in-memory set only has to be correct for races within one process's lifetime on a still-unanswered confirmation, which is exactly what it does. No residual gap at the resume boundary — the two checks are at different scopes and compose correctly.

What the PR doesn't have yet: a direct regression test for the concurrency race itself. The diff's two test changes just update expected deny-path messages — nothing exercises two simultaneous calls into _consume_tool_confirmation (or two concurrent resume calls) for the same function_call_id to assert exactly one wins. Happy to write that against 6627ed4 (asyncio.gather of two confirm-and-dispatch attempts, assert one executes and one no-ops) — same pattern we run internally for our own consumed-set on decision-reference reuse. Say the word and I'll open it as a PR.

[gomission]

@AnvitDevadiga I ran the latest head 4beee3a following @surajksharma07's cross-invocation check in #7148. There is a first-approval regression to resolve before treating the replay result as protection.

The initial confirmation gate already records a function response for the original call id:

{"error":"This tool call requires confirmation, please approve or reject."}

Step 2 now scans the whole current-branch history and adds that pending-response id to responded_fc_ids. The first confirmed=True submission is consequently dropped before dispatch. Using a scripted model, public Runner.run_async, a custom BaseTool, and a local counter:

Input sequence	Main 044a1ec executions	PR 4beee3a executions
Deny once	0	0
Approve once	1	0
Approve, then replay after the completed result is recorded	2	0*
Overlapping approval submissions	2	0*

* On this PR the first call never dispatches, so these zeros do not validate replay suppression. On main, the concurrent probe holds the first tool before it returns, starts the second resume, and observes two distinct InvocationContext objects enter the tool before either completed result is recorded.

The existing tests reproduce the regression without my harness:

python -m pytest -q \
  tests/unittests/flows/llm_flows/tools/test_confirmation.py \
  tests/unittests/runners/test_run_tool_confirmation.py

Main: 31 passed. PR head: 12 failed, 19 passed, including all eight Runner confirmation cases. Both runs used Python 3.12.11, google-genai 2.25.0 and the same dependency environment. The original report and denial workaround are @Tardfyou's work; this check covers the current replay patch, whose denial changes have already been removed.

The deduplication regression needs a positive first-approval control that retains the real pending-confirmation response in history. I haven't changed the implementation. These are local SDK checks, not external-provider exactly-once or process-restart proof. Disclosure: Mission; investigated with Codex.

[gomission]

@AnvitDevadiga Thanks for the correction. I reran the two cited test files on current head 1782ccac (including 399e1ced): 32 passed. The independent positive control now executes the first approved call exactly once, so I can confirm the first-approval regression is fixed in this head.

The same unchanged public-Runner fixture still finds a separate cross-invocation replay gap:

Input sequence	Tool executions on 1782ccac
Deny once	0
Approve once	1
Approve, then resubmit the same confirmation ID in a new invocation	2
Overlapping submissions of that same confirmation ID	2

For sequential replay, the first completed counter response is already in session history before the second submission. For the overlapping case, the first tool is held at an asyncio.Event barrier; a second Runner.run_async invocation enters the tool with a distinct InvocationContext before either tool returns its completed response.

This matches the current boundaries: the latest-user cutoff also excludes a completed response preceding the new duplicate submission, while the consumed set and lock belong to each InvocationContext. The first-approval correction is verified; these results do not support cross-invocation replay suppression yet. Full unchanged fixture below, so the remaining two failures are reproducible.

If our report/reproduction contributes to the landed fix, could you include “First-approval regression reported and independently reproduced by @gomission” in the PR description or final commit message? The implementation credit remains yours, and @Tardfyou retains credit for the original denial report.

These checks use InMemorySessionService, a scripted model and a local counter, with network calls blocked. They establish local SDK behavior, not process-restart durability or external-provider exactly-once execution. Disclosure: Mission; investigated with Codex.

Frozen four-case reproduction: 2 passed, 2 failed on 1782cca

Save as test_confirmation_replay.py in a checkout of 1782ccacdca6c466b2053c0e8563d6eb72509fe5. With the SDK test dependencies installed:

PYTHONPATH=src python -m pytest -q -s test_confirmation_replay.py

The failing cases are sequential_replay and concurrent_replay, both assert 2 == 1.

"""Local public-Runner confirmation replay checks for google/adk-python #7151.

Synthetic BaseLlm and a local counter only. No model, credentials, network,
remote tool, or provider effect. No SDK code is patched. The concurrent case
holds both tool invocations before either returns its function response.
"""

import asyncio
import json
import socket

from google.adk.agents import LlmAgent
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools import BaseTool
from google.genai import types
from pydantic import PrivateAttr
import pytest


@pytest.fixture(autouse=True)
def prohibit_network(monkeypatch):
    def blocked(*args, **kwargs):
        raise AssertionError("Network access is prohibited in this local reproduction")

    monkeypatch.setattr(socket.socket, "connect", blocked)
    monkeypatch.setattr(socket.socket, "connect_ex", blocked)
    monkeypatch.setattr(socket, "getaddrinfo", blocked)


class ScriptedModel(BaseLlm):
    _requested: bool = PrivateAttr(default=False)

    async def generate_content_async(self, llm_request, stream=False):
        if not self._requested:
            self._requested = True
            part = types.Part(
                function_call=types.FunctionCall(
                    name="local_counter", id="original-local-call", args={}
                )
            )
        else:
            part = types.Part(text="done")
        yield LlmResponse(content=types.Content(role="model", parts=[part]))


class CounterTool(BaseTool):
    def __init__(self):
        super().__init__(
            name="local_counter", description="Increment a local test counter"
        )
        self.executions = []
        self.contexts = []
        self.first_started = asyncio.Event()
        self.second_started = asyncio.Event()
        self.release = asyncio.Event()
        self.release.set()

    def _get_declaration(self):
        return types.FunctionDeclaration(name=self.name, description=self.description)

    async def check_require_confirmation(self, args, tool_context):
        return True

    async def run_async(self, *, args, tool_context):
        if tool_context.tool_confirmation is None:
            tool_context.request_confirmation(hint="Allow one local test increment?")
            return {"status": "awaiting_confirmation"}
        # Deliberately no self-enforced verdict check: the main framework
        # gate must handle deny. Observe context identity without modifying it.
        self.contexts.append(tool_context._invocation_context)
        self.executions.append(tool_context.tool_confirmation.confirmed)
        self.first_started.set()
        if len(self.executions) == 2:
            self.second_started.set()
        await self.release.wait()
        return {"count": len(self.executions)}


async def run_case(case):
    tool = CounterTool()
    agent = LlmAgent(
        name="replay_probe", model=ScriptedModel(model="local"), tools=[tool]
    )
    service = InMemorySessionService()
    runner = Runner(app_name="replay_probe", agent=agent, session_service=service)
    session = await service.create_session(app_name="replay_probe", user_id="local")

    async def invoke(message):
        return [
            event
            async for event in runner.run_async(
                user_id="local", session_id=session.id, new_message=message
            )
        ]

    pending_tasks = []
    try:
        initial = await invoke(
            types.Content(role="user", parts=[types.Part(text="request")])
        )
        confirmation_ids = [
            fc.id
            for event in initial
            for fc in event.get_function_calls()
            if fc.name == "adk_request_confirmation"
        ]
        assert len(confirmation_ids) == 1 and not tool.executions
        confirmation_id = confirmation_ids[0]
        initial_responses = [
            fr.response
            for event in initial
            for fr in event.get_function_responses()
            if fr.id == "original-local-call"
        ]

        def decision(confirmed=True):
            return types.Content(
                role="user",
                parts=[
                    types.Part(
                        function_response=types.FunctionResponse(
                            name="adk_request_confirmation",
                            id=confirmation_id,
                            response={"confirmed": confirmed},
                        )
                    )
                ],
            )

        responses_before_replay = None
        concurrent_second_entered = None
        first_resume_no_dispatch = False
        if case == "deny":
            await invoke(decision(False))
        elif case == "approve":
            await invoke(decision())
        elif case == "sequential_replay":
            await invoke(decision())
            stored = await service.get_session(
                app_name="replay_probe", user_id="local", session_id=session.id
            )
            responses_before_replay = sum(
                fr.id == "original-local-call" and "count" in fr.response
                for event in stored.events
                for fr in event.get_function_responses()
            )
            assert responses_before_replay == len(tool.executions)
            await invoke(decision())
        elif case == "concurrent_replay":
            tool.release.clear()
            first = asyncio.create_task(invoke(decision()))
            started = asyncio.create_task(tool.first_started.wait())
            pending_tasks.extend([first, started])
            done, _ = await asyncio.wait(
                {first, started}, timeout=5, return_when=asyncio.FIRST_COMPLETED
            )
            assert done, (
                "First resume neither dispatched nor finished within the timeout"
            )
            if first in done and not tool.first_started.is_set():
                await (
                    first
                )  # Surface an exception if completion was actually a failure.
                first_resume_no_dispatch = True
            stored = await service.get_session(
                app_name="replay_probe", user_id="local", session_id=session.id
            )
            responses_before_replay = sum(
                fr.id == "original-local-call" and "count" in fr.response
                for event in stored.events
                for fr in event.get_function_responses()
            )
            assert responses_before_replay == 0, (
                "First tool must still be awaiting its barrier"
            )
            second = asyncio.create_task(invoke(decision()))
            entered = asyncio.create_task(tool.second_started.wait())
            pending_tasks.extend([second, entered])
            done, _ = await asyncio.wait(
                {second, entered}, timeout=5, return_when=asyncio.FIRST_COMPLETED
            )
            concurrent_second_entered = entered in done
            # A future correct implementation may reject/defer the second call.
            # Release the first regardless, then verify actual execution count.
            tool.release.set()
            await asyncio.wait_for(asyncio.gather(first, second), timeout=5)
            if not entered.done():
                entered.cancel()
                await asyncio.gather(entered, return_exceptions=True)
        else:
            raise AssertionError(case)

        distinct_contexts = len({id(context) for context in tool.contexts})
        if concurrent_second_entered:
            assert distinct_contexts == 2, (
                "The race must involve distinct InvocationContext objects"
            )
        return {
            "case": case,
            "executions": len(tool.executions),
            "verdicts_seen_by_tool": tool.executions,
            "distinct_execution_contexts": distinct_contexts,
            "initial_responses": initial_responses,
            "completed_responses_before_replay": responses_before_replay,
            "first_resume_completed_without_dispatch": first_resume_no_dispatch,
            "second_entered_before_first_returned": concurrent_second_entered,
        }
    finally:
        tool.release.set()
        for task in pending_tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*pending_tasks, return_exceptions=True)
        await runner.close()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "case", ["deny", "approve", "sequential_replay", "concurrent_replay"]
)
async def test_public_runner_confirmation_replay(case):
    result = await run_case(case)
    print("OBSERVED " + json.dumps(result, sort_keys=True))
    assert result["executions"] == (0 if case == "deny" else 1)

[gomission]

@AnvitDevadiga Confirmed on b5bebd0f6e0fad4d58c6eb29b7fdc376f708f663: the same unchanged four-case public-Runner fixture now passes all four cases.

Case	Tool executions
Deny	0
First approval	1
Sequential replay after the completed result is recorded	1
Overlapping submissions while the first tool is held	1

The two cited upstream test files also pass here: 32 passed. In the overlapping case, the second invocation no longer enters the tool before the first returns. Thanks for addressing both failures while preserving the positive approval control.

This verifies the reported local cases using one process and a shared InMemorySessionService; it is not a cross-process or external-provider exactly-once claim. Disclosure: Mission; independently checked with Codex.

[gomission]

@AnvitDevadiga Rechecked 3b691362208e4097650fea41a5051b1d7b4442d2: the unchanged four-case Runner fixture still passes (deny/approve/sequential replay/overlapping replay = 0/1/1/1), and 99 related upstream tests pass here.

I found one regression in the new claim scope: a legitimate first approval in another user's or app's session can be silently consumed when both sessions have the same session ID, branch and original tool-call ID.

The public-Runner fixture below creates two independent sessions on one InMemorySessionService, each with its own generated confirmation-request ID. It holds A's approved tool in flight, submits B's first approval, releases A, then retries B. Both should execute once:

Separation between sessions	After first approvals, A/B	After B retry, A/B
Different session IDs	1 / 1	1 / 1
Different users; same session/tool-call IDs	1 / 0	1 / 0
Different apps; same session/tool-call IDs	1 / 0	1 / 0
Different users and tool-call IDs	1 / 1	1 / 1

The same fixture passes 4/4 on parent 3f514e2ee939bef761215610de2518512cb60a06, but gives 2 failed / 2 passed on this head. Session lookup is scoped by (app_name, user_id, session_id); the new claim key only uses (session_id, branch, function_call_id). B therefore collides with A before its own persisted history is checked. B's user confirmation is already appended, so the new replay filter also suppresses the retry after A finishes.

Including app_name and user_id in both claim creation and release fixes this locally. With only the patch below, the 99 upstream tests, original four cases and new four namespace cases all pass: 107 passed. Restoring the unmodified head restores the two failures. The local patch below is for your existing PR.

Environment: Python 3.12.11, google-genai 2.28.0, pydantic 2.13.5, pytest 9.1.1, pytest-asyncio 1.4.0. One event loop, scripted model, local counters and blocked network connections; no provider or cross-process guarantee. Disclosure: Mission; independently checked with Codex.

Minimal session-key patch

diff --git a/src/google/adk/agents/invocation_context.py b/src/google/adk/agents/invocation_context.py
index 0ba11ffe..c80c8c79 100644
--- a/src/google/adk/agents/invocation_context.py
+++ b/src/google/adk/agents/invocation_context.py
@@ -56,7 +56,7 @@ _claimed_tool_confirmations: dict[
     int,
     tuple[
         weakref.ReferenceType[BaseSessionService],
-        set[tuple[str, str | None, str]],
+        set[tuple[str, str, str, str | None, str]],
     ],
 ] = {}
 
@@ -238,7 +238,13 @@ class InvocationContext(BaseModel):
         service_claims = (service_ref, set())
         _claimed_tool_confirmations[service_key] = service_claims
       claims = service_claims[1]
-      claim_key = (self.session.id, self.branch, function_call_id)
+      claim_key = (
+          self.session.app_name,
+          self.session.user_id,
+          self.session.id,
+          self.branch,
+          function_call_id,
+      )
       if claim_key in claims:
         return False
       claims.add(claim_key)
@@ -305,7 +311,13 @@ class InvocationContext(BaseModel):
 
       claims = service_claims[1]
       claims.difference_update(
-          (self.session.id, branch, function_call_id)
+          (
+              self.session.app_name,
+              self.session.user_id,
+              self.session.id,
+              branch,
+              function_call_id,
+          )
           for function_call_id in function_call_ids
       )
       if not claims:

Runnable public-Runner regression and controls

Save as test_confirmation_namespace.py and run with the PR checkout installed:

python -m pytest -q -s test_confirmation_namespace.py

"""Public-Runner session namespace controls for ADK PR #7151.

One event loop, shared InMemorySessionService, scripted LLM, local effects.
No SDK internals are patched; both decisions are first approvals of distinct
confirmation requests in separate sessions. Network connections are prohibited.
"""

import asyncio
import json
import socket

from google.adk.agents import LlmAgent
from google.adk.models.base_llm import BaseLlm
from google.adk.models.llm_response import LlmResponse
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.adk.tools.function_tool import FunctionTool
from google.genai import types
from pydantic import PrivateAttr
import pytest


@pytest.fixture(autouse=True)
def prohibit_network(monkeypatch):
    def blocked(*args, **kwargs):
        raise AssertionError("Network access is prohibited in this reproduction")

    monkeypatch.setattr(socket.socket, "connect", blocked)
    monkeypatch.setattr(socket.socket, "connect_ex", blocked)
    monkeypatch.setattr(socket, "getaddrinfo", blocked)


class ScriptedModel(BaseLlm):
    call_id: str = "local-call"
    _requested: bool = PrivateAttr(default=False)

    async def generate_content_async(self, llm_request, stream=False):
        if self._requested:
            part = types.Part(text="done")
        else:
            self._requested = True
            part = types.Part(function_call=types.FunctionCall(
                name="local_counter", id=self.call_id, args={}
            ))
        yield LlmResponse(content=types.Content(role="model", parts=[part]))


@pytest.mark.asyncio
@pytest.mark.parametrize("separation", ["session", "user", "app", "user_and_call"])
async def test_independent_first_approvals(separation):
    service = InMemorySessionService()
    namespaces = [("app_a", "user-a", "session-a"), {
        "session": ("app_a", "user-a", "session-b"),
        "user": ("app_a", "user-b", "session-a"),
        "app": ("app_b", "user-a", "session-a"),
        "user_and_call": ("app_a", "user-b", "session-a"),
    }[separation]]
    effects = [0, 0]
    entered = asyncio.Event()
    release = asyncio.Event()
    runners, approvals, tasks = [], [], []

    def make_tool(index):
        async def local_counter() -> dict:
            effects[index] += 1
            if index == 0:
                entered.set()
                await release.wait()
            return {"executions": effects[index]}

        return FunctionTool(func=local_counter, require_confirmation=True)

    async def invoke(index, message):
        _, user_id, session_id = namespaces[index]
        return [event async for event in runners[index].run_async(
            user_id=user_id, session_id=session_id, new_message=message
        )]

    try:
        for index, (app_name, user_id, session_id) in enumerate(namespaces):
            await service.create_session(
                app_name=app_name, user_id=user_id, session_id=session_id
            )
            call_id = "other-call" if separation == "user_and_call" and index else "local-call"
            runners.append(Runner(
                app_name=app_name, session_service=service,
                agent=LlmAgent(name="root_agent", tools=[make_tool(index)],
                               model=ScriptedModel(model="local", call_id=call_id)),
            ))
            events = await invoke(index, types.Content(
                role="user", parts=[types.Part(text="request")]
            ))
            ids = [fc.id for event in events for fc in event.get_function_calls()
                   if fc.name == "adk_request_confirmation"]
            assert len(ids) == 1
            approvals.append(types.Content(role="user", parts=[types.Part(
                function_response=types.FunctionResponse(
                    name="adk_request_confirmation", id=ids[0],
                    response={"confirmed": True},
                )
            )]))
        assert effects == [0, 0]
        assert approvals[0].parts[0].function_response.id != approvals[1].parts[0].function_response.id

        tasks.append(asyncio.create_task(invoke(0, approvals[0])))
        await asyncio.wait_for(entered.wait(), timeout=5)
        # First tool is in flight; the second belongs to an independent session.
        await asyncio.wait_for(invoke(1, approvals[1]), timeout=5)
        after_first_approvals = list(effects)
        release.set()
        await asyncio.wait_for(tasks[0], timeout=5)
        # Check whether a previously blocked approval recovers once A completes.
        await asyncio.wait_for(invoke(1, approvals[1]), timeout=5)
        print("OBSERVED " + json.dumps({
            "separation": separation,
            "namespaces": namespaces,
            "effects_after_first_approvals": after_first_approvals,
            "effects_after_retry": effects,
        }, sort_keys=True))
        assert after_first_approvals == [1, 1]
        assert effects == [1, 1]
    finally:
        release.set()
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        for runner in runners:
            await runner.close()

Related upstream files checked:

python -m pytest -q \
  tests/unittests/flows/llm_flows/tools/test_confirmation.py \
  tests/unittests/runners/test_run_tool_confirmation.py \
  tests/unittests/agents/test_invocation_context.py \
  tests/unittests/agents/test_invocation_context_process_queue.py

[AnvitDevadiga]

Thanks for the detailed reproduction. I rebased onto current main and scoped the in-flight claim and release key by app, user, session, branch, and tool-call ID. I added public-Runner coverage for independent session namespaces and call-ID controls; it verifies one execution per approval, persisted tool responses, and that replay does not emit another confirmation prompt. The four focused suites pass (107 tests), and pre-commit passes. This remains an in-process concurrency guarantee, not a cross-process exactly-once guarantee.