From 054f4bb318c974fb73f2ab3064eff8157f90543b Mon Sep 17 00:00:00 2001 From: ShotaroKataoka Date: Tue, 3 Nov 2026 18:21:32 +0900 Subject: [PATCH 1/3] fix: stop a stopped Cloud turn's in-flight MCP calls from writing to the deck A tool call the agent had already sent kept running on the MCP runtime after a stop and wrote to S3 (dev: run_python saved 4 slides 1.7 s after a cancel; generate_pptx's background previews landed 1.5 s after StopRuntimeSession). - agent/turn_fence.py: each turn mints a generation (time_ns at turn start) and sends it as MCP _meta["sdpm/generation"] on every Presentation Maker client (orchestrator, composer groups; Strands MCPClient meta, so direct compose calls carry it too). It records the decks its tool calls name and, when the turn is stopped (stop event set in the entrypoint's finally, or SIGTERM via stop_guard.begin_turn(on_stop=...)), calls fence_turn once. - servers/remote: fence_turn raises writeFence on the caller's own deck item (conditional UpdateItem, forward only). AwsStorage's S3 client checks the request's fence before every put/delete/copy (one choke point, including the direct _s3.delete_object calls); the fence is bound in offloaded_tool and carried into background threads and upload pools. update_deck is conditional on the fence. Stale requests write nothing and get a "Cancelled" error. Requests without a generation are unchanged. - No IAM change: the advance goes through the MCP runtime role; the agent role keeps no DynamoDB access. Checks: - make lint: All checks passed - make test: 2129 passed, 18 skipped - ash scan --mode local --fail-on-findings: exit 0 - infra unchanged (no npm test / tsc needed) --- CHANGELOG.md | 12 + agent/basic_agent.py | 14 +- agent/factory.py | 19 +- agent/mcp_clients.py | 6 + agent/stop_guard.py | 21 +- agent/turn_fence.py | 115 ++++++++ docs/en/architecture.md | 1 + servers/remote/README.md | 1 + servers/remote/server.py | 68 ++++- servers/remote/server_utils.py | 4 +- servers/remote/storage/aws.py | 46 +++- servers/remote/storage/write_fence.py | 207 ++++++++++++++ servers/remote/tools/html_deck.py | 4 +- tests/agent_stop_fakes.py | 16 +- tests/agent_stop_server.py | 6 +- tests/test_agent_stop_guard.py | 14 +- tests/test_agent_turn_fence.py | 256 ++++++++++++++++++ tests/test_remote_html.py | 4 +- tests/test_remote_write_fence.py | 370 ++++++++++++++++++++++++++ 19 files changed, 1151 insertions(+), 33 deletions(-) create mode 100644 agent/turn_fence.py create mode 100644 servers/remote/storage/write_fence.py create mode 100644 tests/test_agent_turn_fence.py create mode 100644 tests/test_remote_write_fence.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 313483bc..e9c90e9c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -239,6 +239,18 @@ Entries before v0.5.0 were written retroactively as summaries. ### Fixed +- **Stopping a Cloud turn no longer lets its in-flight tool calls change the deck + afterwards.** A tool call the agent had already sent kept running on the MCP runtime and + wrote to S3 after the stop (on dev: `run_python` saved 4 slides 1.7 s after a cancel, and + `generate_pptx`'s previews and thumbnail landed 1.5 s after the stop button). Each turn + now sends a generation number with its Presentation Maker calls (composers included), and + a stopped turn raises the deck's write fence (`writeFence` on the deck item) to it through + the new `fence_turn` tool. The MCP server checks the fence before every S3 write and skips + the writes of a stopped turn (the call returns a "Cancelled" error), and the deck record + update after a build is conditional on it. A turn that ends normally fences nothing, so + its last previews still land when the next message is sent quickly. Calls without a + generation (the local server, other MCP clients) are unchanged. + - **Stopping a Cloud turn no longer leaves it running in the background.** The stop button ends the runtime session, and AgentCore then sends the agent's process SIGTERM. The agent used to wait for the running turn to finish, so the turn and its composers kept calling diff --git a/agent/basic_agent.py b/agent/basic_agent.py index 7766f7d7..fa3c755f 100644 --- a/agent/basic_agent.py +++ b/agent/basic_agent.py @@ -23,6 +23,7 @@ from errors import BAD_INPUT, AUTH, error_event from factory import create_agent from streaming import stream_agent +from turn_fence import TurnFence logger = logging.getLogger("sdpm.agent") @@ -87,8 +88,11 @@ async def agent_stream(payload, context): cancel = asyncio.Event() _cancel_events[session_id] = cancel - # The turn's stop event: set on SIGTERM, observed by Strands and the composers. - stop = stop_guard.begin_turn() + # The turn's write fence: its generation goes with every Presentation Maker call, and + # a stopped turn fences its decks so its in-flight calls no longer write (turn_fence.py). + fence = TurnFence() + # The turn's stop event: set on SIGTERM (which also fences), observed by Strands and the composers. + stop = stop_guard.begin_turn(on_stop=fence.advance) try: os.environ["_CURRENT_SESSION_ID"] = session_id @@ -102,6 +106,7 @@ async def agent_stream(payload, context): jwt_token=jwt_token, chat_model_id=requested_chat_model_id, create_model_id=requested_create_model_id, + fence=fence, ) yield {"mcp_status": mcp_status} @@ -114,6 +119,11 @@ async def agent_stream(payload, context): yield error_event(e) finally: stop_guard.end_turn(stop) + # Stopped (cancel, a newer request, the client gone, SIGTERM): fence the turn's decks. + if stop.is_set(): + fence.advance() + else: + fence.release() if _cancel_events.get(session_id) is cancel: del _cancel_events[session_id] diff --git a/agent/factory.py b/agent/factory.py index c0ca1ead..c4c0ab52 100644 --- a/agent/factory.py +++ b/agent/factory.py @@ -33,6 +33,7 @@ from modes.composer import make_compose_slides from resilience import LoopGuard, gate_tool_call from session import fix_excess_tool_results +from turn_fence import TurnFence from tools.hearing_tool import hearing from tools.web_tools import web_fetch @@ -64,7 +65,7 @@ def _resolve_model_id(requested: str | None, default: str) -> str: _MCP_FACTORIES = [ - lambda jwt_token, session_id="", tool_filters=None: mcp_agentcore_runtime(jwt_token=jwt_token, session_id=session_id, tool_filters=tool_filters), + lambda jwt_token, session_id="", tool_filters=None, **kwargs: mcp_agentcore_runtime(jwt_token=jwt_token, session_id=session_id, tool_filters=tool_filters, **kwargs), lambda jwt_token, session_id="", tool_filters=None: mcp_aws_knowledge(), lambda jwt_token, session_id="", tool_filters=None: mcp_aws_pricing(), ] @@ -221,9 +222,13 @@ def _session_manager(memory_id: str, session_id: str, user_id: str, region: str) # Unified factory # --------------------------------------------------------------------------- -def create_agent(mode: str, user_id: str, session_id: str, jwt_token: str, chat_model_id: str | None = None, create_model_id: str | None = None) -> tuple[Agent, list[dict]]: +def create_agent(mode: str, user_id: str, session_id: str, jwt_token: str, chat_model_id: str | None = None, create_model_id: str | None = None, fence: TurnFence | None = None) -> tuple[Agent, list[dict]]: """Create a Strands Agent for the given mode. + ``fence`` (turn_fence.py): its generation goes into the ``_meta`` of every + Presentation Maker call (orchestrator and composers), and it learns the decks the + turn's tool calls name. + Returns: Tuple of (Configured Strands Agent, MCP status list). """ @@ -251,6 +256,7 @@ def create_agent(mode: str, user_id: str, session_id: str, jwt_token: str, chat_ # MCP servers — Presentation Maker built per turn (it carries the JWT), AWS # Knowledge / AWS Pricing reused from the process; all connected concurrently below. + fence_kwargs = {"meta": fence.meta} if fence is not None else {} built = [] names = [] required_flags = [] @@ -260,7 +266,7 @@ def create_agent(mode: str, user_id: str, session_id: str, jwt_token: str, chat_ else: # Apply tool_filters only to the Presentation Maker server (index 0) filters = {"allowed": cfg.allowed_tools} if (i == 0 and cfg.allowed_tools) else None - built.append(factory_fn(jwt_token, session_id=session_id, tool_filters=filters)) + built.append(factory_fn(jwt_token, session_id=session_id, tool_filters=filters, **fence_kwargs)) names.append(name) required_flags.append(required) @@ -311,7 +317,7 @@ def critical_path(required_clients: list) -> None: # Composers get their own Presentation Maker client per group; it carries the # composer allowlist so a composer sees only the deck tools. composer_filters = {"allowed": MODES["composer"].allowed_tools} - composer_mcp_factory = lambda mcp_session_id="": mcp_agentcore_runtime(jwt_token=jwt_token, session_id=mcp_session_id, tool_filters=composer_filters) # noqa: E731 + composer_mcp_factory = lambda mcp_session_id="": mcp_agentcore_runtime(jwt_token=jwt_token, session_id=mcp_session_id, tool_filters=composer_filters, **fence_kwargs) # noqa: E731 compose_slides = make_compose_slides(mcp_servers, composer_model, composer_mcp_factory, extra_tools=[web_fetch], model_id=resolved_create, user_id=user_id, session_id=session_id) tools.append(compose_slides) @@ -381,6 +387,11 @@ def critical_path(required_clients: list) -> None: # Write gate: no tool runs once the process is stopping (SIGTERM, stop_guard.py) agent.hooks.add_callback(BeforeToolCallEvent, gate_tool_call) + # Write fence: the decks this turn touches, and the client that fences them on a stop + if fence is not None: + agent.hooks.add_callback(BeforeToolCallEvent, fence.note_tool_call) + fence.bind(mcp_servers[0] if mcp_servers else None) + # Cost logger agent.hooks.add_callback(AfterInvocationEvent, log_usage) diff --git a/agent/mcp_clients.py b/agent/mcp_clients.py index 7f3a11cd..4ffd9dbd 100644 --- a/agent/mcp_clients.py +++ b/agent/mcp_clients.py @@ -27,6 +27,7 @@ def mcp_agentcore_runtime( jwt_token: str, session_id: str = "", tool_filters: dict | None = None, + meta: dict | None = None, ) -> MCPClient: """Pattern 1: Amazon Bedrock AgentCore Runtime MCP Server with JWT Bearer authentication. @@ -36,6 +37,9 @@ def mcp_agentcore_runtime( routes by that header, so omitting it makes every request land on a new microVM — see the note below. Empty string disables the header. tool_filters: Optional tool filter dict (e.g. {"allowed": ["tool1", "tool2"]}). + meta: Optional ``_meta`` sent with every tool call (the turn's write-fence + generation, turn_fence.py). It travels in the JSON-RPC body, which the + AgentCore proxy forwards as is; Strands merges trace context into it. """ region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-east-1")) runtime_arn = os.environ["MCP_RUNTIME_ARN"] @@ -62,6 +66,8 @@ def mcp_agentcore_runtime( kwargs: dict = {} if tool_filters: kwargs["tool_filters"] = tool_filters + if meta: + kwargs["meta"] = meta return MCPClient( lambda: streamablehttp_client( diff --git a/agent/stop_guard.py b/agent/stop_guard.py index 21e8d584..cd0cfae1 100644 --- a/agent/stop_guard.py +++ b/agent/stop_guard.py @@ -26,6 +26,7 @@ import sys import threading import time +from typing import Callable from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager @@ -39,18 +40,21 @@ # Re-entrant: the SIGTERM handler takes it on the main thread, which may hold it already. _lock = threading.RLock() -_running: set[threading.Event] = set() +# Running turns: stop event -> what to call when the turn is stopped (or None). +_running: dict[threading.Event, Callable[[], object] | None] = {} _exiting = threading.Event() -def begin_turn() -> threading.Event: +def begin_turn(on_stop: Callable[[], object] | None = None) -> threading.Event: """Register a running turn and return its stop event (set on SIGTERM). + ``on_stop`` runs when SIGTERM stops the turn (on the signal handler's thread, so it + must not block: turn_fence.TurnFence.advance hands its call to a thread). Pair with end_turn() in a finally block. """ stop = threading.Event() with _lock: - _running.add(stop) + _running[stop] = on_stop if writes_closed(): stop.set() # the process is already stopping return stop @@ -58,16 +62,21 @@ def begin_turn() -> threading.Event: def end_turn(stop: threading.Event) -> None: with _lock: - _running.discard(stop) + _running.pop(stop, None) def stop_all() -> int: """Close the write gate and set the stop event of every running turn.""" close_writes() with _lock: - running = list(_running) - for stop in running: + running = list(_running.items()) + for stop, on_stop in running: stop.set() + if on_stop is not None: + try: + on_stop() + except Exception as e: # noqa: BLE001 - the other turns still stop + logger.warning("SIGTERM guard: a turn's stop callback failed: %s", e) return len(running) diff --git a/agent/turn_fence.py b/agent/turn_fence.py new file mode 100644 index 00000000..373997dc --- /dev/null +++ b/agent/turn_fence.py @@ -0,0 +1,115 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""The agent's half of the MCP write fence (servers/remote/storage/write_fence.py). + +Stopping a turn does not stop the Presentation Maker calls it has already sent: they +keep running on the MCP runtime and write to S3 afterwards. So each turn: + - mints a generation (wall-clock nanoseconds at the start of the turn: a later turn, + in this process or in a new microVM, always gets a larger one) and sends it in the + ``_meta`` of every Presentation Maker call, the orchestrator's, the composers' and + compose_slides' direct ones (``meta``, set on each MCP client); + - records the decks its tool calls name (``note_tool_call``, a BeforeToolCallEvent hook); + - when it is stopped (``advance``: cancel, a newer request, the client going away, + SIGTERM), calls ``fence_turn`` with those decks. The server then skips every write + the stopped turn's calls still make to them. + +A turn that ends normally fences nothing, so the background work of its last calls +(previews after generate_pptx) still lands. +""" + +import logging +import threading +import time +import uuid + +logger = logging.getLogger("sdpm.agent") + +GENERATION_META_KEY = "sdpm/generation" # same key as the server's write_fence +FENCE_TOOL = "fence_turn" + +# Consumer id held on the turn's Presentation Maker client until the fence call is done, +# so the Agent's cleanup does not stop the client under it. +_CONSUMER = "sdpm-turn-fence" + + +class TurnFence: + """One turn's generation, the decks it touched, and the call that fences them.""" + + def __init__(self, generation: int | None = None): + self.generation = generation if generation is not None else time.time_ns() + # Re-entrant: advance() also runs from the SIGTERM handler on the main thread, + # which may be inside note_tool_call (hooks run on the event loop's thread). + self._lock = threading.RLock() + self._decks: set[str] = set() + self._client = None + self._done = False + + @property + def meta(self) -> dict: + """The ``_meta`` entry for this turn's MCP calls.""" + return {GENERATION_META_KEY: self.generation} + + @property + def decks(self) -> list[str]: + with self._lock: + return sorted(self._decks) + + def bind(self, client) -> None: + """Use ``client`` (the turn's Presentation Maker MCPClient) for the fence call.""" + with self._lock: + if self._done or client is None: + return + client.add_consumer(_CONSUMER) + self._client = client + + def note_tool_call(self, event) -> None: + """BeforeToolCallEvent hook: remember the deck a tool call names.""" + tool_input = event.tool_use.get("input") + deck_id = tool_input.get("deck_id") if isinstance(tool_input, dict) else None + if isinstance(deck_id, str) and deck_id.strip(): + with self._lock: + self._decks.add(deck_id.strip()) + + def advance(self) -> threading.Thread | None: + """The turn was stopped: fence its decks (once). The call runs on a daemon thread, + so this is safe from a signal handler and from the event loop; returns it.""" + return self._finish(fence=True) + + def release(self) -> None: + """The turn ended without a stop: nothing to fence, let the client go.""" + self._finish(fence=False) + + def _finish(self, fence: bool) -> threading.Thread | None: + with self._lock: + if self._done: + return None + self._done = True + client, decks = self._client, sorted(self._decks) + self._client = None + if client is None: + return None + thread = threading.Thread( + target=self._send, args=(client, decks if fence else []), daemon=True, name="sdpm-turn-fence", + ) + thread.start() + return thread + + def _send(self, client, decks: list[str]) -> None: + try: + if decks: + # Not through call_tool_with_retry: this is the one call made after the + # write gate closed. The client's meta carries this turn's generation. + result = client.call_tool_sync( + tool_use_id=f"fence-{uuid.uuid4().hex[:8]}", name=FENCE_TOOL, arguments={"deck_ids": decks}, + ) + if result.get("status") == "error": + logger.warning("write fence: fence_turn failed for %s: %s", decks, result.get("content")) + else: + logger.info("write fence: fenced generation %s for %s", self.generation, decks) + except Exception as e: # noqa: BLE001 - best effort; the turn is over either way + logger.warning("write fence: fence_turn for %s failed: %s", decks, e) + finally: + try: + client.remove_consumer(_CONSUMER) + except Exception as e: # noqa: BLE001 - best-effort cleanup + logger.warning("write fence: releasing the MCP client failed: %s", e) diff --git a/docs/en/architecture.md b/docs/en/architecture.md index 8ef021fe..82cd37fe 100644 --- a/docs/en/architecture.md +++ b/docs/en/architecture.md @@ -394,6 +394,7 @@ To add custom roles (e.g., team-based access), modify the `resolve_role` functio | `run_python` | Runs in an Amazon Bedrock AgentCore Code Interpreter sandbox | | `search_slides` | Semantic slide search; registered only when an Amazon Bedrock Knowledge Base is configured | | `start_translation` | Not bound: the translate workflow needs a checkout | +| `fence_turn` | Agent plumbing (remote only, in no mode's allowlist): the L4 agent calls it when a turn is stopped, and the stopped turn's in-flight calls then write nothing to those decks. Every call carries the turn's generation in `_meta["sdpm/generation"]`; see `servers/remote/storage/write_fence.py` | ### Agent-Level Tools (Layer 4) diff --git a/servers/remote/README.md b/servers/remote/README.md index 8d4d4df2..472f868c 100644 --- a/servers/remote/README.md +++ b/servers/remote/README.md @@ -45,6 +45,7 @@ Remote MCP server for Amazon Bedrock AgentCore Runtime. Provides all spec-driven | `read_attachment` | Read content from an attached file with byte-offset paging | | `import_attachment` | Import file into deck workspace | | `search_slides` | Semantic slide search (optional, requires Amazon Bedrock KB) | +| `fence_turn` | Agent plumbing: skip the writes a stopped agent turn's in-flight calls still make (`storage/write_fence.py`) | ## Deployment diff --git a/servers/remote/server.py b/servers/remote/server.py index c625aa45..adb3c5c6 100644 --- a/servers/remote/server.py +++ b/servers/remote/server.py @@ -33,6 +33,7 @@ from shared.authz import authorize # noqa: E402 from shared import user_library # noqa: E402 +from storage import write_fence # noqa: E402 from storage.aws import AwsStorage # noqa: E402 from tools import assets, reference, preview, generate # noqa: E402 from tools import sandbox as sandbox_mod # noqa: E402 @@ -66,8 +67,8 @@ def _run_in_background(target, *, name: str) -> None: - """Start ``target`` on a daemon thread. Tests patch this to run inline.""" - threading.Thread(target=target, name=name, daemon=True).start() + """Start ``target`` on a daemon thread, under the request's write fence. Tests patch this to run inline.""" + threading.Thread(target=write_fence.carry(target), name=name, daemon=True).start() # Background work that get_preview may need to wait for. Keyed by deck so a @@ -122,13 +123,21 @@ def offloaded_tool(fn): The original function is returned unchanged so it stays callable (and testable) as a plain function. + + A call that carries the agent's turn generation and names a deck runs under + that deck's write fence (storage/write_fence.py). """ import asyncio import functools import inspect async def _runner(**kwargs): - return await asyncio.to_thread(functools.partial(fn, **kwargs)) + token = _bind_write_fence(kwargs.get("deck_id")) + try: + return await asyncio.to_thread(functools.partial(fn, **kwargs)) + finally: + if token is not None: + write_fence.reset(token) _runner.__name__ = fn.__name__ _runner.__qualname__ = fn.__qualname__ @@ -138,6 +147,30 @@ async def _runner(**kwargs): mcp.tool()(_runner) return fn + +def _request_meta(): + """``_meta`` of the MCP request being handled, or None outside a request.""" + from mcp.server.lowlevel.server import request_ctx + + ctx = request_ctx.get(None) + return ctx.meta if ctx is not None else None + + +def _bind_write_fence(deck_id): + """Bind the write fence for a call that carries a turn generation and names a deck. + + Returns the ContextVar token, or None when the call is not fenced (no generation: + other MCP clients and older agents keep today's behaviour). + """ + generation = write_fence.generation_from_meta(_request_meta()) + if generation is None or not isinstance(deck_id, str) or not deck_id.strip(): + return None + try: + user_id = _get_user_id() + except ValueError: + return None # the tool rejects the call itself + return write_fence.bind(write_fence.Fence(generation=generation, user_id=user_id, deck_id=deck_id)) + # --- HTTP Request ContextVar (for extracting user_id from Runtime header) --- _current_request_headers: ContextVar[dict] = ContextVar("_current_request_headers", default={}) @@ -686,6 +719,33 @@ def import_attachment( ) +# --- Write fence (agent plumbing) --- + + +@offloaded_tool +def fence_turn( + deck_ids: Annotated[list[str], Field(description="Decks the stopped turn worked on.", max_length=50)], +) -> str: + """Agent plumbing, not a slide tool: the agent calls it when one of its turns is + stopped. Writes that the stopped turn's calls still make to these decks are skipped + from now on. The turn is the generation in this request's _meta. + """ + generation = write_fence.generation_from_meta(_request_meta()) + if generation is None: + return json.dumps({"error": f"fence_turn needs the turn generation in _meta[{write_fence.GENERATION_META_KEY!r}]"}) + user_id = _get_user_id() + fenced: list[str] = [] + unchanged: list[str] = [] + for deck_id in dict.fromkeys(d.strip() for d in deck_ids if isinstance(d, str) and 0 < len(d.strip()) <= 64): + # Keyed by the caller's own sub: only the owner's deck item can be fenced. + if write_fence.advance(_storage.table, user_id, deck_id, generation): + fenced.append(deck_id) + else: + unchanged.append(deck_id) + logger.info("fence_turn: generation %s fenced=%s unchanged=%s", generation, fenced, unchanged) + return json.dumps({"generation": generation, "fenced": fenced, "unchanged": unchanged}) + + # --- Generation Tools --- @@ -1236,7 +1296,7 @@ def upload(item: tuple[str, dict]) -> None: _delete_quietly(k) with ThreadPoolExecutor(max_workers=8) as pool: - list(pool.map(upload, sorted(composed.slides.items()))) + list(pool.map(write_fence.carry(upload), sorted(composed.slides.items()))) except Exception: logger.error("compose failed", exc_info=True) finally: diff --git a/servers/remote/server_utils.py b/servers/remote/server_utils.py index abcd633a..3505b537 100644 --- a/servers/remote/server_utils.py +++ b/servers/remote/server_utils.py @@ -9,7 +9,7 @@ import tempfile from pathlib import Path -from storage import Storage +from storage import Storage, write_fence logger = logging.getLogger("sdpm.mcp") @@ -82,7 +82,7 @@ def schedule_webp_background( """ try: threading.Thread( - target=_generate_webp_background, + target=write_fence.carry(_generate_webp_background), args=(deck_id, pptx_path, tmpdir, storage, slugs), kwargs={"user_id": user_id}, name=f"webp-{deck_id}", diff --git a/servers/remote/storage/aws.py b/servers/remote/storage/aws.py index 9530f83d..204ebbf5 100644 --- a/servers/remote/storage/aws.py +++ b/servers/remote/storage/aws.py @@ -16,7 +16,30 @@ from shared import user_library from shared.schema import user_template_key, user_template_sk -from storage import Storage +from storage import Storage, write_fence + + +class _FencedS3: + """The S3 client behind AwsStorage: every write first passes the request's write + fence (storage/write_fence.py). Without a fence (no generation in the request, or no + deck) it is the plain client. Reads and everything else pass straight through.""" + + _WRITES = frozenset({"put_object", "delete_object", "delete_objects", "copy_object"}) + + def __init__(self, client: Any, table: Callable[[], Any]) -> None: + self._client = client + self._table = table + + def __getattr__(self, name: str) -> Any: + attr = getattr(self._client, name) + if name not in self._WRITES: + return attr + + def fenced(*args: Any, **kwargs: Any) -> Any: + write_fence.check(self._table()) + return attr(*args, **kwargs) + + return fenced class AwsStorage(Storage): @@ -44,7 +67,7 @@ def __init__( if s3_client is None and s3_factory is None: raise ValueError("s3_client or s3_factory is required") self._table_obj = table - self._s3_obj = s3_client + self._s3_obj = _FencedS3(s3_client, lambda: self._table) if s3_client is not None else None self._table_factory = table_factory self._s3_factory = s3_factory self._lock = threading.Lock() @@ -71,7 +94,7 @@ def _s3(self) -> Any: if self._s3_obj is None: with self._lock: if self._s3_obj is None: - self._s3_obj = self._s3_factory() # type: ignore[misc] + self._s3_obj = _FencedS3(self._s3_factory(), lambda: self._table) # type: ignore[misc] return self._s3_obj @property @@ -111,6 +134,10 @@ def update_deck(self, deck_id: str, user_id: str, updates: dict) -> dict: Previous values of the updated attributes (UPDATED_OLD) — lets callers clean up superseded S3 objects (e.g. the old pptxS3Key) without an extra read. + + Raises: + write_fence.StaleTurn: The request's turn was stopped (the update is + conditional on the write fence, so it cannot race the fence). """ set_parts = [] remove_parts = [] @@ -132,10 +159,21 @@ def update_deck(self, deck_id: str, user_id: str, updates: dict) -> dict: "Key": {"PK": f"USER#{user_id}", "SK": f"DECK#{deck_id}"}, "UpdateExpression": " ".join(expression_parts), } + fenced = write_fence.update_condition() + if fenced is not None: + fence, condition = fenced + kwargs["ConditionExpression"] = condition["ConditionExpression"] + kwargs["ExpressionAttributeNames"] = condition["ExpressionAttributeNames"] + values.update(condition["ExpressionAttributeValues"]) if values: kwargs["ExpressionAttributeValues"] = values kwargs["ReturnValues"] = "UPDATED_OLD" - response = self._table.update_item(**kwargs) + try: + response = self._table.update_item(**kwargs) + except Exception as e: + if fenced is not None and write_fence.is_conditional_check_failure(e): + raise write_fence.condition_failed(fenced[0]) from e + raise return response.get("Attributes", {}) # --- Presentation JSON (S3) --- diff --git a/servers/remote/storage/write_fence.py b/servers/remote/storage/write_fence.py new file mode 100644 index 00000000..330571c4 --- /dev/null +++ b/servers/remote/storage/write_fence.py @@ -0,0 +1,207 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""Write fence: a stopped agent turn's in-flight tool calls no longer change the deck. + +Stopping a Cloud turn does not stop the MCP calls it has already sent: the server is +stateless (``notifications/cancelled`` reaches another request) and tool bodies run on +threads that cannot be interrupted. Measured on dev: run_python wrote 4 slides 1.7 s +after a cancel, and a generate_pptx call's background previews were uploaded 1.5 s after +a StopRuntimeSession. + +How it works: + - The agent mints a generation per turn (an integer that grows over time) and sends + it in every tool call's ``_meta`` (GENERATION_META_KEY). + - When the turn is stopped, the agent calls ``fence_turn(deck_ids)`` with that + generation. ``advance`` raises the deck item's FENCE_ATTR to it with a conditional + UpdateItem (never lowers it). + - Every S3 write AwsStorage makes for a request that carries a generation and names a + deck checks the deck item first (``check``). A generation at or below the fence is a + stopped turn: the write is skipped and StaleTurn is raised. The deck record update + that accompanies a build (``update_deck``) carries the same test as its condition, so + a write racing the fence cannot repoint the deck at a stale artifact. + +Requests without a generation (the local server, other MCP clients, older agents) are +not fenced. A newer turn is never fenced by an older turn's stop, and a finished turn's +background work (previews after generate_pptx) is not dropped by the next turn's start. + +The fence for a request lives in a ContextVar. ``asyncio.to_thread`` copies it into the +tool's thread; threads the tool starts itself get it through ``carry``. +""" + +from __future__ import annotations + +import contextvars +import logging +import time +from dataclasses import dataclass +from typing import Any, Callable + +from shared.schema import deck_pk, deck_sk + +logger = logging.getLogger("sdpm.mcp") + +# _meta key of the caller's turn generation (MCP _meta keys may carry a prefix/name). +GENERATION_META_KEY = "sdpm/generation" +# Deck item attribute: the newest stopped generation. Writes at or below it are skipped. +FENCE_ATTR = "writeFence" + +# A generation from further in the future than this is rejected: it would keep the +# deck fenced for the owner's own later turns until the clock caught up. +_FUTURE_SLACK_NS = 300 * 10**9 + + +class StaleTurn(RuntimeError): + """The request belongs to a stopped turn; its write was skipped.""" + + +@dataclass +class Fence: + """The fence of one request: its generation and the deck item to check.""" + + generation: int + user_id: str + deck_id: str + stale: bool = False + + def error(self) -> StaleTurn: + return StaleTurn( + f"Cancelled: the agent turn that made this call was stopped, so its writes to " + f"deck {self.deck_id} were skipped." + ) + + +_current: contextvars.ContextVar[Fence | None] = contextvars.ContextVar("sdpm_write_fence", default=None) + + +def valid_generation(value: Any) -> int | None: + """The generation as an int, or None when it is missing, not an integer or from the future.""" + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return None + if value > time.time_ns() + _FUTURE_SLACK_NS: + return None + return value + + +def generation_from_meta(meta: Any) -> int | None: + """Read the generation from a request's ``_meta`` (a pydantic model with extra fields, or a dict).""" + if meta is None: + return None + extra = meta if isinstance(meta, dict) else (getattr(meta, "model_extra", None) or {}) + if GENERATION_META_KEY not in extra: + return None + generation = valid_generation(extra[GENERATION_META_KEY]) + if generation is None: + logger.warning("write fence: ignoring an invalid generation %r", extra[GENERATION_META_KEY]) + return generation + + +def bind(fence: Fence) -> contextvars.Token: + """Make ``fence`` the current request's fence. Pair with ``reset``.""" + return _current.set(fence) + + +def reset(token: contextvars.Token) -> None: + _current.reset(token) + + +def current() -> Fence | None: + return _current.get() + + +def carry(fn: Callable) -> Callable: + """Wrap ``fn`` so it runs under the caller's fence on another thread (Thread targets, + thread-pool tasks). Only the fence is carried, nothing else from the context.""" + fence = _current.get() + if fence is None: + return fn + + def run(*args, **kwargs): + token = _current.set(fence) + try: + return fn(*args, **kwargs) + finally: + _current.reset(token) + + return run + + +def _key(fence: Fence) -> dict: + return {"PK": deck_pk(fence.user_id), "SK": deck_sk(fence.deck_id)} + + +def _skip(fence: Fence, floor: Any) -> StaleTurn: + if not fence.stale: + fence.stale = True + logger.info( + "write fence: deck %s: generation %s is at or below the fence %s; skipping this request's writes", + fence.deck_id, fence.generation, floor, + ) + return fence.error() + + +def check(table: Any) -> None: + """Raise StaleTurn when the current request's turn has been fenced; no-op without a fence. + + One consistent GetItem per write. Once a request is found stale, its later writes are + skipped without another read. + """ + fence = _current.get() + if fence is None: + return + if fence.stale: + raise fence.error() + item = table.get_item( + Key=_key(fence), + ProjectionExpression="#wf", + ExpressionAttributeNames={"#wf": FENCE_ATTR}, + ConsistentRead=True, + ).get("Item") + floor = (item or {}).get(FENCE_ATTR) + if floor is not None and fence.generation <= int(floor): + raise _skip(fence, floor) + + +def update_condition() -> tuple[Fence, dict] | None: + """For a fenced request: the fence and the UpdateItem arguments that make a deck + record update conditional on it. None without a fence.""" + fence = _current.get() + if fence is None: + return None + if fence.stale: + raise fence.error() + return fence, { + "ConditionExpression": "attribute_not_exists(#wf) OR #wf < :wf_gen", + "ExpressionAttributeNames": {"#wf": FENCE_ATTR}, + "ExpressionAttributeValues": {":wf_gen": fence.generation}, + } + + +def condition_failed(fence: Fence) -> StaleTurn: + """The StaleTurn for a conditional deck update that the fence rejected.""" + return _skip(fence, "(conditional update)") + + +def is_conditional_check_failure(error: Exception) -> bool: + response = getattr(error, "response", None) or {} + return response.get("Error", {}).get("Code") == "ConditionalCheckFailedException" + + +def advance(table: Any, user_id: str, deck_id: str, generation: int) -> bool: + """Raise the fence of the caller's own deck to ``generation``. + + Conditional: the item must exist under the caller's key (only an owner fences a + deck) and its fence must be lower. Returns False when nothing changed. + """ + try: + table.update_item( + Key={"PK": deck_pk(user_id), "SK": deck_sk(deck_id)}, + UpdateExpression="SET #wf = :wf_gen", + ConditionExpression="attribute_exists(PK) AND (attribute_not_exists(#wf) OR #wf < :wf_gen)", + ExpressionAttributeNames={"#wf": FENCE_ATTR}, + ExpressionAttributeValues={":wf_gen": generation}, + ) + return True + except Exception as e: + if is_conditional_check_failure(e): + return False + raise diff --git a/servers/remote/tools/html_deck.py b/servers/remote/tools/html_deck.py index a04738c5..941e1882 100644 --- a/servers/remote/tools/html_deck.py +++ b/servers/remote/tools/html_deck.py @@ -42,7 +42,7 @@ preview_key, preview_prefix, ) -from storage import Storage +from storage import Storage, write_fence from tools import template_html_cache logger = logging.getLogger("sdpm.html_deck") @@ -164,7 +164,7 @@ def _one(slug: str) -> tuple[str, str | None]: if not files: return {} with ThreadPoolExecutor(max_workers=_UPLOAD_WORKERS) as pool: - done = list(pool.map(_one, files)) + done = list(pool.map(write_fence.carry(_one), files)) return {slug: key for slug, key in done if key} diff --git a/tests/agent_stop_fakes.py b/tests/agent_stop_fakes.py index bd9e67de..55294b04 100644 --- a/tests/agent_stop_fakes.py +++ b/tests/agent_stop_fakes.py @@ -24,8 +24,11 @@ def __init__(self, path: Path): self.path = Path(path) self._lock = threading.Lock() - def log(self, kind: str, name: str, who: str = "") -> None: - line = json.dumps({"t": time.time(), "kind": kind, "name": name, "who": who}) + def log(self, kind: str, name: str, who: str = "", detail: dict | None = None) -> None: + record = {"t": time.time(), "kind": kind, "name": name, "who": who} + if detail is not None: + record["detail"] = detail + line = json.dumps(record) with self._lock, self.path.open("a", encoding="utf-8") as f: f.write(line + "\n") @@ -92,13 +95,18 @@ class FakeMCP(ToolProvider): hold_loop: write_slide blocks the calling agent's event loop for its whole duration, so nothing else on that loop (Strands' cancel watcher) runs meanwhile. + meta: the client-level ``_meta`` the factory gave it (the turn's write-fence + generation); recorded at construction and with every direct call. """ server_instructions = "" - def __init__(self, recorder: Recorder, who: str, tool_seconds: float = 0.5, hold_loop: bool = False): + def __init__(self, recorder: Recorder, who: str, tool_seconds: float = 0.5, hold_loop: bool = False, + meta: dict | None = None): self.recorder = recorder self.who = who + self.meta = meta + recorder.log("client", "created", who, detail={"meta": meta}) if hold_loop: @tool(name="write_slide") @@ -130,7 +138,7 @@ def stop(self, *exc_info): pass def call_tool_sync(self, tool_use_id, name, arguments, **kwargs): - self.recorder.log("mcp", name, self.who) + self.recorder.log("mcp", name, self.who, detail={"arguments": arguments, "meta": self.meta}) return {"status": "success", "toolUseId": tool_use_id, "content": [{"text": json.dumps(_PAYLOAD)}]} diff --git a/tests/agent_stop_server.py b/tests/agent_stop_server.py index 0892c9e6..9f16c2d0 100644 --- a/tests/agent_stop_server.py +++ b/tests/agent_stop_server.py @@ -24,7 +24,9 @@ import factory # noqa: E402 factory.BedrockModel = lambda **_kwargs: ScriptedModel(recorder) -factory._MCP_FACTORIES = [lambda jwt_token, session_id="", tool_filters=None: FakeMCP(recorder, "parent")] -factory.mcp_agentcore_runtime = lambda **_kwargs: FakeMCP(recorder, "composer") +factory._MCP_FACTORIES = [ + lambda jwt_token, session_id="", tool_filters=None, meta=None: FakeMCP(recorder, "parent", meta=meta), +] +factory.mcp_agentcore_runtime = lambda meta=None, **_kwargs: FakeMCP(recorder, "composer", meta=meta) basic_agent.main(port=int(sys.argv[1])) diff --git a/tests/test_agent_stop_guard.py b/tests/test_agent_stop_guard.py index e10bca72..7883b86b 100644 --- a/tests/test_agent_stop_guard.py +++ b/tests/test_agent_stop_guard.py @@ -104,7 +104,19 @@ def client() -> None: # keeps the connection open, like the Web UI's stream detail = log_path.read_text(errors="replace")[-4000:] assert lived < 4, detail assert {r["kind"] for r in records if r["t"] < sent} >= {"memory", "mcp", "model", "tool"}, records - assert late == [], (late, detail) + # The only call after the signal fences the turn's deck (turn_fence.py), with the + # generation every Presentation Maker client of the turn carries, composers' included. + fences = [r for r in records if r["name"] == "fence_turn"] + assert [r for r in late if r["name"] != "fence_turn"] == [], (late, detail) + assert [(r["kind"], r["who"]) for r in fences] == [("mcp", "parent")], (fences, detail) + assert fences[0]["t"] >= sent + assert fences[0]["detail"]["arguments"] == {"deck_ids": ["deck1"]} + generation = fences[0]["detail"]["meta"]["sdpm/generation"] + clients = [r for r in records if r["kind"] == "client"] + assert {r["who"] for r in clients} == {"parent", "composer"} + assert all(r["detail"]["meta"] == {"sdpm/generation": generation} for r in clients), clients + direct = [r for r in records if r["kind"] == "mcp" and r["name"] != "fence_turn"] + assert direct and all(r["detail"]["meta"] == {"sdpm/generation": generation} for r in direct), direct def test_composers_stop_after_the_parent_invocation_ended(tmp_path, monkeypatch): diff --git a/tests/test_agent_turn_fence.py b/tests/test_agent_turn_fence.py new file mode 100644 index 00000000..c9532384 --- /dev/null +++ b/tests/test_agent_turn_fence.py @@ -0,0 +1,256 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""The agent's half of the MCP write fence (agent/turn_fence.py). + +Every Presentation Maker client of a turn (the orchestrator's, each composer group's) +carries the turn's generation in its ``_meta``; a stopped turn calls fence_turn for the +decks its tool calls named, and a turn that ends normally fences nothing. The SIGTERM +path end to end is in test_agent_stop_guard. +""" + +import threading + + +# Imported lazily: test_mcp_parallel_connect puts agent/ on sys.path when it is imported, +# and agent/tools would then shadow servers/remote's tools package for the test modules +# collected after this one. +def _factory(): + from test_mcp_parallel_connect import _factory as load + + return load() + + +def _basic_agent(): + from test_agent_tool_exposure import _basic_agent as load + + return load() + + +def _run(basic_agent, payload): + from test_agent_tool_exposure import _run as run + + return run(basic_agent, payload) + + +class _Agent: + def __init__(self, **kwargs): + self.messages = [] + self.tool_names = [] + self.system_prompt = "" + + +class _Event: + def __init__(self, tool_input): + self.tool_use = {"name": "run_python", "toolUseId": "t1", "input": tool_input} + + +class _Client: + def __init__(self): + self.calls: list[dict] = [] + self.consumers: list[tuple[str, str]] = [] + + def add_consumer(self, consumer_id, **kwargs): + self.consumers.append(("add", consumer_id)) + + def remove_consumer(self, consumer_id, **kwargs): + self.consumers.append(("remove", consumer_id)) + + def call_tool_sync(self, tool_use_id, name, arguments, **kwargs): + self.calls.append({"name": name, "arguments": arguments, **kwargs}) + return {"status": "success", "content": [{"text": "{}"}]} + + +def _fence_class(): + return _factory().TurnFence + + +def test_a_stopped_turn_fences_the_decks_its_calls_named_once(): + fence = _fence_class()(generation=123) + client = _Client() + fence.bind(client) + fence.note_tool_call(_Event({"deck_id": "d2", "code": "x"})) + fence.note_tool_call(_Event({"deck_id": "d1"})) + fence.note_tool_call(_Event({"deck_id": "d1"})) + fence.note_tool_call(_Event({"query": "no deck"})) + fence.note_tool_call(_Event("not a dict")) + + thread = fence.advance() + thread.join(5) + + # The client's own meta carries the generation; the call names the decks. + assert client.calls == [{"name": "fence_turn", "arguments": {"deck_ids": ["d1", "d2"]}}] + assert fence.meta == {"sdpm/generation": 123} + # Held while the call runs, released after it. + assert client.consumers == [("add", "sdpm-turn-fence"), ("remove", "sdpm-turn-fence")] + assert fence.advance() is None # once: SIGTERM and the turn's finally both call it + assert len(client.calls) == 1 + + +def test_a_turn_that_ends_normally_fences_nothing(): + fence = _fence_class()(generation=5) + client = _Client() + fence.bind(client) + fence.note_tool_call(_Event({"deck_id": "d1"})) + + fence.release() + for t in threading.enumerate(): + if t.name == "sdpm-turn-fence": + t.join(5) + + assert client.calls == [] + assert client.consumers == [("add", "sdpm-turn-fence"), ("remove", "sdpm-turn-fence")] + assert fence.advance() is None + + +def test_a_stop_before_any_deck_or_client_does_nothing(): + fence = _fence_class()() + assert fence.advance() is None # no client bound yet (SIGTERM during create_agent) + client = _Client() + fence.bind(client) # bound after the stop: never held, never called + assert client.consumers == [] and client.calls == [] + + +def test_generations_grow_from_turn_to_turn(): + first, second = _fence_class()(), _fence_class()() + assert 0 < first.generation <= second.generation + + +def test_orchestrator_and_composer_clients_carry_the_generation(monkeypatch): + factory = _factory() + runtime_calls: list[dict] = [] + captured: dict = {} + hooks: list = [] + + class _Hooks: + def add_callback(self, event, callback): + hooks.append(callback) + + class Agent(_Agent): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.hooks = _Hooks() + + client = _Client() + + def runtime(**kwargs): + runtime_calls.append(kwargs) + return client + + def make_compose_slides(mcp_servers, model, composer_mcp_factory, **kwargs): + captured["factory"] = composer_mcp_factory + return object() + + monkeypatch.delenv("MEMORY_ID", raising=False) + monkeypatch.setattr(factory, "mcp_agentcore_runtime", runtime) + monkeypatch.setattr(factory, "mcp_aws_knowledge", lambda: object()) + monkeypatch.setattr(factory, "mcp_aws_pricing", lambda: object()) + monkeypatch.setattr(factory, "_prewarm_mcp_clients", lambda clients, names, required, then=None: (clients, [])) + monkeypatch.setattr(factory, "BedrockModel", lambda **kwargs: object()) + monkeypatch.setattr(factory, "make_compose_slides", make_compose_slides) + monkeypatch.setattr(factory, "Agent", Agent) + monkeypatch.setattr(factory, "collect_mcp_instructions", lambda servers: "") + monkeypatch.setattr(factory, "resolve_parts", lambda *a, **k: ("", [])) + + fence = factory.TurnFence(generation=42) + factory.create_agent("spec", user_id="u1", session_id="s1", jwt_token="jwt", fence=fence) + (orchestrator,) = runtime_calls + runtime_calls.clear() + captured["factory"]("s1-g1") + (composer,) = runtime_calls + + assert orchestrator["meta"] == {"sdpm/generation": 42} + assert composer["meta"] == {"sdpm/generation": 42} + assert composer["session_id"] == "s1-g1" + # The fence learns the turn's decks from the orchestrator's tool calls and fences + # through the orchestrator's client (compose_slides names its deck there). + assert any(getattr(h, "__self__", None) is fence and h.__name__ == "note_tool_call" for h in hooks) + assert ("add", "sdpm-turn-fence") in client.consumers + + # Without a fence (callers other than the entrypoint) the clients get no meta. + runtime_calls.clear() + factory.create_agent("spec", user_id="u1", session_id="s1", jwt_token="jwt") + assert "meta" not in runtime_calls[0] + + +def test_mcp_client_sends_the_meta_with_every_call(monkeypatch): + import sys + + _factory() + mcp_clients = sys.modules["mcp_clients"] + seen: dict = {} + + def fake_client(transport_callable, **kwargs): + seen.update(kwargs) + return object() + + monkeypatch.setattr(mcp_clients, "MCPClient", fake_client) + monkeypatch.setenv("MCP_RUNTIME_ARN", "arn:aws:bedrock-agentcore:ap-northeast-1:111122223333:runtime/x") + mcp_clients.mcp_agentcore_runtime(jwt_token="tok", session_id="s", meta={"sdpm/generation": 7}) + assert seen["meta"] == {"sdpm/generation": 7} + seen.clear() + mcp_clients.mcp_agentcore_runtime(jwt_token="tok", session_id="s") + assert "meta" not in seen + + +class _FenceSpy: + instances: list = [] + + def __init__(self): + self.meta = {"sdpm/generation": 1} + self.events: list[str] = [] + _FenceSpy.instances.append(self) + + def advance(self): + self.events.append("advance") + + def release(self): + self.events.append("release") + + +def _entrypoint(monkeypatch, stream): + basic_agent = _basic_agent() + _FenceSpy.instances = [] + seen: dict = {} + + def fake_create_agent(**kwargs): + seen.update(kwargs) + return object(), [] + + monkeypatch.setattr(basic_agent, "TurnFence", _FenceSpy) + monkeypatch.setattr(basic_agent, "create_agent", fake_create_agent) + monkeypatch.setattr(basic_agent, "stream_agent", stream) + _run(basic_agent, {"prompt": "hello", "runtimeSessionId": "s" * 33}) + (fence,) = _FenceSpy.instances + assert seen["fence"] is fence + return fence + + +def test_a_cancelled_turn_advances_the_fence(monkeypatch): + async def cancelled(agent, user_query, session_id, cancel, cancel_signal): + yield {"keepalive": True} + cancel_signal.set() # what stream_agent does when the turn is cancelled + + assert _entrypoint(monkeypatch, cancelled).events == ["advance"] + + +def test_a_finished_turn_releases_the_fence(monkeypatch): + async def finished(agent, user_query, session_id, cancel, cancel_signal): + yield {"done": True} + + assert _entrypoint(monkeypatch, finished).events == ["release"] + + +def test_sigterm_advances_the_running_turns_fence(monkeypatch): + import sys + + stop_guard = _basic_agent().stop_guard + resilience = sys.modules["resilience"] + monkeypatch.setattr(resilience, "_writes_closed", threading.Event()) + monkeypatch.setattr(stop_guard, "_running", {}) + calls = [] + stop = stop_guard.begin_turn(on_stop=lambda: calls.append("advance")) + try: + assert stop_guard.stop_all() == 1 + finally: + stop_guard.end_turn(stop) + assert stop.is_set() and calls == ["advance"] diff --git a/tests/test_remote_html.py b/tests/test_remote_html.py index ceb6b54f..81997611 100644 --- a/tests/test_remote_html.py +++ b/tests/test_remote_html.py @@ -839,11 +839,11 @@ def test_prepare_rejects_refs_outside_the_deck(tmp_path, cache_root): def test_remote_tool_count_unchanged(): - """20 tools as before; a configured knowledge base adds only search_slides.""" + """21 tools (20 + the agent's fence_turn); a configured knowledge base adds only search_slides.""" import asyncio without_kb = {t.name for t in asyncio.run(remote_server.mcp.list_tools())} - assert len(without_kb) == 20 and "search_slides" not in without_kb + assert len(without_kb) == 21 and "search_slides" not in without_kb and "fence_turn" in without_kb env = { "DECKS_TABLE": "decks-test", "PPTX_BUCKET": "pptx-test", "RESOURCE_BUCKET": "resource-test", diff --git a/tests/test_remote_write_fence.py b/tests/test_remote_write_fence.py new file mode 100644 index 00000000..93491bbb --- /dev/null +++ b/tests/test_remote_write_fence.py @@ -0,0 +1,370 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: MIT-0 +"""MCP write fence (servers/remote/storage/write_fence.py). + +A stopped agent turn's in-flight calls kept writing on dev: run_python put 4 slides +1.7 s after a cancel, and generate_pptx's background previews landed 1.5 s after a +StopRuntimeSession. With the fence, a request that carries a generation at or below the +deck's fence writes nothing; a newer generation writes; a request without one behaves +as before. The deck record update is conditional on the fence, so a write that passed +the check cannot repoint the deck after the fence moved. +""" + +from __future__ import annotations + +import asyncio +import importlib.util +import json +import os +import re +import socket +import threading +import time +from pathlib import Path +from unittest.mock import patch + +import pytest +from botocore.exceptions import ClientError +from mcp.shared.memory import create_connected_server_and_client_session + +from storage import write_fence +from storage.aws import AwsStorage + +_ROOT = Path(__file__).resolve().parents[1] +_SERVER = _ROOT / "servers" / "remote" / "server.py" + +os.environ.setdefault("DECKS_TABLE", "test-table") +os.environ.setdefault("PPTX_BUCKET", "test-pptx") +os.environ.setdefault("RESOURCE_BUCKET", "test-resource") + +_spec = importlib.util.spec_from_file_location("remote_server_write_fence_under_test", _SERVER) +remote_server = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(remote_server) + +USER = "u1" +DECK = "d1" +KEY = ("USER#u1", "DECK#d1") + + +class FakeTable: + """The deck table: GetItem, and UpdateItem with the SET/REMOVE and condition forms + AwsStorage.update_deck and write_fence.advance send.""" + + def __init__(self) -> None: + self.items: dict[tuple[str, str], dict] = {} + self.gets = 0 + self.before_update = None # hook: runs right before an update is applied (races) + + def get_item(self, Key, **kwargs): + self.gets += 1 + item = self.items.get((Key["PK"], Key["SK"])) + return {"Item": dict(item)} if item is not None else {} + + def update_item(self, Key, UpdateExpression, ConditionExpression=None, ExpressionAttributeNames=None, + ExpressionAttributeValues=None, ReturnValues=None): + if self.before_update is not None: + hook, self.before_update = self.before_update, None + hook() + names = ExpressionAttributeNames or {} + values = ExpressionAttributeValues or {} + k = (Key["PK"], Key["SK"]) + item = self.items.get(k) + if ConditionExpression: + fence = (item or {}).get(write_fence.FENCE_ATTR) + ok = fence is None or fence < values[":wf_gen"] + if ConditionExpression.startswith("attribute_exists(PK)"): + ok = ok and item is not None + if not ok: + raise ClientError({"Error": {"Code": "ConditionalCheckFailedException"}}, "UpdateItem") + item = self.items.setdefault(k, {"PK": k[0], "SK": k[1]}) + old = {} + for clause in re.findall(r"(SET|REMOVE) (.*?)(?= REMOVE | SET |$)", UpdateExpression): + verb, body = clause + for part in body.split(","): + part = part.strip() + if verb == "SET": + name, value = (s.strip() for s in part.split("=")) + name = names.get(name, name) + if name in item: + old[name] = item[name] + item[name] = values[value] + else: + name = names.get(part, part) + if name in item: + old[name] = item.pop(name) + return {"Attributes": old} + + +class FakeS3: + class exceptions: # noqa: N801 - mirrors the boto3 client attribute + class NoSuchKey(Exception): + pass + + def __init__(self) -> None: + self.writes: list[tuple[str, str]] = [] + + def put_object(self, Bucket, Key, Body, **kwargs): + self.writes.append(("put", Key)) + + def delete_object(self, Bucket, Key): + self.writes.append(("delete", Key)) + + def delete_objects(self, Bucket, Delete): + self.writes.extend(("delete", o["Key"]) for o in Delete["Objects"]) + + +@pytest.fixture() +def store(): + table, s3 = FakeTable(), FakeS3() + table.items[KEY] = {"PK": KEY[0], "SK": KEY[1], "name": "deck"} + return AwsStorage(table=table, s3_client=s3, pptx_bucket="pptx", resource_bucket="res"), table, s3 + + +def _fenced(generation: int): + return write_fence.bind(write_fence.Fence(generation=generation, user_id=USER, deck_id=DECK)) + + +def _write_everything(storage: AwsStorage) -> None: + storage.upload_file(key=f"decks/{DECK}/slides/a.json", data=b"{}") + storage.put_slide_json(DECK, "b", {}) + storage.put_deck_json(DECK, {}) + storage._s3.delete_object(Bucket="pptx", Key=f"previews/{DECK}/old.webp") # the direct deletes + storage.delete_object_keys([f"previews/{DECK}/x.webp"]) + + +# --- The fence at the storage choke point --- + + +def test_without_a_generation_every_write_goes_through(store): + storage, table, s3 = store + table.items[KEY][write_fence.FENCE_ATTR] = 10**30 # even a fenced deck: not this request's business + _write_everything(storage) + storage.update_deck(DECK, USER, {"updatedAt": "t"}) + assert len(s3.writes) == 5 + assert table.items[KEY]["updatedAt"] == "t" + assert table.gets == 0 + + +def test_a_current_generation_writes(store): + storage, table, s3 = store + table.items[KEY][write_fence.FENCE_ATTR] = 100 + token = _fenced(101) + try: + _write_everything(storage) + old = storage.update_deck(DECK, USER, {"pptxS3Key": "pptx/d1/new.pptx"}) + finally: + write_fence.reset(token) + assert len(s3.writes) == 5 + assert table.items[KEY]["pptxS3Key"] == "pptx/d1/new.pptx" + assert old == {} + + +def test_a_deck_never_fenced_accepts_any_generation(store): + storage, _, s3 = store + token = _fenced(1) + try: + storage.upload_file(key=f"decks/{DECK}/x", data=b"") + finally: + write_fence.reset(token) + assert s3.writes == [("put", f"decks/{DECK}/x")] + + +@pytest.mark.parametrize("generation", [99, 100]) +def test_a_stopped_generation_writes_nothing(store, generation): + storage, table, s3 = store + table.items[KEY][write_fence.FENCE_ATTR] = 100 + token = _fenced(generation) + try: + with pytest.raises(write_fence.StaleTurn, match="Cancelled"): + storage.upload_file(key=f"decks/{DECK}/slides/a.json", data=b"{}") + # Every later write of the request is skipped too, without another read. + gets = table.gets + for write in ( + lambda: storage.put_slide_json(DECK, "b", {}), + lambda: storage._s3.delete_object(Bucket="pptx", Key="previews/d1/old.webp"), + lambda: storage.delete_object_keys(["previews/d1/x.webp"]), + lambda: storage.update_deck(DECK, USER, {"thumbnailS3Key": "previews/d1/t.webp"}), + ): + with pytest.raises(write_fence.StaleTurn): + write() + assert table.gets == gets + finally: + write_fence.reset(token) + assert s3.writes == [] + assert "thumbnailS3Key" not in table.items[KEY] + + +def test_the_deck_update_cannot_race_the_fence(store): + storage, table, s3 = store + table.items[KEY][write_fence.FENCE_ATTR] = 100 + token = _fenced(101) + try: + storage.upload_file(key="pptx/d1/new.pptx", data=b"") # passed the check: generation 101 is current + # The turn is stopped between the check and the record update. + table.before_update = lambda: write_fence.advance(table, USER, DECK, 101) + with pytest.raises(write_fence.StaleTurn): + storage.update_deck(DECK, USER, {"pptxS3Key": "pptx/d1/new.pptx", "updatedAt": "t"}) + with pytest.raises(write_fence.StaleTurn): # and the rest of the request stops writing + storage._s3.delete_object(Bucket="pptx", Key="pptx/d1/old.pptx") + finally: + write_fence.reset(token) + assert "pptxS3Key" not in table.items[KEY] + assert s3.writes == [("put", "pptx/d1/new.pptx")] + + +def test_the_fence_only_moves_forward_and_only_on_the_owners_deck(): + table = FakeTable() + table.items[KEY] = {"PK": KEY[0], "SK": KEY[1]} + assert write_fence.advance(table, USER, DECK, 200) is True + assert write_fence.advance(table, USER, DECK, 150) is False + assert write_fence.advance(table, USER, DECK, 200) is False + assert table.items[KEY][write_fence.FENCE_ATTR] == 200 + # Another user's sub names a different item: nothing exists there, nothing is created. + assert write_fence.advance(table, "u2", DECK, 300) is False + assert ("USER#u2", "DECK#d1") not in table.items + + +@pytest.mark.parametrize("meta, expected", [ + (None, None), + ({}, None), + ({"sdpm/generation": 5}, 5), + ({"sdpm/generation": "5"}, None), + ({"sdpm/generation": True}, None), + ({"sdpm/generation": 0}, None), + ({"sdpm/generation": time.time_ns() + 10**15}, None), # from the future (~12 days) +]) +def test_generation_from_meta(meta, expected): + assert write_fence.generation_from_meta(meta) == expected + + +def test_threads_the_request_starts_keep_its_fence(store, monkeypatch): + storage, table, s3 = store + table.items[KEY][write_fence.FENCE_ATTR] = 100 + import server_utils + from tools import html_deck + + seen: list = [] + done = threading.Event() + + def background(): + seen.append(write_fence.current()) + done.set() + + def webp(*args, **kwargs): + background() + + monkeypatch.setattr(server_utils, "_generate_webp_background", webp) + token = _fenced(50) + try: + remote_server._run_in_background(background, name="t") # run_python's compose + previews + assert done.wait(5) + done.clear() + server_utils.schedule_webp_background(DECK, Path("x"), Path("y"), storage, []) # generate_pptx's previews + assert done.wait(5) + # html_deck's upload pool: the stale uploads are skipped (and reported as failed). + files = {"a": Path(__file__), "b": Path(__file__)} + assert html_deck._upload_all(DECK, files, lambda slug: f"previews/{DECK}/{slug}.html", "text/html", storage) == {} + finally: + write_fence.reset(token) + assert [f.generation for f in seen] == [50, 50] + assert s3.writes == [] + + +# --- Through the MCP server --- + + +def _call(tool: str, arguments: dict, meta: dict | None): + async def run(): + async with create_connected_server_and_client_session(remote_server.mcp) as session: + return await session.call_tool(tool, arguments, meta=meta) + + return asyncio.run(run()) + + +@pytest.fixture() +def served(store): + storage, table, s3 = store + with ( + patch.object(remote_server, "_storage", storage), + patch.object(remote_server, "_get_user_id", return_value=USER), + patch.object(remote_server, "_check_deck_access", return_value=None), + ): + yield table, s3 + + +_CODE = {"deck_id": DECK, "code": "print(1)", "name": "snippet"} + + +def test_a_stopped_turns_call_is_cancelled_and_writes_nothing(served): + table, s3 = served + table.items[KEY][write_fence.FENCE_ATTR] = 100 + result = _call("code_to_slide", _CODE, {"sdpm/generation": 100}) + assert result.isError + assert "Cancelled" in result.content[0].text + assert s3.writes == [] + + +def test_a_current_turns_call_writes(served): + table, s3 = served + table.items[KEY][write_fence.FENCE_ATTR] = 100 + result = _call("code_to_slide", _CODE, {"sdpm/generation": 101}) + assert not result.isError, result.content[0].text + assert s3.writes == [("put", f"decks/{DECK}/includes/snippet.json")] + + +def test_a_call_without_a_generation_writes_as_before(served): + table, s3 = served + table.items[KEY][write_fence.FENCE_ATTR] = 10**30 + result = _call("code_to_slide", _CODE, None) + assert not result.isError, result.content[0].text + assert s3.writes == [("put", f"decks/{DECK}/includes/snippet.json")] + + +def test_fence_turn_fences_the_callers_decks_with_the_requests_generation(served): + table, _ = served + result = _call("fence_turn", {"deck_ids": [DECK, DECK, "not-mine"]}, {"sdpm/generation": 300}) + assert json.loads(result.content[0].text) == {"generation": 300, "fenced": [DECK], "unchanged": ["not-mine"]} + assert table.items[KEY][write_fence.FENCE_ATTR] == 300 + # The call it fenced: the same generation now writes nothing. + assert _call("code_to_slide", _CODE, {"sdpm/generation": 300}).isError + + +def test_fence_turn_needs_a_generation(served): + table, _ = served + result = _call("fence_turn", {"deck_ids": [DECK]}, None) + assert "error" in json.loads(result.content[0].text) + assert write_fence.FENCE_ATTR not in table.items[KEY] + + +def test_the_generation_survives_strands_mcpclient_over_stateless_http(served): + """The agent's path: Strands MCPClient with a client-level meta, streamable HTTP, the + server's stateless_http FastMCP app. (The AgentCore proxy forwards the JSON-RPC body, + _meta included; dev verification covers that hop.)""" + import uvicorn + from mcp.client.streamable_http import streamablehttp_client + from strands.tools.mcp import MCPClient + + table, s3 = served + table.items[KEY][write_fence.FENCE_ATTR] = 100 + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + server = uvicorn.Server(uvicorn.Config(remote_server.mcp.streamable_http_app(), host="127.0.0.1", port=port, log_level="warning")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + try: + deadline = time.time() + 20 + while not server.started: + assert time.time() < deadline, "server did not start" + time.sleep(0.05) + for generation, written in ((100, False), (101, True)): + s3.writes.clear() + client = MCPClient(lambda: streamablehttp_client(f"http://127.0.0.1:{port}/mcp"), + meta={"sdpm/generation": generation}) + with client: + result = client.call_tool_sync(tool_use_id="t1", name="code_to_slide", arguments=_CODE) + assert (result["status"] == "success") is written, result + assert bool(s3.writes) is written + finally: + server.should_exit = True + thread.join(10) From 942b81961384385e9acbc9029fe5ccbd0d10add8 Mon Sep 17 00:00:00 2001 From: ShotaroKataoka Date: Tue, 3 Nov 2026 18:36:21 +0900 Subject: [PATCH 2/3] fix(agent): wait for a stopped turn's fence_turn call before the process exits on SIGTERM The fence call ran on a daemon thread that nothing waited for. A turn that shut down quickly (PR-1's dev runs exited 0.5-1.1 s after SIGTERM) let the process exit before fence_turn reached the MCP runtime, leaving the deck unfenced. A non-daemon thread would not help: after the lifespan shutdown uvicorn restores the default SIGTERM handler and re-raises the signal. stop_guard.hold() now keeps the thread TurnFence.advance() returns (from stop_all on SIGTERM, and from agent_stream's finally for a cancel followed by a stop), and the app lifespan's shutdown joins the held threads until the watchdog's exit (SIGTERM + 3 s). Tests: the SIGTERM end-to-end test's fake fence_turn now takes 1 s and must still be recorded (it failed before this fix); a unit test covers the wait and its bound. Checks run on this commit: - make lint: All checks passed - make test: 2130 passed, 18 skipped - ash scan --mode local --fail-on-findings: exit 0 - infra unchanged (no infra npm test / tsc needed) --- agent/basic_agent.py | 3 ++- agent/stop_guard.py | 47 ++++++++++++++++++++++++++++++---- agent/turn_fence.py | 3 ++- tests/agent_stop_fakes.py | 7 ++++- tests/agent_stop_server.py | 4 ++- tests/test_agent_stop_guard.py | 11 ++++++-- tests/test_agent_turn_fence.py | 39 ++++++++++++++++++++++++++++ 7 files changed, 103 insertions(+), 11 deletions(-) diff --git a/agent/basic_agent.py b/agent/basic_agent.py index fa3c755f..12a170a9 100644 --- a/agent/basic_agent.py +++ b/agent/basic_agent.py @@ -120,8 +120,9 @@ async def agent_stream(payload, context): finally: stop_guard.end_turn(stop) # Stopped (cancel, a newer request, the client gone, SIGTERM): fence the turn's decks. + # hold: a SIGTERM that follows (StopRuntimeSession after a cancel) waits for the call. if stop.is_set(): - fence.advance() + stop_guard.hold(fence.advance()) else: fence.release() if _cancel_events.get(session_id) is cancel: diff --git a/agent/stop_guard.py b/agent/stop_guard.py index cd0cfae1..8eee13d4 100644 --- a/agent/stop_guard.py +++ b/agent/stop_guard.py @@ -11,14 +11,18 @@ The guard, installed through the app's lifespan so it runs before uvicorn's own handler: 1. closes the write gate (resilience.close_writes): tool calls, compose's direct MCP calls and Memory events (GatedMemorySessionManager) are dropped from then on; - 2. sets every running turn's stop event, which Strands and the composers observe; + 2. sets every running turn's stop event, which Strands and the composers observe, and + runs its stop callback (the turn's fence_turn call, turn_fence.py); 3. starts a watchdog that flushes OpenTelemetry and exits the process after EXIT_AFTER_S, whatever is still running (a composer thread blocked in a read); 4. closes the process's resident MCP clients (AWS Knowledge, and AWS Pricing's stdio subprocess) on a side thread, best effort: the watchdog's exit bounds it. -basic_agent passes GRACEFUL_SHUTDOWN_S to uvicorn so it stops waiting before that. +basic_agent passes GRACEFUL_SHUTDOWN_S to uvicorn so it stops waiting before that. The +lifespan's shutdown then waits, up to the watchdog's exit, for the stop calls still +running (hold / wait_held), since a turn can end before its fence_turn reaches the server. """ +import asyncio import contextlib import logging import os @@ -43,13 +47,19 @@ # Running turns: stop event -> what to call when the turn is stopped (or None). _running: dict[threading.Event, Callable[[], object] | None] = {} _exiting = threading.Event() +# Calls a stopped turn still has to send (turn_fence's fence_turn), and when the watchdog +# exits the process. The lifespan's shutdown waits for them until then: right after it, +# uvicorn re-raises SIGTERM under the default handler, which ends the process at once. +_held: set[threading.Thread] = set() +_exit_at: float | None = None def begin_turn(on_stop: Callable[[], object] | None = None) -> threading.Event: """Register a running turn and return its stop event (set on SIGTERM). ``on_stop`` runs when SIGTERM stops the turn (on the signal handler's thread, so it - must not block: turn_fence.TurnFence.advance hands its call to a thread). + must not block: turn_fence.TurnFence.advance hands its call to a thread and returns + it, and the process waits for a returned thread before it exits; see hold()). Pair with end_turn() in a finally block. """ stop = threading.Event() @@ -65,6 +75,27 @@ def end_turn(stop: threading.Event) -> None: _running.pop(stop, None) +def hold(thread: object) -> None: + """Keep the process alive for ``thread`` (if it is one) when it shuts down, up to the + watchdog's exit.""" + if not isinstance(thread, threading.Thread): + return + with _lock: + _held.difference_update([t for t in _held if not t.is_alive()]) + _held.add(thread) + + +def wait_held() -> None: + """Wait for the held threads, until the watchdog's exit (EXIT_AFTER_S without SIGTERM).""" + with _lock: + threads = list(_held) + deadline = _exit_at if _exit_at is not None else time.monotonic() + EXIT_AFTER_S + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + if any(t.is_alive() for t in threads): + logger.warning("SIGTERM guard: exiting with %d stop call(s) still running", sum(t.is_alive() for t in threads)) + + def stop_all() -> int: """Close the write gate and set the stop event of every running turn.""" close_writes() @@ -74,7 +105,7 @@ def stop_all() -> int: stop.set() if on_stop is not None: try: - on_stop() + hold(on_stop()) except Exception as e: # noqa: BLE001 - the other turns still stop logger.warning("SIGTERM guard: a turn's stop callback failed: %s", e) return len(running) @@ -127,6 +158,10 @@ def _close_resident_clients() -> None: def _on_sigterm(previous): def handler(signum, frame): + global _exit_at + if not _exiting.is_set(): # before stop_all, so wait_held already sees the deadline + with _lock: + _exit_at = time.monotonic() + EXIT_AFTER_S stopped = stop_all() if not _exiting.is_set(): # signal handlers run on the main thread only: no race here _exiting.set() @@ -148,6 +183,8 @@ def install() -> None: @contextlib.asynccontextmanager async def lifespan(app): - """App lifespan: uvicorn has installed its signal handlers by the time this starts.""" + """App lifespan: uvicorn has installed its signal handlers by the time this starts. + Its shutdown waits for the stop calls still running (wait_held).""" install() yield + await asyncio.to_thread(wait_held) diff --git a/agent/turn_fence.py b/agent/turn_fence.py index 373997dc..33fb052a 100644 --- a/agent/turn_fence.py +++ b/agent/turn_fence.py @@ -72,7 +72,8 @@ def note_tool_call(self, event) -> None: def advance(self) -> threading.Thread | None: """The turn was stopped: fence its decks (once). The call runs on a daemon thread, - so this is safe from a signal handler and from the event loop; returns it.""" + so this is safe from a signal handler and from the event loop; returns it (for + stop_guard.hold, so a shutting-down process waits for it).""" return self._finish(fence=True) def release(self) -> None: diff --git a/tests/agent_stop_fakes.py b/tests/agent_stop_fakes.py index 55294b04..676972dd 100644 --- a/tests/agent_stop_fakes.py +++ b/tests/agent_stop_fakes.py @@ -97,15 +97,18 @@ class FakeMCP(ToolProvider): so nothing else on that loop (Strands' cancel watcher) runs meanwhile. meta: the client-level ``_meta`` the factory gave it (the turn's write-fence generation); recorded at construction and with every direct call. + fence_seconds: fence_turn is recorded only after this long, like a request that + reaches a cold MCP runtime late; a process that exits sooner never records it. """ server_instructions = "" def __init__(self, recorder: Recorder, who: str, tool_seconds: float = 0.5, hold_loop: bool = False, - meta: dict | None = None): + meta: dict | None = None, fence_seconds: float = 0.0): self.recorder = recorder self.who = who self.meta = meta + self.fence_seconds = fence_seconds recorder.log("client", "created", who, detail={"meta": meta}) if hold_loop: @@ -138,6 +141,8 @@ def stop(self, *exc_info): pass def call_tool_sync(self, tool_use_id, name, arguments, **kwargs): + if name == "fence_turn": + time.sleep(self.fence_seconds) self.recorder.log("mcp", name, self.who, detail={"arguments": arguments, "meta": self.meta}) return {"status": "success", "toolUseId": tool_use_id, "content": [{"text": json.dumps(_PAYLOAD)}]} diff --git a/tests/agent_stop_server.py b/tests/agent_stop_server.py index 9f16c2d0..c8db45c6 100644 --- a/tests/agent_stop_server.py +++ b/tests/agent_stop_server.py @@ -25,7 +25,9 @@ factory.BedrockModel = lambda **_kwargs: ScriptedModel(recorder) factory._MCP_FACTORIES = [ - lambda jwt_token, session_id="", tool_filters=None, meta=None: FakeMCP(recorder, "parent", meta=meta), + lambda jwt_token, session_id="", tool_filters=None, meta=None: FakeMCP( + recorder, "parent", meta=meta, fence_seconds=float(os.environ.get("FAKE_FENCE_SECONDS", "0")), + ), ] factory.mcp_agentcore_runtime = lambda meta=None, **_kwargs: FakeMCP(recorder, "composer", meta=meta) diff --git a/tests/test_agent_stop_guard.py b/tests/test_agent_stop_guard.py index 7883b86b..bf288172 100644 --- a/tests/test_agent_stop_guard.py +++ b/tests/test_agent_stop_guard.py @@ -14,6 +14,7 @@ import asyncio import base64 import json +import os import signal import socket import subprocess @@ -33,6 +34,8 @@ # The signal reaches the handler on the server's main thread within milliseconds; a call # recorded later than this after the test sent it came after the guard closed the gate. SLACK_S = 0.05 +# How long the fake fence_turn takes in the SIGTERM test (below stop_guard.EXIT_AFTER_S). +FENCE_SECONDS = 1.0 def _free_port() -> int: @@ -58,10 +61,13 @@ def test_sigterm_during_compose_stops_every_write_and_exits(tmp_path): port = _free_port() recorder = Recorder(tmp_path / "calls.jsonl") log_path = tmp_path / "server.log" + # fence_turn takes longer than the turn takes to shut down (about 0.5-1.1 s on dev), + # like a cold MCP runtime: the process must still wait for it before it exits. + env = {**os.environ, "FAKE_FENCE_SECONDS": str(FENCE_SECONDS)} with log_path.open("w") as log: server = subprocess.Popen( [sys.executable, str(ROOT / "tests" / "agent_stop_server.py"), str(port), str(recorder.path)], - cwd=AGENT_DIR, stdout=log, stderr=subprocess.STDOUT, + cwd=AGENT_DIR, stdout=log, stderr=subprocess.STDOUT, env=env, ) try: def up() -> bool: @@ -109,7 +115,7 @@ def client() -> None: # keeps the connection open, like the Web UI's stream fences = [r for r in records if r["name"] == "fence_turn"] assert [r for r in late if r["name"] != "fence_turn"] == [], (late, detail) assert [(r["kind"], r["who"]) for r in fences] == [("mcp", "parent")], (fences, detail) - assert fences[0]["t"] >= sent + assert fences[0]["t"] >= sent + FENCE_SECONDS assert fences[0]["detail"]["arguments"] == {"deck_ids": ["deck1"]} generation = fences[0]["detail"]["meta"]["sdpm/generation"] clients = [r for r in records if r["kind"] == "client"] @@ -202,6 +208,7 @@ def test_sigterm_closes_the_resident_mcp_clients_without_holding_the_handler(mon monkeypatch.setattr(resilience, "_writes_closed", threading.Event()) monkeypatch.setattr(stop_guard, "_exiting", threading.Event()) + monkeypatch.setattr(stop_guard, "_exit_at", None) monkeypatch.setattr(mcp_clients, "_RESIDENT_CLIENTS", {}) watchdog = threading.Event() monkeypatch.setattr(stop_guard, "_exit_after", lambda delay: watchdog.set()) # no os._exit here diff --git a/tests/test_agent_turn_fence.py b/tests/test_agent_turn_fence.py index c9532384..615aa718 100644 --- a/tests/test_agent_turn_fence.py +++ b/tests/test_agent_turn_fence.py @@ -254,3 +254,42 @@ def test_sigterm_advances_the_running_turns_fence(monkeypatch): finally: stop_guard.end_turn(stop) assert stop.is_set() and calls == ["advance"] + + +def test_shutdown_waits_for_the_stop_call_until_the_watchdog_exit(monkeypatch): + # The turn can end (and uvicorn finish) before fence_turn reaches the server; the + # lifespan's shutdown then waits for the call (wait_held), but not past the watchdog. + import sys + import time + + stop_guard = _basic_agent().stop_guard + resilience = sys.modules["resilience"] + monkeypatch.setattr(resilience, "_writes_closed", threading.Event()) + monkeypatch.setattr(stop_guard, "_running", {}) + monkeypatch.setattr(stop_guard, "_held", set()) + monkeypatch.setattr(stop_guard, "_exit_at", None) + sent = threading.Event() + + def advance(): + thread = threading.Thread(target=lambda: (time.sleep(0.3), sent.set()), daemon=True) + thread.start() + return thread + + stop = stop_guard.begin_turn(on_stop=advance) + try: + stop_guard.stop_all() + finally: + stop_guard.end_turn(stop) + stop_guard.wait_held() + assert sent.is_set() + + stuck = threading.Event() + hung = threading.Thread(target=stuck.wait, daemon=True) + hung.start() + stop_guard.hold(hung) + stop_guard.hold(None) # a turn with nothing to fence + monkeypatch.setattr(stop_guard, "_exit_at", time.monotonic() + 0.2) + started = time.monotonic() + stop_guard.wait_held() + assert time.monotonic() - started < 1.0 + stuck.set() From c0941959271997a136b3c28e80407392a9714203 Mon Sep 17 00:00:00 2001 From: ShotaroKataoka Date: Tue, 3 Nov 2026 19:10:05 +0900 Subject: [PATCH 3/3] fix(agent): copy turn_fence.py into the agent image The agent image COPYs its modules by name and missed the new turn_fence.py, so the dev agent (v266) failed at import: ModuleNotFoundError: No module named 'turn_fence'. Checks: make lint: All checks passed; make test: 2130 passed, 18 skipped --- agent/Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent/Dockerfile b/agent/Dockerfile index da39a4e5..ccec2eea 100644 --- a/agent/Dockerfile +++ b/agent/Dockerfile @@ -27,7 +27,7 @@ RUN uv export --frozen --directory /tmp/workspace --package sdpm-agent --no-dev RUN useradd -m -u 1000 bedrock_agentcore USER bedrock_agentcore -COPY agent/basic_agent.py agent/cost_logger.py agent/errors.py agent/factory.py agent/composition.py agent/mcp_clients.py agent/message_hooks.py agent/model_profiles.py agent/resilience.py agent/session.py agent/stop_guard.py agent/streaming.py ./ +COPY agent/basic_agent.py agent/cost_logger.py agent/errors.py agent/factory.py agent/composition.py agent/mcp_clients.py agent/message_hooks.py agent/model_profiles.py agent/resilience.py agent/session.py agent/stop_guard.py agent/streaming.py agent/turn_fence.py ./ COPY agent/prompts/ prompts/ COPY agent/modes/ modes/ COPY agent/tools/ tools/