Merge pull request #38 from rezaho/feature/tracing-streaming

feat(framework): AG-UI event stream translator (Session 06)

Author: rezaho

docs/architecture/framework/aggui-custom-events.md

@@ -1,6 +1,6 @@
 # Framework features required for Spren v0.3
 
-Two framework features block Spren v0.3: the **NDJSON streaming tracing writer** (Session 01, ✅ shipped) and the **AG-UI event stream translator** (Session 06, scoped — required before Spren Session 04 starts). Everything else Spren v0.3 needs is already in the framework after the unified-barrier merge.
+Two framework features block Spren v0.3: the **NDJSON streaming tracing writer** (Session 01, ✅ shipped) and the **AG-UI event stream translator** (Session 06, ✅ shipped). Everything else Spren v0.3 needs is already in the framework after the unified-barrier merge.
 
 ## Outcome
 
@@ -17,7 +17,7 @@ When Spren v0.3 ships:
 | # | Title | Status | Outcome |
 |---|-------|--------|---------|
 | 01 | [NDJSON streaming tracing writer](./sessions/v0.3.0/01-ndjson-streaming-tracing-writer.md) | ✅ shipped | Replace `coordination/tracing/writers/json_writer.py` (or its current canonical path) with a streaming NDJSON writer. EventBus subscriber appends one JSON object per line per event. Reader interface returns a hierarchical TraceTree on demand. Framework regression suite green. |
-| 06 | [AG-UI event stream translator](./sessions/v0.3.0/06-aggui-translator.md) | scoped | Ship an optional adapter (`marsys.transport.aggui` or equivalent namespace) that translates `EventBus` events into AG-UI-format events. Pydantic models for each AG-UI event type pinned to a specific AG-UI version. `AGUIEventStream(orchestra, run_id) -> AsyncIterator[AGUIEvent]` interface. `schema_version: int = 1` on every event for forward compatibility. Tests against a real multi-agent run. |
+| 06 | [AG-UI event stream translator](./sessions/v0.3.0/06-aggui-translator.md) | ✅ shipped | Optional adapter at `marsys.coordination.aggui` (note: not `marsys.transport.aggui`; observability + API design docs updated separately). Translates `EventBus` events into AG-UI-format events via 22 typed mappers + 13 `marsys.*` Custom events with strict Pydantic validation + `MarsysRunState` snapshot/delta. `AGUIEventStream(translator)` consumes `orchestra.aggui_translator`. Schema-version handshake emitted as leading `Custom("marsys.aggui.handshake")`. Drop-newest backpressure with `Custom("marsys.stream.lagged")` notification. New `AssistantMessageEvent` on the EventBus unlocks `TextMessage*` triple. Optional dep `pip install 'marsys[aggui]'`. 69 new tests; integration test covers a 3-agent synthetic workflow with SDK round-trip + SSE round-trip. |
 
 ## Why these two for v0.3
 

docs/implementation/framework/sessions/v0.3.0/06-aggui-translator.md

