fix(langgraph): stop discarding a subagent's explicit manually_emit_state

Follows the protocol correction in the parent commit. An audit revision added a
subagent guard to the `manually_emit_state` handler that DROPPED the payload --
neither recording nor emitting it -- justified by the rule that only the parent
agent owns state. That rule was invented, and the guard silently discarded a
state write the caller had explicitly asked for.

Pre-audit behaviour is restored: the value is recorded and the snapshot emitted,
as it was before (verified against 3386757, where this handler had no subagent
guard at all). The dispatch chokepoint stamps subagent_run_id on the snapshot,
which is provenance -- it records who wrote the state, while the state itself
stays run-scoped.

The two tests pinning the drop are replaced by one asserting the restored
behaviour: the snapshot is emitted, carries the subagent's id, and the value is
recorded. Both directions are covered, since silently dropping and silently
mis-scoping are different failures.

The other two state suppressions are LEFT ALONE, because they are pre-existing
and their real reason still holds: node-exit and checkpoint snapshots carry a
subgraph's PARTIAL view of the run's state document, so emitting one
mid-delegation would overwrite the whole state with a fragment. Their comments
are corrected only to stop citing a protocol prohibition that does not exist --
this is a producer-side choice about the integration's state model, and the
distinction matters for anyone reading the code to learn what the protocol
requires.

Verified: 468 LangGraph tests (down one, as two tests became one).

Author: mme

integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -200,9 +200,10 @@ def derive_subagent_context(
 # has to reconstruct attribution from a prior event's messageId/toolCallId).
 # Excludes only the run-lifecycle events, MESSAGES_SNAPSHOT (its messages are
 # tagged individually), and the SUBAGENT_* lifecycle events (which carry their
-# own subagent_run_id explicitly). STATE_SNAPSHOT/STATE_DELTA are never emitted for a
-# subagent (see the state-suppression in the stream loop), so they stay here only
-# for the parent's benefit.
+# own subagent_run_id explicitly). STATE_SNAPSHOT/STATE_DELTA are attributable per
+# the protocol design; this integration emits them for a subagent only on the
+# explicit `manually_emit_state` path, since node-exit and checkpoint snapshots
+# would carry a partial subgraph view (see the suppressions in the stream loop).
 _SUBAGENT_ATTRIBUTABLE_EVENT_TYPES = frozenset({
     EventType.TEXT_MESSAGE_START, EventType.TEXT_MESSAGE_CHUNK,
     EventType.TEXT_MESSAGE_CONTENT, EventType.TEXT_MESSAGE_END,
@@ -1009,9 +1010,13 @@ async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[Pr
                             # the forthcoming state update.
                             self.active_run["state_reliable"] = False
                     elif self.active_run.get("current_subagent_run_id"):
-                        # Subagents don't emit STATE_SNAPSHOT — only the parent
-                        # agent's state is surfaced. The subagent's messages still
-                        # reach the client via MESSAGES_SNAPSHOT, so nothing is lost.
+                        # Node-exit snapshots are suppressed while a subagent is
+                        # active because a subgraph's state is a PARTIAL view of the
+                        # run's document -- emitting it would overwrite the whole
+                        # state with a fragment. This is a producer-side choice about
+                        # this integration's state model, not a protocol rule:
+                        # attributed state events are legal. The subagent's messages
+                        # still reach the client via MESSAGES_SNAPSHOT.
                         # The lane survives a close (see reconcile_subagents), so a
                         # trailing event from a finished subagent stays suppressed.
                         pass
@@ -2283,32 +2288,22 @@ def _chunk_get(c: Any, key: str, default: Any = None) -> Any:
                 )
 
             elif event["name"] == CustomEventNames.ManuallyEmitState:
-                # Only the parent owns state, so this is suppressed inside a
-                # subagent exactly like the node-exit and checkpoint snapshots
-                # are. Without the guard the dispatch chokepoint would stamp the
-                # subagent's id onto a STATE_SNAPSHOT (STATE_* is in
-                # _SUBAGENT_ATTRIBUTABLE_EVENT_TYPES), and the client applies
-                # STATE_SNAPSHOT to the shared state without consulting
-                # subagent_run_id — so a subagent's partial state would land as if
-                # the parent had sent it.
+                # No subagent guard here, deliberately. An audit revision dropped
+                # this payload when a subagent was active, on the grounds that only
+                # the parent owns state. That was an invented rule -- the protocol
+                # lists STATE_* as attributable -- and dropping silently discarded
+                # a state write the user had explicitly asked for.
                 #
-                # The payload is DROPPED, not merely left unemitted. Recording it
-                # would defer the same violation by one event rather than prevent
-                # it: `manually_emitted_state` is run-global, and the stream loop
-                # reads it back as `updated_state` at the next node exit and emits
-                # it as a snapshot. That deferred snapshot carries no subagent_run_id,
-                # so it would reach the consumer looking like the parent's own
-                # state — the exact outcome this guard exists to stop.
-                if not self.active_run.get("current_subagent_run_id"):
-                    self.active_run["manually_emitted_state"] = event["data"]
-                    yield self._dispatch_event(
-                        StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(self.active_run["manually_emitted_state"]), raw_event=event)
-                    )
-                else:
-                    logger.debug(
-                        "Dropping manually_emit_state from subagent %s; only the parent owns state",
-                        self.active_run.get("current_subagent_run_id"),
-                    )
+                # `manually_emit_state` is an explicit request to set the run's
+                # state, and `manually_emitted_state` is run-global by design. A
+                # subagent making that request means it, so it is recorded and
+                # emitted like any other. The dispatch chokepoint stamps
+                # subagent_run_id on the snapshot, which is provenance: it records
+                # who wrote the state, while the state itself stays run-scoped.
+                self.active_run["manually_emitted_state"] = event["data"]
+                yield self._dispatch_event(
+                    StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(self.active_run["manually_emitted_state"]), raw_event=event)
+                )
 
             yield self._dispatch_event(
                 CustomEvent(type=EventType.CUSTOM, name=event["name"], value=event["data"], raw_event=event)
@@ -2771,8 +2766,10 @@ async def get_state_and_messages_snapshots(self, config: RunnableConfig) -> Asyn
             state_values = {}
         else:
             state_values = state.values
-        # Only the parent agent emits STATE_SNAPSHOT; while a subagent is active
-        # its (subgraph) state is not surfaced. The MESSAGES_SNAPSHOT below is
+        # Checkpoint snapshots are suppressed while a subagent is active for the
+        # same reason as the node-exit ones: a subgraph's state is a partial view of
+        # the run's document. A producer-side choice, not a protocol rule. The
+        # MESSAGES_SNAPSHOT below is
         # still emitted and carries the subagent's messages (merged + tagged), so
         # attribution and history survive without leaking subagent state.
         if not self.active_run.get("current_subagent_run_id"):

integrations/langgraph/python/tests/test_emit_events.py

@@ -117,16 +117,18 @@ async def test_manually_emit_state(self):
         assert agent.active_run["manually_emitted_state"] == {"counter": 42}
 
     @pytest.mark.asyncio
-    async def test_manually_emit_state_suppressed_inside_subagent(self):
-        """State belongs to the parent, so a subagent must not emit STATE_SNAPSHOT.
-
-        The two automatic paths (node-exit and checkpoint snapshots) already gate
-        on current_subagent_run_id. This manual path did not, and because
-        STATE_SNAPSHOT is in _SUBAGENT_ATTRIBUTABLE_EVENT_TYPES the dispatch
-        chokepoint would then stamp the subagent's id onto it — producing exactly
-        the event the design forbids. The client applies STATE_SNAPSHOT to the
-        shared state without consulting subagent_run_id, so a subagent calling the
-        manual helper would overwrite the parent's state with its own.
+    async def test_manually_emit_state_inside_subagent_is_emitted_and_attributed(self):
+        """A subagent's explicit manually_emit_state IS recorded and emitted.
+
+        An audit revision dropped this payload when a subagent was active, on the
+        grounds that only the parent owns state. That rule is not in the protocol --
+        the design lists STATE_SNAPSHOT / STATE_DELTA as attributable -- and dropping
+        silently discarded a state write the caller had explicitly requested.
+
+        `manually_emit_state` means "set the run's state", and a subagent calling it
+        means it. The snapshot goes out and the dispatch chokepoint stamps
+        subagent_run_id on it, which is provenance: it records WHO wrote the state,
+        while the state itself remains run-scoped.
         """
         agent = self._make_agent()
         agent.active_run["current_subagent_run_id"] = "tools:s1"
@@ -141,36 +143,15 @@ async def test_manually_emit_state_suppressed_inside_subagent(self):
             events.append(ev)
 
         event_types = [e.type for e in events]
-        assert EventType.STATE_SNAPSHOT not in event_types
-        # The CUSTOM passthrough still goes out, so the subagent's signal is not
-        # swallowed — only the state application is withheld.
+        assert EventType.STATE_SNAPSHOT in event_types
         assert EventType.CUSTOM in event_types
 
-    @pytest.mark.asyncio
-    async def test_manually_emit_state_inside_subagent_does_not_leak_into_parent_state(self):
-        """Suppressing the snapshot is not enough — the value must not be recorded.
-
-        Withholding the immediate STATE_SNAPSHOT while still storing the payload in
-        the run-global `manually_emitted_state` only defers the violation: the
-        stream loop reads that key back as `updated_state` on the next parent node
-        exit and emits it as an UNATTRIBUTED snapshot, so the consumer applies the
-        subagent's partial state as the parent's. The suppression has to drop the
-        value, not just delay it.
-        """
-        agent = self._make_agent()
-        agent.active_run["current_subagent_run_id"] = "tools:s1"
-        agent.active_run["active_subagents"] = {}
-        event = {
-            "event": LangGraphEventTypes.OnCustomEvent.value,
-            "name": CustomEventNames.ManuallyEmitState.value,
-            "data": {"counter": 42},
-        }
-        async for _ in agent._handle_single_event(event, {}):
-            pass
-
-        assert agent.active_run["manually_emitted_state"] is None, (
-            "a subagent's manually-emitted state must not be recorded; the stream "
-            "loop would re-emit it as the parent's state on the next node exit"
+        snapshot = next(e for e in events if e.type == EventType.STATE_SNAPSHOT)
+        assert snapshot.subagent_run_id == "tools:s1", (
+            "the snapshot must carry the subagent's id as provenance"
+        )
+        assert agent.active_run["manually_emitted_state"] == {"counter": 42}, (
+            "an explicit state write must be recorded, not discarded"
         )
 
     @pytest.mark.asyncio