From 849831bb202b5c32576078a794ddd1e421fa0ebd Mon Sep 17 00:00:00 2001
From: opieter-aws <opieter@amazon.com>
Date: Tue, 21 Jul 2026 12:35:38 -0400
Subject: [PATCH 1/5] fix(multiagent): preserve shared context and cumulative
 accounting across serialize/deserialize

---
 strands-py/src/strands/multiagent/graph.py    | 13 ++++-
 strands-py/src/strands/multiagent/swarm.py    | 13 ++++-
 .../tests/strands/multiagent/test_graph.py    | 49 +++++++++++++++++
 .../tests/strands/multiagent/test_swarm.py    | 54 +++++++++++++++++++
 4 files changed, 127 insertions(+), 2 deletions(-)

diff --git a/strands-py/src/strands/multiagent/graph.py b/strands-py/src/strands/multiagent/graph.py
index 5e9506adf0..1388d6dd18 100644
--- a/strands-py/src/strands/multiagent/graph.py
+++ b/strands-py/src/strands/multiagent/graph.py
@@ -56,7 +56,7 @@
 from ..types.multiagent import MultiAgentInput
 from ..types.session import decode_bytes_values, encode_bytes_values
 from ..types.traces import AttributeValue
-from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
+from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status, _parse_metrics, _parse_usage
 
 logger = logging.getLogger(__name__)
 
@@ -1275,6 +1275,10 @@ def serialize_state(self) -> dict[str, Any]:
             "next_nodes_to_execute": next_nodes,
             "current_task": encode_bytes_values(self.state.task),
             "execution_order": [n.node_id for n in self.state.execution_order],
+            "accumulated_usage": self.state.accumulated_usage,
+            "accumulated_metrics": self.state.accumulated_metrics,
+            "execution_count": self.state.execution_count,
+            "execution_time": self.state.execution_time,
             "_internal_state": {
                 "interrupt_state": self._interrupt_state.to_dict(),
             },
@@ -1415,6 +1419,13 @@ def _from_dict(self, payload: dict[str, Any]) -> None:
         # Task
         self.state.task = decode_bytes_values(payload.get("current_task", self.state.task))
 
+        # Cumulative accounting: restore so the timeout budget (should_continue) and the reported
+        # totals stay correct across a resume, rather than resetting to zero.
+        self.state.accumulated_usage = _parse_usage(payload.get("accumulated_usage") or {})
+        self.state.accumulated_metrics = _parse_metrics(payload.get("accumulated_metrics") or {})
+        self.state.execution_count = int(payload.get("execution_count") or 0)
+        self.state.execution_time = int(payload.get("execution_time") or 0)
+
         # next nodes to execute
         next_nodes = [self.nodes[nid] for nid in (payload.get("next_nodes_to_execute") or []) if nid in self.nodes]
         self._resume_next_nodes = next_nodes
diff --git a/strands-py/src/strands/multiagent/swarm.py b/strands-py/src/strands/multiagent/swarm.py
index 28711809d3..4eb6892dda 100644
--- a/strands-py/src/strands/multiagent/swarm.py
+++ b/strands-py/src/strands/multiagent/swarm.py
@@ -56,7 +56,7 @@
 from ..types.multiagent import MultiAgentInput
 from ..types.session import decode_bytes_values, encode_bytes_values
 from ..types.traces import AttributeValue
-from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status
+from .base import MultiAgentBase, MultiAgentResult, NodeResult, Status, _parse_metrics, _parse_usage
 
 logger = logging.getLogger(__name__)
 
@@ -993,6 +993,9 @@ def serialize_state(self) -> dict[str, Any]:
             "node_results": {k: v.to_dict() for k, v in self.state.results.items()},
             "next_nodes_to_execute": next_nodes,
             "current_task": encode_bytes_values(self.state.task),
