[Bug] google_adk: a retried stateful MCP session Activity silently opens a new, empty MCP session #310

Summary

TemporalStatefulMcpToolSet keeps one McpToolset open in its -server-session Activity for
the whole async with block. When Temporal retries that Activity, the retry calls the toolset
factory again, and the next tool call opens a new MCP session. Later tool calls in the same run
go to that new, empty session and succeed. The Workflow gets wrong results and no error.

The docstrings promise an error instead. TemporalStatefulMcpToolSetProvider says "the caller
must handle the case where that worker fails, as Temporal cannot seamlessly recreate the lost
toolset. Failure surfaces as an ApplicationError with type="DedicatedWorkerFailure""
(_mcp.py),
and TemporalStatefulMcpToolSet says the connection "is reused across every tool call in that
block"
(_mcp.py).

The session Activity runs with Temporal's default retry policy
(connect).
Its body refuses to start only while the same run's toolset is still registered in the process
(_mcp.py).
Nothing tells a retry apart from the first attempt, and ADK opens the MCP session lazily on the
first list or call.

Three ways to trigger the retry, the first two with no configuration:

• the host Worker shuts down gracefully, as in a deploy: the Activity is cancelled with
  worker_shutdown and retried on another Worker;
• the block outlives the session Activity's start-to-close timeout (one hour by default);
• server_session_config sets a heartbeat timeout and the host Worker dies.

Versions

