Resume fails with "Replay divergence detected" on SqliteSessionService when a parallel branch finishes during a request_input model call #7695

import asyncio
import tempfile
import time
from pathlib import Path

from google.adk.agents import Context, LlmAgent
from google.adk.apps import App
from google.adk.apps._configs import ResumabilityConfig
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.sqlite_session_service import SqliteSessionService
from google.adk.tools import request_input
from google.adk.workflow import START, JoinNode, Workflow, node
from google.genai import types


class FakeLlm(BaseLlm):
    model: str = "fake"
    calls: int = 0

    async def generate_content_async(self, llm_request, stream=False):
        self.calls += 1
        if self.calls == 1:
            # First call: asks the user for input after 1s.
            await asyncio.sleep(1)
            ask = types.FunctionCall(id="ask-1", name="adk_request_input", args={"message": "Which one?"})
            yield LlmResponse(content=types.Content(role="model", parts=[types.Part(function_call=ask)]))
        else:
            # After the resume: a call slower than the 15s replay barrier timeout.
            await asyncio.sleep(16)
            yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="done")]))


agent = LlmAgent(name="agent", model=FakeLlm(), instruction="Ask if unsure.", tools=[request_input])


@node(rerun_on_resume=True)
async def asker(ctx: Context, node_input):
    yield await ctx.run_node(agent, node_input="question")


@node(rerun_on_resume=True)
async def side_branch(ctx: Context, node_input):
    await asyncio.sleep(0.5)  # finishes while the asker's model call is still running
    yield "side done"


@node
async def end(ctx: Context, node_input):
    yield node_input


async def main():
    workflow = Workflow(name="root", edges=[(START, (asker, side_branch), JoinNode(name="join"), end)])
    app = App(name="repro", root_agent=workflow, resumability_config=ResumabilityConfig(is_resumable=True))
    sessions = SqliteSessionService(str(Path(tempfile.mkdtemp()) / "sessions.db"))
    runner = Runner(app=app, session_service=sessions)
    session = await sessions.create_session(app_name="repro", user_id="u")

    async def send(part):
        message = types.Content(role="user", parts=[part])
        return [e async for e in runner.run_async(user_id="u", session_id=session.id, new_message=message)]

    events = await send(types.Part(text="go"))
    ask = next(c for e in events for c in e.get_function_calls())

    reloaded = await sessions.get_session(app_name="repro", user_id="u", session_id=session.id)
    print("order after reload:", [e.node_info.path for e in reloaded.events if e.output is not None or e.get_function_calls()])

    reply = types.FunctionResponse(id=ask.id, name=ask.name, response={"result": "the first"})
    started = time.time()
    try:
        await send(types.Part(function_response=reply))
        print("resume OK")
    except Exception as exc:
        print(f"resume failed after {time.time() - started:.1f}s: {exc}")


asyncio.run(main())

class RestampAsk(BasePlugin):
    async def on_event_callback(self, *, invocation_context, event):
        if any(c.name == "adk_request_input" for c in event.get_function_calls()):
            event.timestamp = platform_time.get_time()