+            "accumulated_usage": self.state.accumulated_usage,
+            "accumulated_metrics": self.state.accumulated_metrics,
+            "execution_time": self.state.execution_time,
             "context": {
                 "shared_context": getattr(self.state.shared_context, "context", {}) or {},
                 "handoff_node": self.state.handoff_node.node_id if self.state.handoff_node else None,
@@ -1036,12 +1039,20 @@ def deserialize_state(self, payload: dict[str, Any]) -> None:
 
     def _from_dict(self, payload: dict[str, Any]) -> None:
         self.state.completion_status = Status(payload["status"])
+        # Point the state's shared context at the swarm-owned object, matching the identity a fresh run
+        # establishes in stream_async, so the serialize path (reads self.state.shared_context) and the
+        # node-input builder (reads self.shared_context) stay in sync after resume.
+        self.state.shared_context = self.shared_context
         # Hydrate completed nodes & results
         context = payload["context"] or {}
         self.shared_context.context = context.get("shared_context") or {}
         self.state.handoff_message = context.get("handoff_message")
         self.state.handoff_node = self.nodes[context["handoff_node"]] if context.get("handoff_node") else None
 
+        self.state.accumulated_usage = _parse_usage(payload.get("accumulated_usage") or {})
+        self.state.accumulated_metrics = _parse_metrics(payload.get("accumulated_metrics") or {})
+        self.state.execution_time = int(payload.get("execution_time") or 0)
+
         self.state.node_history = [self.nodes[nid] for nid in (payload.get("node_history") or []) if nid in self.nodes]
 
         raw_results = payload.get("node_results") or {}
diff --git a/strands-py/tests/strands/multiagent/test_graph.py b/strands-py/tests/strands/multiagent/test_graph.py
index 22eae11bba..fd227dfe9b 100644
--- a/strands-py/tests/strands/multiagent/test_graph.py
+++ b/strands-py/tests/strands/multiagent/test_graph.py
@@ -2046,6 +2046,55 @@ async def test_graph_persisted(mock_strands_tracer, mock_use_span):
     assert "test_node" in final_state["node_results"]
 
 
+def test_graph_serialize_deserialize_serialize_preserves_cumulative_state():
+    """Resuming then re-serializing must not lose cumulative accounting.
+
+    Guards the serialize -> deserialize -> serialize round-trip on the resume path: the second
+    serialize must equal the first for accumulated_usage / accumulated_metrics / execution_count /
+    execution_time. Before the fix these were never persisted, so a resumed graph reset them to zero,
+    corrupting the timeout budget (should_continue) and the totals reported in GraphResult.
+    """
+    builder = GraphBuilder()
+    builder.add_node(create_mock_agent("test_agent"), "test_node")
+    builder.set_entry_point("test_node")
+    graph = builder.build()
+
+    payload = {
+        "type": "graph",
+        "id": "default_graph",
+        "status": "executing",
+        "completed_nodes": [],
+        "failed_nodes": [],
+        "interrupted_nodes": [],
+        "node_results": {},
+        "next_nodes_to_execute": ["test_node"],
+        "current_task": "resume me",
+        "execution_order": [],
+        "accumulated_usage": {"inputTokens": 11, "outputTokens": 22, "totalTokens": 33},
+        "accumulated_metrics": {"latencyMs": 44},
+        "execution_count": 3,
+        "execution_time": 555,
+        "_internal_state": {"interrupt_state": {"activated": False, "context": {}, "interrupts": {}}},
+    }
+
+    graph.deserialize_state(payload)
+
+    # Cumulative accounting is restored, not reset to zero.
+    assert graph.state.accumulated_usage == {"inputTokens": 11, "outputTokens": 22, "totalTokens": 33}
+    assert graph.state.accumulated_metrics == {"latencyMs": 44}
+    assert graph.state.execution_count == 3
+    assert graph.state.execution_time == 555
+
+    serialize1 = graph.serialize_state()
+    graph.deserialize_state(serialize1)
+    serialize2 = graph.serialize_state()
+
+    assert serialize2["accumulated_usage"] == serialize1["accumulated_usage"]
+    assert serialize2["accumulated_metrics"] == serialize1["accumulated_metrics"]
+    assert serialize2["execution_count"] == serialize1["execution_count"]
+    assert serialize2["execution_time"] == serialize1["execution_time"]
+
+
 @pytest.mark.parametrize(
     ("cancel_node", "cancel_message"),
     [(True, "node cancelled by user"), ("custom cancel message", "custom cancel message")],
diff --git a/strands-py/tests/strands/multiagent/test_swarm.py b/strands-py/tests/strands/multiagent/test_swarm.py
index cb0414b42f..5f930fd604 100644
--- a/strands-py/tests/strands/multiagent/test_swarm.py
+++ b/strands-py/tests/strands/multiagent/test_swarm.py
@@ -1212,6 +1212,60 @@ async def test_swarm_persistence(mock_strands_tracer, mock_use_span):
     assert "test_agent" in final_state["node_results"]
 
 
+def test_swarm_serialize_deserialize_serialize_preserves_state():
+    """Resuming then re-serializing must not lose shared context or cumulative accounting.
+
+    Guards the serialize -> deserialize -> serialize round-trip on the resume path: the second
+    serialize must equal the first for shared_context and the cumulative fields. Before the fix,
+    deserialize restored shared_context onto self.shared_context while serialize read
+    self.state.shared_context, so the next serialize emitted {}; and execution_time /
+    accumulated_usage / accumulated_metrics were never persisted, resetting to zero on resume.
+    """
+    agent = create_mock_agent("first")
+    swarm = Swarm([agent])
+
+    payload = {
+        "type": "swarm",
+        "id": "default_swarm",
+        "status": "executing",
+        "node_history": [],
+        "node_results": {},
+        "next_nodes_to_execute": ["first"],
+        "current_task": "resume me",
+        "accumulated_usage": {"inputTokens": 11, "outputTokens": 22, "totalTokens": 33},
+        "accumulated_metrics": {"latencyMs": 44},
+        "execution_time": 555,
+        "context": {
+            "shared_context": {"first": {"fact": "persist-me"}},
+            "handoff_node": None,
+            "handoff_message": None,
+        },
+        "_internal_state": {"interrupt_state": {"activated": False, "context": {}, "interrupts": {}}},
+    }
+
+    swarm.deserialize_state(payload)
+
+    # The swarm-owned and state-owned shared contexts must be the same restored object.
+    assert swarm.shared_context.context == {"first": {"fact": "persist-me"}}
+    assert swarm.state.shared_context.context == {"first": {"fact": "persist-me"}}
+    assert swarm.state.shared_context is swarm.shared_context
+
+    # Cumulative accounting is restored, not reset to zero.
+    assert swarm.state.accumulated_usage == {"inputTokens": 11, "outputTokens": 22, "totalTokens": 33}
+    assert swarm.state.accumulated_metrics == {"latencyMs": 44}
+    assert swarm.state.execution_time == 555
+
+    serialize1 = swarm.serialize_state()
+    swarm.deserialize_state(serialize1)
+    serialize2 = swarm.serialize_state()
+
+    assert serialize2["context"]["shared_context"] == serialize1["context"]["shared_context"]
+    assert serialize2["context"]["shared_context"] == {"first": {"fact": "persist-me"}}
+    assert serialize2["accumulated_usage"] == serialize1["accumulated_usage"]
+    assert serialize2["accumulated_metrics"] == serialize1["accumulated_metrics"]
+    assert serialize2["execution_time"] == serialize1["execution_time"]
+
+
 @pytest.mark.asyncio
 async def test_swarm_handle_handoff():
     first_agent = create_mock_agent("first")

From ac77e2501b701fbfedbe9de3ab55504d3e53f542 Mon Sep 17 00:00:00 2001
From: opieter-aws <opieter@amazon.com>
Date: Tue, 21 Jul 2026 14:32:18 -0400
Subject: [PATCH 2/5] adress comments

---
 strands-py/src/strands/multiagent/swarm.py        |  8 ++++----
 strands-py/tests/strands/multiagent/test_graph.py |  9 ++++-----
 strands-py/tests/strands/multiagent/test_swarm.py | 10 ++++------
 3 files changed, 12 insertions(+), 15 deletions(-)

diff --git a/strands-py/src/strands/multiagent/swarm.py b/strands-py/src/strands/multiagent/swarm.py
index 4eb6892dda..45995d6894 100644
--- a/strands-py/src/strands/multiagent/swarm.py
+++ b/strands-py/src/strands/multiagent/swarm.py
@@ -1010,14 +1010,14 @@ def deserialize_state(self, payload: dict[str, Any]) -> None:
         """Restore swarm state from a session dict and prepare for execution.
 
         This method handles two scenarios:
-        1. If the persisted status is COMPLETED, FAILED resets all nodes and graph state
-           to allow re-execution from the beginning.
+        1. If the payload omits next_nodes_to_execute (a terminal or fresh state), resets all
+           nodes and swarm state to allow re-execution from the beginning.
         2. Otherwise, restores the persisted state and prepares to resume execution
-           from the next ready nodes.
+           from the next node.
 
         Args:
             payload: Dictionary containing persisted state data including status,
-                    completed nodes, results, and next nodes to execute.
+                    node history, results, and next nodes to execute.
         """
         if "_internal_state" in payload:
             internal_state = payload["_internal_state"]
diff --git a/strands-py/tests/strands/multiagent/test_graph.py b/strands-py/tests/strands/multiagent/test_graph.py
index fd227dfe9b..8e429c8bf0 100644
--- a/strands-py/tests/strands/multiagent/test_graph.py
+++ b/strands-py/tests/strands/multiagent/test_graph.py
@@ -2047,12 +2047,11 @@ async def test_graph_persisted(mock_strands_tracer, mock_use_span):
 
 
 def test_graph_serialize_deserialize_serialize_preserves_cumulative_state():
-    """Resuming then re-serializing must not lose cumulative accounting.
+    """serialize -> deserialize -> serialize is value-preserving on the resume path.
 
-    Guards the serialize -> deserialize -> serialize round-trip on the resume path: the second
-    serialize must equal the first for accumulated_usage / accumulated_metrics / execution_count /
-    execution_time. Before the fix these were never persisted, so a resumed graph reset them to zero,
-    corrupting the timeout budget (should_continue) and the totals reported in GraphResult.
+    Guarantees that a resumed graph re-serializes the same cumulative accounting (accumulated_usage /
+    accumulated_metrics / execution_count / execution_time) it was restored with, so the timeout
+    budget (should_continue) and the totals reported in GraphResult reflect the whole run.
     """
     builder = GraphBuilder()
     builder.add_node(create_mock_agent("test_agent"), "test_node")
diff --git a/strands-py/tests/strands/multiagent/test_swarm.py b/strands-py/tests/strands/multiagent/test_swarm.py
index 5f930fd604..e2d9abbfe3 100644
--- a/strands-py/tests/strands/multiagent/test_swarm.py
+++ b/strands-py/tests/strands/multiagent/test_swarm.py
@@ -1213,13 +1213,11 @@ async def test_swarm_persistence(mock_strands_tracer, mock_use_span):
 
 
 def test_swarm_serialize_deserialize_serialize_preserves_state():
-    """Resuming then re-serializing must not lose shared context or cumulative accounting.
+    """serialize -> deserialize -> serialize is value-preserving on the resume path.
 
-    Guards the serialize -> deserialize -> serialize round-trip on the resume path: the second
-    serialize must equal the first for shared_context and the cumulative fields. Before the fix,
-    deserialize restored shared_context onto self.shared_context while serialize read
-    self.state.shared_context, so the next serialize emitted {}; and execution_time /
-    accumulated_usage / accumulated_metrics were never persisted, resetting to zero on resume.
+    Guarantees that a resumed swarm re-serializes the same shared_context and cumulative
+    accounting (execution_time / accumulated_usage / accumulated_metrics) it was restored with,
+    and that the restored SwarmState and Swarm share the same shared_context object.
     """
     agent = create_mock_agent("first")
     swarm = Swarm([agent])

From ef85b239f21c019ba16fdba1fb6720f25365bc85 Mon Sep 17 00:00:00 2001
From: opieter-aws <opieter@amazon.com>
Date: Wed, 22 Jul 2026 12:20:21 -0400
Subject: [PATCH 3/5] Adress comments

---
 strands-py/src/strands/multiagent/base.py     | 38 +++++++++
 strands-py/src/strands/multiagent/graph.py    |  7 +-
 strands-py/src/strands/multiagent/swarm.py    |  6 +-
 .../tests/strands/multiagent/test_graph.py    | 77 +++++++++++++++++++
 .../tests/strands/multiagent/test_swarm.py    | 21 +++++
 5 files changed, 144 insertions(+), 5 deletions(-)

diff --git a/strands-py/src/strands/multiagent/base.py b/strands-py/src/strands/multiagent/base.py
index fc2186e27e..a7083cf910 100644
--- a/strands-py/src/strands/multiagent/base.py
+++ b/strands-py/src/strands/multiagent/base.py
@@ -4,6 +4,7 @@
 """
 
 import logging
+import time
 import warnings
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Mapping
@@ -187,6 +188,43 @@ class MultiAgentBase(ABC):
     """
 
     id: str
+    # Wall-clock start of the active invocation, or None when no invocation is running. Set at
+    # invocation start; folded into the orchestrator's committed execution-time total exactly once
+    # at finalization.
+    _invocation_start_time: float | None = None
+
+    def _execution_time_with_active_interval(self, committed_time: int) -> int:
+        """Committed execution time (ms) plus the active invocation's in-flight interval.
+
+        The active interval is folded into the committed total only once, at finalization, so any
+        read that must reflect elapsed time — checkpoints, result building — adds it on here.
+
+        Args:
+            committed_time: Execution time in milliseconds already committed by prior invocations.
+
+        Returns:
+            committed_time plus the current invocation's elapsed milliseconds, or committed_time
+            unchanged when no invocation is running.
+        """
+        if self._invocation_start_time is None:
+            return committed_time
+        return committed_time + round((time.time() - self._invocation_start_time) * 1000)
+
+    def _commit_active_interval(self, committed_time: int) -> int:
+        """Fold the active invocation's interval into committed_time and end the interval.
+
+        Idempotent: with no active interval this returns committed_time unchanged, so calling it at
+        finalization never double-counts.
+
+        Args:
+            committed_time: Execution time in milliseconds already committed by prior invocations.
+
+        Returns:
+            The new committed total including the interval that just ended.
+        """
+        total = self._execution_time_with_active_interval(committed_time)
+        self._invocation_start_time = None
+        return total
 
     @abstractmethod
     async def invoke_async(
diff --git a/strands-py/src/strands/multiagent/graph.py b/strands-py/src/strands/multiagent/graph.py
index 1388d6dd18..cc1e3735ff 100644
--- a/strands-py/src/strands/multiagent/graph.py
+++ b/strands-py/src/strands/multiagent/graph.py
@@ -646,6 +646,7 @@ async def stream_async(
 
         # Initialize state
         start_time = time.time()
+        self._invocation_start_time = start_time
         if not self._resume_from_session and not self._interrupt_state.activated:
             # Initialize state
             self.state = GraphState(
@@ -697,7 +698,7 @@ async def stream_async(
                 self.state.status = Status.FAILED
                 raise
             finally:
-                self.state.execution_time += round((time.time() - start_time) * 1000)
+                self.state.execution_time = self._commit_active_interval(self.state.execution_time)
                 await self.hooks.invoke_callbacks_async(AfterMultiAgentInvocationEvent(self))
                 self._resume_from_session = False
                 self._resume_next_nodes.clear()
@@ -1249,7 +1250,7 @@ def _build_result(self, interrupts: list[Interrupt]) -> GraphResult:
             accumulated_usage=self.state.accumulated_usage,
             accumulated_metrics=self.state.accumulated_metrics,
             execution_count=self.state.execution_count,
-            execution_time=self.state.execution_time,
+            execution_time=self._execution_time_with_active_interval(self.state.execution_time),
             total_nodes=self.state.total_nodes,
             completed_nodes=len(self.state.completed_nodes),
             failed_nodes=len(self.state.failed_nodes),
@@ -1278,7 +1279,7 @@ def serialize_state(self) -> dict[str, Any]:
             "accumulated_usage": self.state.accumulated_usage,
             "accumulated_metrics": self.state.accumulated_metrics,
             "execution_count": self.state.execution_count,
-            "execution_time": self.state.execution_time,
+            "execution_time": self._execution_time_with_active_interval(self.state.execution_time),
             "_internal_state": {
                 "interrupt_state": self._interrupt_state.to_dict(),
             },
diff --git a/strands-py/src/strands/multiagent/swarm.py b/strands-py/src/strands/multiagent/swarm.py
index 45995d6894..ebee607861 100644
--- a/strands-py/src/strands/multiagent/swarm.py
+++ b/strands-py/src/strands/multiagent/swarm.py
@@ -410,6 +410,8 @@ async def stream_async(
                 shared_context=self.shared_context,
             )
 
+        self._invocation_start_time = self.state.start_time
+
         span = self.tracer.start_multiagent_span(task, "swarm", custom_trace_attributes=self.trace_attributes)
         with trace_api.use_span(span, end_on_exit=True):
             interrupts = []
@@ -435,7 +437,7 @@ async def stream_async(
                 self.state.completion_status = Status.FAILED
                 raise
             finally:
-                self.state.execution_time += round((time.time() - self.state.start_time) * 1000)
+                self.state.execution_time = self._commit_active_interval(self.state.execution_time)
                 await self.hooks.invoke_callbacks_async(AfterMultiAgentInvocationEvent(self, invocation_state))
                 self._resume_from_session = False
 
@@ -995,7 +997,7 @@ def serialize_state(self) -> dict[str, Any]:
             "current_task": encode_bytes_values(self.state.task),
             "accumulated_usage": self.state.accumulated_usage,
             "accumulated_metrics": self.state.accumulated_metrics,
-            "execution_time": self.state.execution_time,
+            "execution_time": self._execution_time_with_active_interval(self.state.execution_time),
             "context": {
                 "shared_context": getattr(self.state.shared_context, "context", {}) or {},
                 "handoff_node": self.state.handoff_node.node_id if self.state.handoff_node else None,
diff --git a/strands-py/tests/strands/multiagent/test_graph.py b/strands-py/tests/strands/multiagent/test_graph.py
index 8e429c8bf0..c3f30f7265 100644
--- a/strands-py/tests/strands/multiagent/test_graph.py
+++ b/strands-py/tests/strands/multiagent/test_graph.py
@@ -2094,6 +2094,83 @@ def test_graph_serialize_deserialize_serialize_preserves_cumulative_state():
     assert serialize2["execution_time"] == serialize1["execution_time"]
 
 
+@pytest.mark.asyncio
+async def test_graph_execution_time_reflects_active_invocation(mock_strands_tracer, mock_use_span):
+    """The final GraphResult includes the current invocation's interval on top of restored prior time.
+
+    execution_time is committed to state once, at finalization. GraphResult is built before that
+    commit, so it must fold in the in-flight interval itself — and finalization must not double-count.
+    """
+    # Monotonic fake clock advanced explicitly; robust to how many times time.time() is called.
+    clock = {"now": 1000.0}
+
+    with patch("strands.multiagent.graph.time.time", lambda: clock["now"]):
+        builder = GraphBuilder()
+        builder.add_node(create_mock_agent("test_agent"), "test_node")
+        builder.set_entry_point("test_node")
+        graph = builder.build()
+
+        # Resume from a checkpoint that already accrued 555ms in a prior invocation.
+        graph.deserialize_state(
+            {
+                "type": "graph",
+                "id": "default_graph",
+                "status": "executing",
+                "completed_nodes": [],
+                "failed_nodes": [],
+                "interrupted_nodes": [],
+                "node_results": {},
+                "next_nodes_to_execute": ["test_node"],
+                "current_task": "resume me",
+                "execution_order": [],
+                "accumulated_usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+                "accumulated_metrics": {"latencyMs": 0},
+                "execution_count": 0,
+                "execution_time": 555,
+                "_internal_state": {"interrupt_state": {"activated": False, "context": {}, "interrupts": {}}},
+            }
+        )
+
+        # The node advances the clock by 1200ms while it "runs".
+        async def advancing_stream(*args, **kwargs):
+            clock["now"] += 1.2
+            yield {"result": graph.nodes["test_node"].executor.return_value}
+
+        graph.nodes["test_node"].executor.stream_async = Mock(side_effect=advancing_stream)
+
+        result = await graph.invoke_async("resume me")
+
+    # 555 restored + 1200 in-flight = 1755ms; committed exactly once.
+    tru_result_time = result.execution_time
+    exp_result_time = 1755
+    assert tru_result_time == exp_result_time
+    assert graph.state.execution_time == exp_result_time
+
+
+@pytest.mark.asyncio
+async def test_graph_checkpoint_persists_in_flight_execution_time(mock_strands_tracer, mock_use_span):
+    """A mid-run per-node checkpoint persists elapsed time so a resumed run keeps its timeout budget.
+
+    Guards the crash-restart path: the AfterNodeCall session sync serializes before the invocation's
+    finally commits the interval, so serialize_state must fold the in-flight interval into
+    execution_time rather than persisting the stale pre-invocation value (which would reset the budget).
+    """
+    clock = {"now": 2000.4}
+
+    builder = GraphBuilder()
+    builder.add_node(create_mock_agent("test_agent"), "test_node")
+    builder.set_entry_point("test_node")
+    graph = builder.build()
+
+    with patch("strands.multiagent.graph.time.time", lambda: clock["now"]):
+        # Marker set at invocation start; a checkpoint taken 400ms in must reflect that interval.
+        graph._invocation_start_time = 2000.0
+        graph.state.status = Status.EXECUTING
+        checkpoint = graph.serialize_state()
+
+    assert checkpoint["execution_time"] == 400
+
+
 @pytest.mark.parametrize(
     ("cancel_node", "cancel_message"),
     [(True, "node cancelled by user"), ("custom cancel message", "custom cancel message")],
diff --git a/strands-py/tests/strands/multiagent/test_swarm.py b/strands-py/tests/strands/multiagent/test_swarm.py
index e2d9abbfe3..d5d49dcfcb 100644
--- a/strands-py/tests/strands/multiagent/test_swarm.py
+++ b/strands-py/tests/strands/multiagent/test_swarm.py
@@ -1264,6 +1264,27 @@ def test_swarm_serialize_deserialize_serialize_preserves_state():
     assert serialize2["execution_time"] == serialize1["execution_time"]
 
 
+def test_swarm_checkpoint_persists_in_flight_execution_time():
+    """A mid-run per-node checkpoint persists elapsed time so a resumed swarm keeps its timeout budget.
+
+    Guards the crash-restart path: the AfterNodeCall session sync serializes before the invocation's
+    finally commits the interval, so serialize_state must fold the in-flight interval into
+    execution_time rather than persisting the stale pre-invocation value (which would reset the budget).
+    """
+    clock = {"now": 3000.6}
+
+    swarm = Swarm([create_mock_agent("first")])
+
+    with patch("strands.multiagent.swarm.time.time", lambda: clock["now"]):
+        # Marker set at invocation start; a checkpoint taken 600ms in must reflect that interval.
+        swarm._invocation_start_time = 3000.0
+        swarm.state.completion_status = Status.EXECUTING
+        swarm.state.current_node = swarm.nodes["first"]
+        checkpoint = swarm.serialize_state()
+
+    assert checkpoint["execution_time"] == 600
+
+
 @pytest.mark.asyncio
 async def test_swarm_handle_handoff():
     first_agent = create_mock_agent("first")

From 6c9ca250b5f55752ad66a3bd68aa627c32ab9e0a Mon Sep 17 00:00:00 2001
From: opieter-aws <opieter@amazon.com>
Date: Wed, 22 Jul 2026 13:13:45 -0400
Subject: [PATCH 4/5] Fix timing issue

---
 strands-py/src/strands/multiagent/graph.py    |  3 ++-
 strands-py/src/strands/multiagent/swarm.py    |  4 +--
 .../tests/strands/multiagent/test_graph.py    | 26 +++++++++++++++++++
 .../tests/strands/multiagent/test_swarm.py    | 23 ++++++++++++++++
 4 files changed, 53 insertions(+), 3 deletions(-)

diff --git a/strands-py/src/strands/multiagent/graph.py b/strands-py/src/strands/multiagent/graph.py
index cc1e3735ff..f9a95cfc1e 100644
--- a/strands-py/src/strands/multiagent/graph.py
+++ b/strands-py/src/strands/multiagent/graph.py
@@ -646,7 +646,6 @@ async def stream_async(
 
         # Initialize state
         start_time = time.time()
-        self._invocation_start_time = start_time
         if not self._resume_from_session and not self._interrupt_state.activated:
             # Initialize state
             self.state = GraphState(
@@ -665,6 +664,8 @@ async def stream_async(
         with trace_api.use_span(span, end_on_exit=True):
             interrupts = []
 
+            self._invocation_start_time = start_time
+
             try:
                 logger.debug(
                     "max_node_executions=<%s>, execution_timeout=<%s>s, node_timeout=<%s>s | graph execution config",
diff --git a/strands-py/src/strands/multiagent/swarm.py b/strands-py/src/strands/multiagent/swarm.py
index ebee607861..a4e57b7bad 100644
--- a/strands-py/src/strands/multiagent/swarm.py
+++ b/strands-py/src/strands/multiagent/swarm.py
@@ -410,12 +410,12 @@ async def stream_async(
                 shared_context=self.shared_context,
             )
 
-        self._invocation_start_time = self.state.start_time
-
         span = self.tracer.start_multiagent_span(task, "swarm", custom_trace_attributes=self.trace_attributes)
         with trace_api.use_span(span, end_on_exit=True):
             interrupts = []
 
+            self._invocation_start_time = self.state.start_time
+
             try:
                 current_node = cast(SwarmNode, self.state.current_node)
                 logger.debug("current_node=<%s> | starting swarm execution with node", current_node.node_id)
diff --git a/strands-py/tests/strands/multiagent/test_graph.py b/strands-py/tests/strands/multiagent/test_graph.py
index c3f30f7265..4ede242543 100644
--- a/strands-py/tests/strands/multiagent/test_graph.py
+++ b/strands-py/tests/strands/multiagent/test_graph.py
@@ -2171,6 +2171,32 @@ async def test_graph_checkpoint_persists_in_flight_execution_time(mock_strands_t
     assert checkpoint["execution_time"] == 400
 
 
+@pytest.mark.asyncio
+async def test_graph_tracing_setup_failure_does_not_leak_timer(mock_strands_tracer, mock_use_span):
+    """A tracing setup failure must not leave the invocation timer running.
+
+    The timer starts inside the span context so its clearing finally is guaranteed to run. If span
+    setup raises before then, no interval is started, and a later serialize_state must not accrue
+    wall time against an abandoned invocation.
+    """
+    clock = {"now": 1000.0}
+    mock_strands_tracer.start_multiagent_span.side_effect = RuntimeError("span setup failed")
+
+    builder = GraphBuilder()
+    builder.add_node(create_mock_agent("test_agent"), "test_node")
+    builder.set_entry_point("test_node")
+    graph = builder.build()
+
+    with patch("strands.multiagent.graph.time.time", lambda: clock["now"]):
+        with pytest.raises(RuntimeError, match="span setup failed"):
+            await graph.invoke_async("go")
+        clock["now"] = 1000.5  # 500ms later
+        checkpoint = graph.serialize_state()
+
+    assert graph._invocation_start_time is None
+    assert checkpoint["execution_time"] == 0
+
+
 @pytest.mark.parametrize(
     ("cancel_node", "cancel_message"),
     [(True, "node cancelled by user"), ("custom cancel message", "custom cancel message")],
diff --git a/strands-py/tests/strands/multiagent/test_swarm.py b/strands-py/tests/strands/multiagent/test_swarm.py
index d5d49dcfcb..99f5feedb4 100644
--- a/strands-py/tests/strands/multiagent/test_swarm.py
+++ b/strands-py/tests/strands/multiagent/test_swarm.py
@@ -1285,6 +1285,29 @@ def test_swarm_checkpoint_persists_in_flight_execution_time():
     assert checkpoint["execution_time"] == 600
 
 
+@pytest.mark.asyncio
+async def test_swarm_tracing_setup_failure_does_not_leak_timer(mock_strands_tracer, mock_use_span):
+    """A tracing setup failure must not leave the invocation timer running.
+
+    The timer starts inside the span context so its clearing finally is guaranteed to run. If span
+    setup raises before then, no interval is started, and a later serialize_state must not accrue
+    wall time against an abandoned invocation.
+    """
+    clock = {"now": 2000.0}
+    mock_strands_tracer.start_multiagent_span.side_effect = RuntimeError("span setup failed")
+
+    swarm = Swarm([create_mock_agent("first")])
+
+    with patch("strands.multiagent.swarm.time.time", lambda: clock["now"]):
+        with pytest.raises(RuntimeError, match="span setup failed"):
+            await swarm.invoke_async("go")
+        clock["now"] = 2000.5  # 500ms later
+        checkpoint = swarm.serialize_state()
+
+    assert swarm._invocation_start_time is None
+    assert checkpoint["execution_time"] == 0
+
+
 @pytest.mark.asyncio
 async def test_swarm_handle_handoff():
     first_agent = create_mock_agent("first")

From e2f68d239e825ae26793a44175d47a00066a16e2 Mon Sep 17 00:00:00 2001
From: opieter-aws <opieter@amazon.com>
Date: Thu, 23 Jul 2026 12:00:49 -0400
Subject: [PATCH 5/5] address reviewer feedback

---
 strands-py/src/strands/multiagent/base.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/strands-py/src/strands/multiagent/base.py b/strands-py/src/strands/multiagent/base.py
index a7083cf910..67bf92d9aa 100644
--- a/strands-py/src/strands/multiagent/base.py
+++ b/strands-py/src/strands/multiagent/base.py
@@ -191,7 +191,11 @@ class MultiAgentBase(ABC):
     # Wall-clock start of the active invocation, or None when no invocation is running. Set at
     # invocation start; folded into the orchestrator's committed execution-time total exactly once
     # at finalization.
-    _invocation_start_time: float | None = None
+    _invocation_start_time: float | None
+
+    def __init__(self) -> None:
+        """Initialize base multi-agent state."""
+        self._invocation_start_time = None
 
     def _execution_time_with_active_interval(self, committed_time: int) -> int:
         """Committed execution time (ms) plus the active invocation's in-flight interval.