• temporalio-google-adk 0.0.1 with google-adk 2.13.0, temporalio 1.35.0
• main at db14ec3 with its lock: google-adk 2.12.0, temporalio 1.34.0
• _mcp.py is the same in 0.0.1 and main
• mcp 1.30.0 (the script's server uses FastMCP, which mcp 2 renamed), Python 3.13.16,
  temporal server start-dev from Temporal CLI 1.9.1

Reproduction

One file, run against a local temporal server start-dev. In one process it runs a stateful MCP
server over streamable HTTP (put stores a value per MCP session, get returns it or EMPTY), a
Worker with GoogleAdkPlugin and TemporalStatefulMcpToolSetProvider, a fake model, and an ADK
agent. Turn 1 stores kept, the Workflow sleeps 45 seconds, and turn 2 reads it back. The control
run keeps the default session config. The other run sets the session Activity's start-to-close
timeout to 15 seconds, so Temporal retries it during the sleep. It exits 1 when turn 2 reads a new,
empty session.

pip install 'temporalio-google-adk[mcp]==0.0.1' 'mcp<2'
GOOGLE_API_USE_CLIENT_CERTIFICATE=false PYTHONWARNINGS=ignore python repro.py 127.0.0.1:7233

repro.py:

"""A retried stateful MCP session Activity silently opens a new, empty MCP session.

Run: python repro.py <temporal-address>

One process runs a stateful MCP server (``put`` stores a value per MCP session,
``get`` returns it or ``EMPTY``), a Worker with ``GoogleAdkPlugin`` and a
``TemporalStatefulMcpToolSetProvider``, and a Workflow whose ADK agent stores a
value in turn 1, sleeps 45 seconds, and reads it back in turn 2. The control
run keeps the default session config. The other run sets the session Activity's
start-to-close timeout to 15 seconds (the default is one hour), so Temporal
retries it during the sleep. Exits 1 if turn 2 reads a new, empty session.
"""

import asyncio
import json
import sys
import threading
import uuid
from datetime import timedelta

from temporalio import workflow
from temporalio.client import Client, WorkflowFailureError
from temporalio.worker import Worker

with workflow.unsafe.imports_passed_through():
    import uvicorn
    from google.adk import Agent
    from google.adk.models import BaseLlm, LLMRegistry
    from google.adk.models.llm_response import LlmResponse
    from google.adk.runners import InMemoryRunner
    from google.adk.tools.mcp_tool import McpToolset
    from google.adk.tools.mcp_tool.mcp_session_manager import StreamableHTTPConnectionParams
    from google.genai import types
    from mcp.server.fastmcp import Context, FastMCP

    from temporalio.google_adk import (
        GoogleAdkPlugin,
        TemporalModel,
        TemporalStatefulMcpToolSet,
        TemporalStatefulMcpToolSetProvider,
    )
    from temporalio.workflow import ActivityConfig

PORT = 7606
server = FastMCP("state", json_response=True, session_idle_timeout=None, log_level="WARNING")
store: dict[str, str] = {}
sessions: list[str] = []


@server.tool()
def put(value: str, ctx: Context) -> str:
    sid = ctx.request_context.request.headers["mcp-session-id"]
    sessions.append(f"put on session {sid[:8]}")
    store[sid] = value
    return "stored"


@server.tool()
def get(ctx: Context) -> str:
    sid = ctx.request_context.request.headers["mcp-session-id"]
    sessions.append(f"get on session {sid[:8]}")
    return store.get(sid, "EMPTY")


class Model(BaseLlm):
    """Calls the tool named in the user message, then answers with its output or error."""

    @classmethod
    def supported_models(cls) -> list[str]:
        return ["repro-model"]

    async def generate_content_async(self, llm_request, stream=False):
        last = llm_request.contents[-1]
        outputs = [p.function_response.response for p in last.parts or [] if p.function_response]
        if outputs:
            out = outputs[0]
            part = types.Part(text=out["content"][0]["text"] if "content" in out else json.dumps(out))
        else:
            ask = json.loads(last.parts[0].text)
            part = types.Part(function_call=types.FunctionCall(name=ask["tool"], args=ask["args"]))
        yield LlmResponse(content=types.Content(role="model", parts=[part]))


@workflow.defn
class Repro:
    @workflow.run
    async def run(self, session_timeout_seconds: int | None) -> list[str]:
        session_config = None
        if session_timeout_seconds:
            session_config = ActivityConfig(start_to_close_timeout=timedelta(seconds=session_timeout_seconds))
        async with TemporalStatefulMcpToolSet("state", server_session_config=session_config) as toolset:
            agent = Agent(name="repro", model=TemporalModel("repro-model"), tools=[toolset])
            runner = InMemoryRunner(agent=agent, app_name="repro")
            session = await runner.session_service.create_session(app_name="repro", user_id="u")

            async def turn(tool: str, args: dict) -> str:
                message = types.Content(role="user", parts=[types.Part(text=json.dumps({"tool": tool, "args": args}))])
                texts = []
                async for event in runner.run_async(user_id="u", session_id=session.id, new_message=message):
                    texts += [p.text for p in (event.content.parts if event.content else []) or [] if p.text]
                return texts[-1]

            first = await turn("put", {"value": "kept"})
            await workflow.sleep(45)
            second = await turn("get", {})
        return [first, second]


def factory(_):
    return McpToolset(connection_params=StreamableHTTPConnectionParams(url=f"http://127.0.0.1:{PORT}/mcp"))


async def main() -> int:
    kwargs = {"host": "127.0.0.1", "port": PORT, "log_level": "warning"}
    threading.Thread(target=uvicorn.run, args=(server.streamable_http_app(),), kwargs=kwargs, daemon=True).start()
    LLMRegistry.register(Model)
    plugin = GoogleAdkPlugin(toolset_providers=[TemporalStatefulMcpToolSetProvider("state", factory)])
    client = await Client.connect(sys.argv[1], plugins=[plugin])
    queue = f"repro-{uuid.uuid4().hex[:6]}"
    defect = False
    async with Worker(client, task_queue=queue, workflows=[Repro]):
        for name, timeout in [("control, default session config", None), ("session start-to-close 15 s", 15)]:
            sessions.clear()
            try:
                result = await client.execute_workflow(Repro.run, timeout, id=f"{queue}-{timeout}", task_queue=queue)
                print(f"{name}: turn 1 put -> {result[0]}, turn 2 get -> {result[1]} ({', '.join(sessions)})")
                defect = defect or result[1] == "EMPTY"
            except WorkflowFailureError as e:
                cause = e.cause
                while getattr(cause, "cause", None) is not None:
                    cause = cause.cause
                print(f"{name}: run failed with {type(cause).__name__} {getattr(cause, 'type', '')}: {cause} ({', '.join(sessions)})")
    if defect:
        print("DEFECT: turn 2 read a new, empty MCP session and the run completed with no error")
    return 1 if defect else 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))

Expected

Turn 2 reads kept from the same session. If the session cannot be kept, the run fails with
ApplicationError(type="DedicatedWorkerFailure"), as the docstrings say.

Observed

On 0.0.1, and the same on main:

control, default session config: turn 1 put -> stored, turn 2 get -> kept (put on session 1fb1e368, get on session 1fb1e368)
session start-to-close 15 s: turn 1 put -> stored, turn 2 get -> EMPTY (put on session c02813b4, get on session 497d1de2)
DEFECT: turn 2 read a new, empty MCP session and the run completed with no error

