[campoy]

What was changed

A retried stateful MCP session Activity no longer opens a new, empty MCP session in the middle of a run. If an earlier attempt opened a session, the retry fails with a non-retryable ApplicationError(type="DedicatedWorkerFailure"), as the TemporalStatefulMcpToolSetProvider docstring says. Retries in which no session was opened still build the toolset, except after a lost host or a timed-out attempt (below).

• The -server-session Activity wraps create_session() on the toolset's MCPSessionManager, which ADK's McpToolset and McpTool call before every request. After ADK returns a session, and before ADK sends a request on it, the wrapper records mcp-session-opened in the Activity's heartbeat details with activity.heartbeat(). A toolset with no session manager (such as a test fake) records it before its first list or call.
• Heartbeat details carry over between attempts. An attempt that finds the marker fails with DedicatedWorkerFailure, non_retryable=True, before it calls the factory.
• Only the Worker writes these details, so cancels, pauses, resets and heartbeat timeouts reach the session Activity as before. The Worker sends its latest details with the attempt's failure, including the failure it reports for a graceful shutdown, so a graceful deploy passes the marker on.
• The Worker does not let the plugin wait for a heartbeat: it sends the marker when its throttle allows (up to 30 seconds later without a heartbeat timeout), 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, carrying mcp-session-may-be-open, and every heartbeat carries it until a session opens. Before the first session opens, the session Activity reads DescribeWorkflowExecution until the server shows that heartbeat. A retry that finds it is refused too.
• An attempt that ends without opening a session (a factory or first connection that failed, a graceful shutdown) clears that marker, and the Worker sends the clear with the failure, so those retries still build the toolset. The Worker encodes each heartbeat in its own task, so heartbeats can reach it out of order. The attempt therefore waits for any heartbeat still being sent before it sends the clear. "May be open" and "opened" both refuse a retry, so their order does not matter.
• Workflow side: get_tools and stateful tool calls fail with DedicatedWorkerFailure once the session Activity has ended, with the session's failure as the cause. The caller does not wait for an operation that is still pending when the session ends: it fails at once, and the operation is cancelled. With cancellation_type=WAIT_CANCELLATION_COMPLETED the operation is not cancelled but left to its own timeout, because no worker may be left to acknowledge the cancel, and an unacknowledged cancel keeps the Workflow from being evicted. Such an operation stays on the name@run_id queue, which a later async with block with the same name also serves. So the session worker refuses, with a non-retryable DedicatedWorkerFailure, any operation scheduled before its own session Activity: a run numbers its activities in order, and that number is the default activity ID. No operation is scheduled after the session ended. They no longer depend on an optional schedule_to_start_timeout, which a caller-supplied config may leave out. There is no version gate (see Replay below).
• The google_adk README describes this retry behaviour and the heartbeat limitation.

The marker is stored after ADK opens a session, not before the first operation: a first connection that fails raises before ADK returns a session, so its retry keeps working. Recording it after the operation would let a host lost after a tool call's side effect be rebuilt.

Earlier versions of this PR sent the marker with awaited client heartbeats as well as the Worker's. With two writers, a Worker heartbeat sent before the session opened could reach the server after the marker and erase it. Without a heartbeat timeout, resets also stopped reaching the session. This version has one writer and keeps the Worker's heartbeats.

What callers see after a refused retry: the run's next tool call, or one already waiting, fails with DedicatedWorkerFailure without waiting for a timeout. On the agent route, ADK catches that error while preparing the turn, logs it, and runs the turn without the toolset (google/adk-python does this for any toolset that fails to load). The run then fails when the async with block exits, because cleanup() re-raises the session Activity's failure as an ActivityError whose cause is the DedicatedWorkerFailure.

Breaking change (callers): a factory whose state survives being rebuilt, for example one that reloads it from storage, recovers today when the session Activity is retried. After this change, once a session was opened, that run gets DedicatedWorkerFailure instead. This also happens on a graceful deploy and after an Activity reset that keeps heartbeat details (reset_heartbeat=False). It also gets that error when its host is lost, or the attempt times out, before any session opened. Before upgrading, such a caller can move to the stateless TemporalMcpToolSetProvider and TemporalMcpToolSet, which the provider's docstring prefers when no persistent connection is needed, or catch the failure around the async with block and open a new TemporalStatefulMcpToolSet block.