@@ -10,6 +10,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- **AG-UI event stream translator** (Framework Session 06; `marsys.coordination.aggui`): a framework-internal adapter that subscribes to `EventBus` and emits AG-UI protocol events as an async iterator, so any UI that speaks AG-UI (Spren SSE, MARSYS Cloud, Studio, third-party clients) can render a live MARSYS run with no per-consumer translation logic. Public exports: `AGGUITranslator` (the EventBus subscriber; constructed inside `Orchestra._wire_event_bus()` so resumed sessions also produce streams), `AGUIEventStream(translator)` (async iterator), `aggui_event_to_sse(event) -> str` (thin wrapper around `ag_ui.encoder.EventEncoder`), `AGGUIConfig` (`enabled: bool = False`, `queue_max_size: int = 10000`), `MarsysRunState` (typed snapshot of branches / barriers / plans / total_steps). Optional dependency: `pip install 'marsys[aggui]'` (pulls `ag-ui-protocol==0.1.18` and `jsonpatch>=1.33`).
+- **`AssistantMessageEvent`** (`coordination/status/events.py`): new EventBus event emitted by the agent layer immediately after `model.arun()` returns (symmetric with the existing `AgentMessagesPreparedEvent` on the input side). Carries `agent_name`, `step_number`, `step_span_id`, `message_id` (ULID), `content`, `tool_calls`, and `finish_reason`. Unlocks AG-UI's `TextMessageStart`/`Content`/`End` triple — without this event, the translator would emit empty `TextMessageContent`. `TraceCollector` subscribes (14th handler) and stores assistant content via the existing content-addressed `MessageStore` pattern, attaching `output_message_ref` to the enclosing step span. The error path at `agents.py:3146` does NOT emit this event; `ErrorEvent` (mapped to `Custom("marsys.error")`) carries the failure signal.
+- **Custom event registry** with strict Pydantic validation (`coordination/aggui/custom_events.py`): 13 framework-internal Custom events — `marsys.aggui.handshake` (stream-level protocol-version handshake), `marsys.stream.lagged` (backpressure-overflow drop notification), `marsys.error` / `marsys.resource.limit` / `marsys.generation.metadata` / `marsys.branch.created` / `marsys.branch.completed` / `marsys.parallel.group` / `marsys.convergence` / `marsys.user_interaction.{pending,resolved,timeout}` / `marsys.memory.compaction`. Strict validation: every emitted Custom event is validated against its registered Pydantic model — schema drift raises immediately. Unknown Custom names raise `KeyError`.
+- **`MarsysRunState`** (`coordination/aggui/state.py`): typed Pydantic snapshot of run state, carried by AG-UI `StateSnapshot` and `StateDelta` events. v0.3 schema fields: `branches: dict[branch_id, BranchState]` (with `current_agent` updated on each `AgentStartEvent` — surfaces "which branch is running which agent" live), `barriers: dict[barrier_id, BarrierState]` (trimmed v0.3 — `barrier_id, status, rendezvous_node, group_id, successful_count, total_count`; `arrived_count`/`resolver_branch`/full candidates set deferred to a future session that emits dedicated barrier events), `plans: dict[agent_name, PlanState]`, `total_steps: int` (incremented on each `AgentCompleteEvent`). State deltas use RFC 6902 JSON Patch via the `jsonpatch>=1.33` library.
+- **Backpressure: drop-newest + lagged catch-up** (`coordination/aggui/translator.py`): bounded `asyncio.Queue(maxsize=10000)`. On overflow, the new event is dropped and `_lagged_count` increments. The next successful enqueue prefixes a `Custom("marsys.stream.lagged", value={"count": cumulative_drops})` event and resets the counter. Drop-NEWEST (NOT drop-oldest) preserves prefix coherence of AG-UI's `TextMessageStart`/`Content`/`End` ordering invariant — drop-oldest would let consumers see a `Content` for a `Start` they never received.
+- **Exhaustive event-mapping registry** (`coordination/aggui/mapping.py`): three buckets — `DISPATCH` (22 active mappers), `INTERNAL_ONLY` (`AgentMessagesPreparedEvent`, `MemoryResetEvent`), `NOT_YET_EMITTED` (`ValidationDecisionEvent`, `BranchEvent` — defined but never emitted; future PR that adds emission must move them out). The exhaustive test (`test_exhaustive_mapping.py`) walks `coordination/status/events.py` + `coordination/tracing/events.py` + `coordination/events.py` + `agents/memory.py`, discovers every `*Event` class across multiple base-class lineages (`MemoryResetEvent` does not inherit `StatusEvent`), and asserts each is in `EVENT_REGISTRY`. Adding a new event class without a disposition fails the test.
+- **Auto-generated Custom events doc** (`docs/architecture/framework/aggui-custom-events.md`): JSON Schemas for every entry in `CUSTOM_EVENT_REGISTRY`, generated by `packages/framework/scripts/generate_aggui_custom_events_doc.py` from the Pydantic source. CI test (`tests/coordination/aggui/test_doc_generation.py`) re-runs the generator and diffs against the checked-in markdown — fails on drift. Single source of truth: the Pydantic models.
+- **`ExecutionConfig.aggui: AGGUIConfig`** field (`coordination/config.py`): default off. When `enabled=True`, `Orchestra._wire_event_bus()` constructs the translator as a peer subscriber to `TraceCollector`. Wired in `_wire_event_bus()` (not just `_initialize_components()`) so resume sessions (Framework Session 03's pause/resume) also produce AG-UI streams.
+
 - **Pause/resume snapshot API on `Orchestra`** (Framework Session 03; ADR-007): `Orchestra.pause_session(session_id) -> None`, `Orchestra.resume_session(session_id) -> OrchestraResult`, `Orchestra.list_paused_sessions() -> list[PausedSessionMetadata]`, `Orchestra.discard_paused_session(session_id) -> None`. The pause path quiesces the running orchestrator at a tick boundary, takes a deep-copy snapshot, and writes it atomically via the configured `StorageBackend`. Resume reads the snapshot, verifies `framework_version` (exact-string match; mismatch raises `IncompatibleSnapshotError`), reconstructs a fresh `Orchestrator`, and continues dispatch through to terminal state.
 - **`StateSnapshot` Pydantic model** (`coordination/state/snapshot.py`): on-disk wire shape for paused workflows. Mirrors every field of the live `Orchestrator` mutable state (`branches`, `barriers`, `convergence_barriers`, `runnable`, `_fire_queue`, `root_barrier_id`, `_workflow_error`, `_completed_emitted`, `_user_interactions`, `_user_interaction_inflight`) plus `framework_version`, `session_id`, `topology_digest`, `created_at`, `paused_at`. JSON-encoded; the schema is exposed via `StateSnapshot.model_json_schema()`.
 - **`StorageBackend` Protocol + `FileStorageBackend`** (`coordination/state/storage.py`): generic snapshot persistence abstraction. The Protocol method set is `read(key) -> bytes`, `write(key, bytes)` (atomic), `delete(key)`, `list_with_metadata() -> list[StorageEntry]`, `expire_older_than(timedelta) -> int`. `FileStorageBackend` ships in the framework with atomic-write semantics: write-temp + `fsync(fd)` + `os.replace` + `fsync(parent_dir_fd)` (POSIX; the parent-dir fsync is a no-op on Windows but `os.replace` itself is atomic via `MoveFileEx` semantics). Cloud / CI backends will satisfy the Protocol structurally — no inheritance from a framework class.
@@ -18,6 +27,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - **`Orchestra.__init__(storage_backend=, snapshot_retention=)`** kwargs: `storage_backend` accepts any `StorageBackend` instance; default constructs a `FileStorageBackend` rooted at `~/.marsys/runs` (or `MARSYS_DATA_DIR/runs` if the env var is set). `snapshot_retention` defaults to `timedelta(days=30)` and drives a one-shot retention sweeper that runs at construction time on the active event loop.
 - **`Orchestra._active_orchestrators: dict[str, Orchestrator]`** field: populated when `Orchestra.execute()` constructs a new orchestrator and popped in the `finally` block. Lookup mechanism for `pause_session(session_id)`.
 - **Snapshot-layer exceptions** (`coordination/state/errors.py`): `IncompatibleSnapshotError`, `SnapshotCorruptionError`, `SnapshotNotFoundError`, `SnapshotSerializationError`. Re-exports `SessionNotFoundError` and `StateError` from the framework's exception hierarchy for callers using `from marsys.coordination.state import ...`.
+- **`WorkflowDefinition` canonical wire shape** (`coordination/topology/serialize.py`, `coordination/serialize.py`, `agents/serialize.py`, `models/serialize.py`, `coordination/topology/exceptions.py`): Pydantic round-trip serializer for the runtime topology + agents + execution config. Spec models (`WorkflowDefinition` envelope; `TopologySpec`, `NodeSpec`, `EdgeSpec`, `PatternConfigSpec` for topology; `AgentSpec`; `ExecutionConfigSpec` + nested `ConvergencePolicyConfigSpec`/`TracingConfigSpec`/`StatusConfigSpec`; `ModelConfigSpec`) mirror the runtime dataclasses one-for-one. Two pure conversion functions (`workflow_to_pydantic(orchestra, topology) → WorkflowDefinition`, `pydantic_to_topology(spec, tool_registry) → Topology`) round-trip between the spec and the runtime objects. Round-trip is property-tested with Hypothesis (≥100 examples) over every `NodeType` × `EdgeType` × `EdgePattern` × `PatternType` cell. Pattern provenance survives via `Topology.metadata["original_pattern"]` (a `PatternConfigConverter.convert` one-line addition). Tools are referenced by name on the wire (`AgentSpec.tools: list[str]`); the runtime callables are resolved from a caller-supplied `tool_registry`, with `UnknownToolError` on misses. `WorkflowDefinition` runs a `model_validator(mode="after")` cross-reference validator catching dangling `agent_ref` and edge endpoints at storage time.
+- **`ModelConfigSpec`** (`models/serialize.py`): storage-boundary mirror of `marsys.models.ModelConfig` with no `api_key` field and no env-resolving validator. Workflow definitions never persist secrets; consumers materialize a runnable `ModelConfig` at execution time via `runtime_model_config_from_spec(spec, api_key=...)`. `model_config_spec_from_runtime` is the round-trip helper that drops `api_key` deterministically.
+- **`topology_equals(a, b) -> bool`** (`coordination/topology/serialize.py`): semantic-equality helper that compares topologies as multisets over `(source, target, edge_type, bidirectional, pattern, metadata)` plus the topology-level metadata dict. Required because `Edge.__eq__` only compares `(source, target, edge_type)` and silently ignores `bidirectional`, `pattern`, `metadata`. Use this in round-trip tests and downstream diff UIs.
+- **`workflow_definition_schema()`** (`coordination/topology/serialize.py`): returns the canonical JSON Schema (dialect JSON Schema 2020-12) for non-Python consumers (`ajv`, `jsonschema`, `datamodel-code-generator`, MARSYS Cloud, MARSYS Studio, CI integrations). A custom `GenerateJsonSchema` subclass injects the `$schema` URI into the output; a fail-fast assertion catches silent Pydantic dialect-default drift.
+- **`UnknownToolError`** / **`NonSerializableTopologyError`** (`coordination/topology/exceptions.py`): hard-failure exceptions for the serializer. `UnknownToolError` carries the offending tool name + agent name and points callers at the `tool_registry` parameter. `NonSerializableTopologyError` fires when `workflow_to_pydantic` meets a `DeterministicNode` (det-nodes carry execution-runtime state the wire shape does not capture).
 - **`TelemetrySink` ABC** (`packages/framework/src/marsys/coordination/tracing/sink.py`): generic seam for forwarding closed spans to external observability backends (Spren daemon, LangSmith, Phoenix, Langfuse, MARSYS Cloud, custom HTTP). Two abstract async methods: `publish_span(span)` called once per span close, `close()` called once at run end. Adapters live outside the framework as third-party packages and translate the framework's `Span` shape to whatever vendor API they target. Errors are caught + logged at the `TraceCollector` boundary; one bad sink does not stop the run or block other sinks.
 - **`SecretRedactor`** (`coordination/tracing/redactor.py`): scrubs known-secret keys from span attribute payloads at the `TraceCollector._stream_span` chokepoint. Default deny-list (case-insensitive, word-boundary match): `api_key, apikey, token, authorization, auth, secret, password, bearer, cookie, session, credential`. Word boundaries treat `_`/`-`/non-alphanumeric chars as separators so `auth_token` redacts but `prompt_tokens` (an LLM token-count metric) does not. Walks `span.attributes`, every `event['attributes']` dict in `span.events`, every `link['attributes']` dict in `span.links`. Mutates in place — all consumers (NDJSON writer, vendor sinks, in-memory `TraceTree`) see the same redacted view. `NoRedactionRedactor` opt-out variant for callers that explicitly accept the leak risk.
 - **`TracingConfig.sinks: list = []`** and **`TracingConfig.redactor: SecretRedactor | None = None`** fields. Sinks register alongside the default `NDJSONTraceWriter`; if `redactor` is None the default `SecretRedactor()` instantiates lazily inside `TraceCollector`. Zero new kwargs on `Orchestra` — sinks plumb through `ExecutionConfig.tracing`, eliminating any TRUNK-CRITICAL surface change.

docs/implementation/framework/sessions/v0.3.0/06-aggui-translator/acceptance.md

@@ -64,6 +64,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+aggui = [
+    "ag-ui-protocol==0.1.18",
+    "jsonpatch>=1.33",
+]
+
 local-models = [
     "torch>=2.6.0",
     "torchvision>=0.23.0",
@@ -114,6 +119,8 @@ test = [
     "pytest>=8.4.1",
     "pytest-asyncio>=1.0.0",
     "pytest-mock>=3.14.1",
+    "hypothesis>=6.0.0",
+    "jsonschema>=4.23.0",
 ]
 
 [project.urls]

docs/implementation/framework/v0.3-spren-support.md

@@ -0,0 +1,91 @@
+"""Generate ``docs/architecture/framework/aggui-custom-events.md`` from the
+Pydantic models in ``coordination.aggui.custom_events``.
+
+Single source of truth: ``CUSTOM_EVENT_REGISTRY``. The generated markdown is
+checked into the repo and validated against re-generation in CI by
+``tests/coordination/aggui/test_doc_generation.py``.
+
+Run from the repo root:
+
+.. code-block:: bash
+
+    python packages/framework/scripts/generate_aggui_custom_events_doc.py
+
+Or directly with the package on PYTHONPATH:
+
+.. code-block:: bash
+
+    python -m scripts.generate_aggui_custom_events_doc
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+DOC_PATH = REPO_ROOT / "docs" / "architecture" / "framework" / "aggui-custom-events.md"
+
+# Ensure package is importable when run as a script
+sys.path.insert(0, str(REPO_ROOT / "packages" / "framework" / "src"))
+
+
+def render() -> str:
+    from marsys.coordination.aggui.custom_events import CUSTOM_EVENT_REGISTRY
+
+    lines: list[str] = []
+    lines.append("# AG-UI `Custom` Events — `marsys.*`")
+    lines.append("")
+    lines.append(
+        "**AUTO-GENERATED — DO NOT EDIT BY HAND.** "
+        "Regenerate via "
+        "`python packages/framework/scripts/generate_aggui_custom_events_doc.py`. "
+        "Source: `packages/framework/src/marsys/coordination/aggui/custom_events.py`."
+    )
+    lines.append("")
+    lines.append(
+        "AG-UI's `Custom` event is the documented escape hatch for protocol-internal "
+        "events that don't fit the standard lifecycle. MARSYS uses it for "
+        "framework-specific lifecycle signals (branch / parallel-group / convergence / "
+        "error / resource-limit / user-interaction / memory-compaction) and for the "
+        "stream-level handshake."
+    )
+    lines.append("")
+    lines.append(
+        "Every event below is validated at emission time against the Pydantic model "
+        "in `coordination/aggui/custom_events.py`. Validation failure raises — "
+        "catches schema drift at the source. Consumers that want lenient parsing "
+        "wrap the iterator in `try/except` at their boundary."
+    )
+    lines.append("")
+    lines.append("---")
+    lines.append("")
+    for name in sorted(CUSTOM_EVENT_REGISTRY.keys()):
+        model = CUSTOM_EVENT_REGISTRY[name]
+        schema = model.model_json_schema()
+        lines.append(f"## `{name}`")
+        lines.append("")
+        if model.__doc__:
+            doc = " ".join(line.strip() for line in model.__doc__.splitlines() if line.strip())
+            lines.append(doc)
+            lines.append("")
+        lines.append("### JSON Schema")
+        lines.append("")
+        lines.append("```json")
+        lines.append(json.dumps(schema, indent=2, sort_keys=True))
+        lines.append("```")
+        lines.append("")
+    return "\n".join(lines) + "\n"
+
+
+def main() -> int:
+    content = render()
+    DOC_PATH.parent.mkdir(parents=True, exist_ok=True)
+    DOC_PATH.write_text(content, encoding="utf-8")
+    print(f"Wrote {DOC_PATH}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())