The script exits with status 1.

In the Worker log, attempts 2 and 3 of the session Activity fail with Cannot connect to an already running toolset while attempt 1 is still cleaning up, and attempt 4 calls the factory again; its
toolset opened the session that served get.

Controls

• The control run, with no retry during the sleep, reads its own value (first line).
• The same happens with a graceful Worker shutdown between the turns, with no configuration at all:
  the session Activity is retried on the new Worker about one second later, and turn 2 reads
  EMPTY from a session the new attempt opened. A test in the PR does this with a stdio MCP server.
• A retry before any session was opened is fine and must keep working: a factory that fails, or a
  first connection that fails, is retried today and builds a working toolset.

Impact

State held by a stateful MCP server (a browser session, an open transaction, a scratch file, a
login) is lost in the middle of a run, and the agent goes on with empty state. Nothing reports it:
tool calls succeed and the run completes. An ordinary deploy is enough to trigger it.

Proposed fix

Fail a retried session Activity with a non-retryable DedicatedWorkerFailure when an earlier
attempt opened an MCP session, and keep the other retries:

1. The session Activity records a marker in its own heartbeat details, with activity.heartbeat(),
   when its toolset opens an MCP session. It wraps create_session() on the toolset's
   MCPSessionManager, which McpToolset and McpTool call before every request, and records
   the marker after ADK opens a session and before ADK sends a request on it. A toolset with no
   session manager records it before its first operation.
2. Heartbeat details carry over to the next attempt. An attempt that finds the marker fails with
   ApplicationError(type="DedicatedWorkerFailure", non_retryable=True) before it calls the
   factory.
3. Only the Worker writes the details, and it sends an attempt's latest details with its failure,
   also for a graceful shutdown. Every heartbeat after the marker carries it, so an earlier
   heartbeat cannot erase it. Cancels, pauses and resets reach the session Activity as before.
4. The Worker does not let the plugin wait for a heartbeat, and it can hold the marker heartbeat
   for up to 30 seconds, so a host can be lost after opening a session but before the marker
   reaches the server. So every attempt sends its first heartbeat at once, marked "a session may
   be open", and keeps that mark until a session opens. Before the first session opens, the
   session Activity reads DescribeWorkflowExecution until the server shows that heartbeat. A
   retry that finds the mark is refused too, even if no session was opened. An attempt that ends
   without opening a session clears it, after its earlier heartbeats finish encoding so none of
   them can replace the clearing one. So a factory or first connection that failed, or a graceful
   shutdown before any session, is still retried. The Worker's client needs permission to describe
   the Workflow.
5. In Workflow code, get_tools() and stateful tool calls also watch the session Activity. If it
   ends first, a waiting list or call fails at once with DedicatedWorkerFailure, with the
   session's failure as the cause, instead of after its own schedule-to-start or start-to-close
   timeout. Its Activity is cancelled, unless its cancellation type is
   WAIT_CANCELLATION_COMPLETED: the lost Worker would never confirm that cancel. A list or call
   started after the session ended fails at once.
6. Every async with block of a run uses the same task queue. The Worker refuses a list or call
   that the server scheduled before the current session Activity, because it belongs to an earlier
   block, which has ended.

Recording the marker before the first operation, or failing every retry, would also refuse retries
in which no session existed, such as a first connection that failed.

With this change, the script's second run fails with DedicatedWorkerFailure. On the agent route,
ADK runs turn 2 without the toolset, whose tool listing now fails, and logs it; the run fails when
the async with block exits.

This changes callers: a factory whose state survives being rebuilt, for example one that reloads
it from storage, recovers today and gets DedicatedWorkerFailure after the fix. A lost or timed-out
attempt is refused even before a session opened. Such a caller can move to the stateless
TemporalMcpToolSetProvider, or catch the failure around the block and open a new
TemporalStatefulMcpToolSet block. The refusal before any session opened can go once the Python
SDK offers a heartbeat the plugin can wait for.

It also changes replay, with no version gate: Workflow code changes (item 5), so a run whose
session Activity ended while it still used the toolset does not replay across the upgrade or a
rollback. Such runs are already failing with DedicatedWorkerFailure or a timeout; let them end
before upgrading. Finished and waiting runs whose session did not end replay on both builds.

#311 implements it.