Temporary limitation: a retry after a lost host or a timed-out attempt is refused even when no session was opened, because the plugin cannot wait for a Worker heartbeat. This will be replaced once the Python SDK offers a heartbeat the plugin can wait for.

Also: a reset with reset_heartbeat=True clears the marker, so the next attempt builds the toolset. The session Activity sends one more heartbeat at its start, and before its first session it calls DescribeWorkflowExecution once or a few times, so the Worker's client needs permission to describe the Workflow. No public name, option or default changes.

Replay: Workflow code changes, without a version gate:

• Histories without a session that ended while the run used the toolset replay on both builds: completed runs, and runs waiting with their session open.
• A history in which the session Activity ended before a later list or call, or while one waited, does not replay across builds, in either direction. It fails with [TMPRL1100], or with an ActivityMachine error for a history recorded before this change in which an operation was waiting.
• Such runs are already failing with DedicatedWorkerFailure or a timeout. A Worker deploy that replays one of them mid-run will fail its Workflow task until it runs on the build that recorded it.
• Neither upgrade nor rollback protects the run that crosses it. The run whose session Activity is running when you upgrade recorded no marker, so its retry at that upgrade still rebuilds the toolset; the earlier code ignores the marker.
• The refusal of an earlier block's operations is worker-side only and changes no Workflow command, input or task queue.

Why?

See #310. Nothing told a retry of the session Activity apart from its first attempt, so a graceful Worker shutdown, the one-hour start-to-close timeout, or a heartbeat timeout with a lost host gave the run a new, empty MCP session, and later tool calls succeeded against it with no error.

Checklist

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

2. How was this tested:
   Supplemental local runs in Docker, with temporalio 1.34.0, google-adk 2.12.0 and mcp 1.30.0 from the plugin's lock and Temporal CLI 1.9.1:

   • On this change (on top of cb83258): make sync, make format, make test PYTEST_ARGS='-n 4' (197 passed, 5 skipped), make lint, make build and scripts/ci/check_conventions.py pass. After a typing fix in the new test, make lint and test_stateful_mcp.py (20 passed, 3 runs) were rerun.
   • test_stateful_operation_of_ended_session_not_run_by_next_block (new): the session's host shuts down gracefully after a session opened, a put('OLD-PAYLOAD') waits on the queue, the retried session is refused, and the Workflow catches the failure and opens a second block with the same name, whose get must return EMPTY.
     • With this change, both cases pass (4 runs): with WAIT_CANCELLATION_COMPLETED the second block's worker refuses the old put; with TRY_CANCEL the put is cancelled.
     • With cb83258's _mcp.py, the WAIT_CANCELLATION_COMPLETED case fails: get returns OLD-PAYLOAD. The TRY_CANCEL case passes.
   • test_stateful_session_after_host_killed[heartbeat_timeout_call_waiting] (new): a call-tool with WAIT_CANCELLATION_COMPLETED and a 2-minute start-to-close timeout is running a tool that never returns when the session host is killed. The retried session is refused, and the run must fail within 10 seconds.
     • With this change it passed in the suite run and 3 more runs.
     • With ee42451's _mcp.py, the run never failed (pytest timeout after 120 s): the cancelled call waited for an acknowledgment, and the Worker logged Eviction processed, but 2 tasks remain.
     • A first version of this fix that requested the cancel and did not wait for it also failed: the run raised the refusal, but eviction then hung on the cancelled operation. A minimal script on the same SDK shows the same: a requested WAIT_CANCELLATION_COMPLETED cancel that is never acknowledged blocks eviction, while an operation left uncancelled does not.
   • Workflow cancellation during a running call-tool (scratch test, default cancellation type): the run ends WORKFLOW_EXECUTION_CANCELED with the same history events on ee42451 and on this change.
   • test_stateful_operations_fail_after_session_refused (start-to-close-only config, the session on its own host queue that shuts down gracefully after a session opened):
     • With a list-tools waiting, the refused retry cancels it and the run fails with the refusal.
     • With get_tools called after the refusal, nothing is scheduled and the run fails at once.
     • With this change, both pass. With 8b32464's _mcp.py, both timed out.
   • test_stateful_session_after_worker_replaced[first_connection_fails_delayed_heartbeats]:
     • Setup: no heartbeat timeout, a real McpToolset whose first connection fails, "may be open" heartbeats that take 1.5 s to encode, and failure attributes encoded through a 3 s codec. The Worker shuts down gracefully while a heartbeat is being encoded.
     • With the wait before the clear, the retry builds the toolset and the run completes (5 of 5, earlier round).
     • Without the wait before the clear, the retry is refused (3 of 3); with 8b32464's _mcp.py it is refused too.
   • test_stateful_session_after_host_killed[no_heartbeat_timeout] kills the session host, with no heartbeat timeout, after a session opened behind a throttled heartbeat. On 3b9ea40 the retried attempt built a new toolset; on 8b32464 and with this change it is refused without a factory call.
   • On main at 88745dc, with only _HEARTBEAT_INTERVAL and _SESSION_OPENED added so the tests can run, using 679d134's tests:
     • The activity-reset test fails with DID NOT RAISE WorkflowFailureError.
     • The host-kill test fails with assert 1 == 0 (the new host called the factory).
     • The delayed-marker test sees no marker.
     • The worker-replacement cases with a session opened time out waiting for the retried session Activity to fail.
   • [Bug] google_adk: a retried stateful MCP session Activity silently opens a new, empty MCP session #310's script (run on 8b32464) fails as reported on 0.0.1 (google-adk 2.13.0) and on main. On 8b32464 its control reads its value, and the retried run fails with DedicatedWorkerFailure because the server held mcp-session-may-be-open; get never reached the MCP server.
   • Replay against ee42451: the host-kill and refused-operation histories (default cancellation type) recorded on either build replay on both. A history in which a WAIT_CANCELLATION_COMPLETED call was running when the session ended, recorded on this change, fails on ee42451 with [TMPRL1100], as described under Replay above. The later worker-side refusal of an earlier block's operations changed no Workflow code, so it was not replayed again.
   • Replay on 8b32464 (earlier round): completed and waiting runs recorded on 8b32464 and on this change replay on both builds. Refused runs in which a later or waiting list-tools met the ended session fail across builds, as described under Replay above, and replay on the build that recorded them. Live carries through a Worker replacement were not rerun for this change.
   • CI's Python matrix, macOS and Windows were not run locally.

3. Any docs updates needed?
   The google_adk README now covers the stateful toolsets' retry behaviour, how operations fail after the session ends (including WAIT_CANCELLATION_COMPLETED, and that a later block with the same name does not run such an operation), and the heartbeat limitation. Release notes should name the caller change above.

Not covered:

• google-adk 2.13.0 with this version, and the streamable HTTP and SSE transports in the tests (they use stdio).
• A host lost while the first heartbeat is in flight, before the read-back confirms it, opened no session yet, so its retry builds a new toolset without losing state.
• The refusal after a lost host with no session opened, and the read-back, can be dropped once the Python SDK offers a heartbeat the plugin can wait for: the marker could then be stored before ADK uses the session.
• A Workflow cancelled while a WAIT_CANCELLATION_COMPLETED operation runs on a lost host still waits for that operation's timeout, as on main.
• A WAIT_CANCELLATION_COMPLETED operation that had started on a lost host and is retried after its timeout keeps its activity ID, so a later block refuses it by the same check. This was reasoned from the code, not run.
• Three internals the plugin relies on:
  • The marker relies on McpToolset._mcp_session_manager and MCPSessionManager.create_session(). These are private to ADK but the same from google-adk 2.4.0 to 2.13.0. If ADK renames them, the toolset falls back to recording the marker before its first operation, which also refuses a retry after a first connection that failed.
  • Waiting for heartbeats still being sent relies on the SDK starting an asyncio task per activity.heartbeat() call, as temporalio 1.34.0 does. If it stops doing so, there is nothing to wait for.
  • Refusing an earlier block's operations relies on default activity IDs, which number a run's activities in scheduling order. If config or server_session_config sets a non-numeric activity_id, the IDs are not compared and such an operation can run against a later block, as before.
• The OpenAI Agents plugin's stateful MCP server (python/openai_agents/src/temporalio/openai_agents/_mcp.py) has the same session Activity code and is not changed here.

Open questions: none.

Draft while the caller change is reviewed.