diff --git a/packages/opal-backend/opal_backend/dev/main.py b/packages/opal-backend/opal_backend/dev/main.py index ed2140167a0..6790f179ac0 100644 --- a/packages/opal-backend/opal_backend/dev/main.py +++ b/packages/opal-backend/opal_backend/dev/main.py @@ -106,7 +106,7 @@ def _make_backend(token: str, origin: str = "") -> HttpBackendClient: ) # Graph session infrastructure (Heartstone). -_graph_session_store = InMemoryGraphSessionStore() +_graph_session_store = InMemoryGraphSessionStore(event_bus=_event_bus) _graph_runner = GraphRunner( store=_graph_session_store, event_bus=_event_bus, diff --git a/packages/opal-backend/opal_backend/graph_runner.py b/packages/opal-backend/opal_backend/graph_runner.py index fc98374c8a7..ccfb5169a4f 100644 --- a/packages/opal-backend/opal_backend/graph_runner.py +++ b/packages/opal-backend/opal_backend/graph_runner.py @@ -132,6 +132,18 @@ async def run_node( except Exception as exc: await self._handle_node_error(session_id, node_id, exc) + def update_auth( + self, session_id: str, + access_token: str, origin: str, + ) -> None: + """Refresh the stored OAuth token for a session. + + Called by the resume endpoint to inject a fresh token before + re-running nodes — the original token from session creation + may have expired. + """ + self._session_auth[session_id] = (access_token, origin) + async def resume_node( self, session_id: str, interaction_id: str, response: dict[str, Any], @@ -456,22 +468,36 @@ async def _get_node_type( async def _check_graph_complete( self, session_id: str, ) -> None: - """Emit graphComplete if all nodes are done.""" - if await self._store.is_graph_complete(session_id): + """Emit graphComplete or graphError if all nodes are done.""" + if not await self._store.is_graph_complete(session_id): + return + + # Check whether any node failed. + failed_errors = await self._store.get_failed_errors(session_id) + + if failed_errors: + error_msg = "; ".join(failed_errors) + await self._store.set_status(session_id, "failed") + event = { + "type": "graphError", + "sessionId": session_id, + "error": error_msg, + } + await self._store.append_event(session_id, event) + await self._event_bus.publish(session_id, event) + else: outputs = await self._store.get_graph_outputs(session_id) await self._store.set_status(session_id, "completed") - await self._store.append_event(session_id, { + event = { "type": "graphComplete", "sessionId": session_id, "outputs": outputs, - }) - await self._event_bus.publish(session_id, { - "type": "graphComplete", - "sessionId": session_id, - "outputs": outputs, - }) - # Close the event bus for this session. - await self._event_bus.close(session_id) + } + await self._store.append_event(session_id, event) + await self._event_bus.publish(session_id, event) + + # Close the event bus for this session. + await self._event_bus.close(session_id) async def _emit_node_error( self, session_id: str, node_id: str, error: str, diff --git a/packages/opal-backend/opal_backend/graph_session_store.py b/packages/opal-backend/opal_backend/graph_session_store.py index b6cd18a4e5e..2bae5f83401 100644 --- a/packages/opal-backend/opal_backend/graph_session_store.py +++ b/packages/opal-backend/opal_backend/graph_session_store.py @@ -21,10 +21,28 @@ __all__ = [ "GraphSessionStore", + "SessionSummary", "SuspendedNodeState", ] +@dataclass +class SessionSummary: + """Lightweight summary of a graph session. + + Attributes: + session_id: Unique identifier for the session. + graph_id: The graph this session belongs to. + status: Current session status (running, suspended, completed, etc.). + created_at: Unix timestamp of session creation. + """ + + session_id: str + graph_id: str + status: str + created_at: float + + @dataclass class SuspendedNodeState: """State saved when a node suspends (waiting for user input). @@ -52,13 +70,14 @@ class GraphSessionStore(Protocol): # ── Plan Storage ── async def create( - self, session_id: str, plan: GraphPlan, + self, session_id: str, plan: GraphPlan, graph_id: str, *, headless_inputs: dict[str, Any] | None = None, ) -> None: """Store plan with initial dependency counts. Args: + graph_id: Graph identifier for session indexing. headless_inputs: Optional mapping of node_id → LLMContent for headless mode. When set, input nodes auto-resolve using pre-supplied values instead of suspending. @@ -147,6 +166,12 @@ async def is_graph_complete( """True when all nodes are completed (or skipped).""" ... + async def get_failed_errors( + self, session_id: str, + ) -> list[str]: + """Return error messages from all failed nodes (empty if none).""" + ... + async def get_graph_outputs( self, session_id: str, ) -> dict[str, Any]: @@ -178,6 +203,20 @@ async def get_events( """Return events with index > after.""" ... + # ── Session Management ── + + async def list_sessions( + self, graph_id: str, + ) -> list[SessionSummary]: + """Return all sessions for a graph, sorted by creation time descending.""" + ... + + async def delete_session( + self, session_id: str, + ) -> bool: + """Delete a session and its stored state. Returns True if found.""" + ... + # ── Status ── async def get_status( diff --git a/packages/opal-backend/opal_backend/local/graph_session_router.py b/packages/opal-backend/opal_backend/local/graph_session_router.py index 6241161094b..f8dc8e77e7a 100644 --- a/packages/opal-backend/opal_backend/local/graph_session_router.py +++ b/packages/opal-backend/opal_backend/local/graph_session_router.py @@ -6,10 +6,13 @@ Parallel to ``session_router.py`` for agent sessions. This module provides the HTTP endpoints for graph execution: - POST /v1beta1/graphSessions/new → start a graph run - GET /v1beta1/graphSessions/{id} → SSE stream (replay + live) - GET /v1beta1/graphSessions/{id}/status → lightweight status - POST /v1beta1/graphSessions/{id}:cancel → cancel running tasks + POST /v1beta1/graphSessions/new → start a graph run + GET /v1beta1/graphSessions → SSE monitor (session list + status) + GET /v1beta1/graphSessions/{id} → SSE stream (replay + live) + GET /v1beta1/graphSessions/{id}/status → lightweight status + POST /v1beta1/graphSessions/{id}:cancel → cancel running tasks + POST /v1beta1/graphSessions/{id}:resume → resume suspended node + DELETE /v1beta1/graphSessions/{id} → delete session """ from __future__ import annotations @@ -26,7 +29,7 @@ from ..graph_condense import condense from ..graph_plan import create_plan from ..graph_runner import GraphRunner -from ..graph_session_store import GraphSessionStore +from ..graph_session_store import GraphSessionStore, SessionSummary from ..graph_types import GraphDescriptor from ..task_scheduler import TaskScheduler @@ -66,7 +69,7 @@ async def create_graph_session(request: Request) -> JSONResponse: """ body = await request.json() graph_data = body.get("graph") - graph_id = body.get("graphId") + graph_id = body.get("graphId", "") if not graph_data and not graph_id: return JSONResponse( @@ -78,8 +81,8 @@ async def create_graph_session(request: Request) -> JSONResponse: access_token = body.get("accessToken", "") origin = request.headers.get("origin", "") - # Load graph from Drive if graphId is provided. - if graph_id: + # Load graph from Drive if graphId is provided without inline graph. + if graph_id and not graph_data: if not access_token: return JSONResponse( {"error": "'accessToken' is required when using 'graphId'"}, @@ -119,13 +122,58 @@ async def create_graph_session(request: Request) -> JSONResponse: headless_inputs = inputs if mode == "headless" else None # Store the plan and kick off initial tasks. - await store.create(session_id, plan, headless_inputs=headless_inputs) + await store.create( + session_id, plan, graph_id, + headless_inputs=headless_inputs, + ) await runner.start_graph( session_id, access_token=access_token, origin=origin, ) return JSONResponse({"sessionId": session_id}) + @router.get("") + async def monitor_graph_sessions( + request: Request, graphId: str = "", + ) -> EventSourceResponse: + """SSE: initial session list + live status updates. + + Provides: + - Initial snapshot of all sessions for the given graph + - Live status change events (sessionStatusChange, sessionCreated, + sessionDeleted) + + Lightweight — no event replay, no output data. + """ + async def event_generator(): + # Initial snapshot: one sessionStatus event per existing session. + sessions = await store.list_sessions(graphId) + for s in sessions: + yield { + "event": "sessionStatus", + "data": json.dumps({ + "sessionId": s.session_id, + "status": s.status, + "createdAt": s.created_at, + }), + } + + # Live updates: same shape, same event name. + channel = f"graph:{graphId}" + subscription = event_bus.subscribe(channel) + try: + async for item in subscription: + if await request.is_disconnected(): + break + yield {"event": "sessionStatus", "data": json.dumps(item)} + finally: + event_bus.unsubscribe(channel, subscription) + + return EventSourceResponse( + content=event_generator(), + media_type="text/event-stream", + ) + @router.get("/{session_id}") async def stream_graph_events( request: Request, session_id: str, after: int = -1, @@ -160,6 +208,17 @@ async def event_generator(): } next_index += len(events) + # Emit a replay-complete marker so clients know stored + # events are done and live events follow. + yield { + "event": "event", + "id": str(next_index), + "data": json.dumps( + {"type": "replayComplete", "index": next_index}, + ), + } + next_index += 1 + # Phase 2: subscribe for live events if still running. current_status = await store.get_status(session_id) if current_status in ("running", "suspended"): @@ -238,6 +297,18 @@ async def resume_graph_session( response = body.get("response", {}) + # Refresh the access token so continued execution uses + # a valid credential (the original token from session + # creation may have expired during the suspend). + access_token = body.get("accessToken", "") + if not access_token: + auth_header = request.headers.get("authorization", "") + if auth_header.startswith("Bearer "): + access_token = auth_header[len("Bearer "):] + if access_token: + origin = request.headers.get("origin", "") + runner.update_auth(session_id, access_token, origin) + try: await runner.resume_node( session_id, interaction_id, response, @@ -251,4 +322,24 @@ async def resume_graph_session( "sessionId": session_id, "status": "running", }) + @router.delete("/{session_id}") + async def delete_graph_session(session_id: str) -> JSONResponse: + """Delete a graph session and all its stored state.""" + status = await store.get_status(session_id) + if status is None: + return JSONResponse( + {"error": "Graph session not found"}, status_code=404, + ) + # Cancel if still running. + if status in ("running", "suspended"): + await scheduler.cancel(session_id) + await event_bus.close(session_id) + + deleted = await store.delete_session(session_id) + if not deleted: + return JSONResponse( + {"error": "Graph session not found"}, status_code=404, + ) + return JSONResponse({"ok": True}) + return router diff --git a/packages/opal-backend/opal_backend/local/graph_session_store_impl.py b/packages/opal-backend/opal_backend/local/graph_session_store_impl.py index c21c72090ea..3abddd00ace 100644 --- a/packages/opal-backend/opal_backend/local/graph_session_store_impl.py +++ b/packages/opal-backend/opal_backend/local/graph_session_store_impl.py @@ -9,11 +9,13 @@ from __future__ import annotations +import time from dataclasses import dataclass, field from typing import Any +from ..event_bus import EventBus from ..graph_types import GraphPlan -from ..graph_session_store import SuspendedNodeState +from ..graph_session_store import SessionSummary, SuspendedNodeState __all__ = ["InMemoryGraphSessionStore"] @@ -41,6 +43,9 @@ class _SessionState: interaction_index: dict[str, str] = field(default_factory=dict) # Headless mode: node_id → pre-supplied LLMContent value. headless_inputs: dict[str, Any] | None = None + # Session management. + graph_id: str = "" + created_at: float = 0.0 class InMemoryGraphSessionStore: @@ -50,18 +55,23 @@ class InMemoryGraphSessionStore: at worst. No external dependencies. """ - def __init__(self) -> None: + def __init__(self, *, event_bus: EventBus | None = None) -> None: self._sessions: dict[str, _SessionState] = {} + self._graph_index: dict[str, list[str]] = {} + self._event_bus = event_bus # ── Plan Storage ── async def create( - self, session_id: str, plan: GraphPlan, + self, session_id: str, plan: GraphPlan, graph_id: str, *, headless_inputs: dict[str, Any] | None = None, ) -> None: """Store plan with initial dependency counts.""" - state = _SessionState(plan=plan, headless_inputs=headless_inputs) + state = _SessionState( + plan=plan, headless_inputs=headless_inputs, + graph_id=graph_id, created_at=time.time(), + ) # Compute per-node pending dep counts from plan stages. for stage in plan.stages: @@ -75,6 +85,17 @@ async def create( state.nodes[info.node.id] = ns self._sessions[session_id] = state + self._graph_index.setdefault(graph_id, []).append(session_id) + + if self._event_bus: + await self._event_bus.publish( + f"graph:{graph_id}", + { + "sessionId": session_id, + "status": state.status, + "createdAt": state.created_at, + }, + ) async def get_plan( self, session_id: str, @@ -211,6 +232,19 @@ async def is_graph_complete( for ns in state.nodes.values() ) + async def get_failed_errors( + self, session_id: str, + ) -> list[str]: + """Return error messages from all failed nodes.""" + state = self._sessions.get(session_id) + if not state: + return [] + return [ + ns.error + for ns in state.nodes.values() + if ns.status == "failed" and ns.error + ] + async def get_graph_outputs( self, session_id: str, ) -> dict[str, Any]: @@ -320,3 +354,60 @@ async def set_status( state = self._sessions.get(session_id) if state: state.status = status + if self._event_bus: + await self._event_bus.publish( + f"graph:{state.graph_id}", + { + "sessionId": session_id, + "status": status, + "createdAt": state.created_at, + }, + ) + + # ── Session Management ── + + async def list_sessions( + self, graph_id: str, + ) -> list[SessionSummary]: + """Return all sessions for a graph, sorted by creation time descending.""" + session_ids = self._graph_index.get(graph_id, []) + summaries = [] + for sid in session_ids: + state = self._sessions.get(sid) + if state: + summaries.append(SessionSummary( + session_id=sid, + graph_id=state.graph_id, + status=state.status, + created_at=state.created_at, + )) + # Sort by creation time descending (newest first). + summaries.sort(key=lambda s: s.created_at, reverse=True) + return summaries + + async def delete_session( + self, session_id: str, + ) -> bool: + """Delete a session and its stored state. Returns True if found.""" + state = self._sessions.pop(session_id, None) + if state is None: + return False + + # Remove from graph index. + ids = self._graph_index.get(state.graph_id, []) + if session_id in ids: + ids.remove(session_id) + if not ids: + del self._graph_index[state.graph_id] + + if self._event_bus: + await self._event_bus.publish( + f"graph:{state.graph_id}", + { + "sessionId": session_id, + "status": "deleted", + "createdAt": state.created_at, + }, + ) + + return True diff --git a/packages/opal-backend/tests/test_agent_integration.py b/packages/opal-backend/tests/test_agent_integration.py index 6f9983de2e4..2592e9e707a 100644 --- a/packages/opal-backend/tests/test_agent_integration.py +++ b/packages/opal-backend/tests/test_agent_integration.py @@ -82,7 +82,7 @@ async def test_agent_events_forwarded(self): runner._scheduler = scheduler plan = _agent_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") events: list[dict] = [] subscriber = bus.subscribe("s1") @@ -120,7 +120,7 @@ async def test_agent_output_flows_downstream(self): runner._scheduler = scheduler plan = _agent_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -160,7 +160,7 @@ async def test_suspend_emits_input_required(self): runner._scheduler = scheduler plan = _agent_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") events: list[dict] = [] subscriber = bus.subscribe("s1") @@ -202,7 +202,7 @@ async def test_resume_completes_graph(self): runner._scheduler = scheduler plan = _agent_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") # Start and wait for suspend. subscriber = bus.subscribe("s1") @@ -248,7 +248,7 @@ async def failing_agent(**kwargs): runner._scheduler = scheduler plan = _agent_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") events: list[dict] = [] subscriber = bus.subscribe("s1") diff --git a/packages/opal-backend/tests/test_graph_session_router.py b/packages/opal-backend/tests/test_graph_session_router.py index 515091f6ce6..d28f1eb67af 100644 --- a/packages/opal-backend/tests/test_graph_session_router.py +++ b/packages/opal-backend/tests/test_graph_session_router.py @@ -324,6 +324,123 @@ def test_resume_not_found_session(self): assert resp.status_code == 404 +class TestReplayCompleteMarker: + def test_replay_complete_in_stream(self): + app, *_ = _create_app() + client = TestClient(app) + + resp = client.post( + "/v1beta1/graphSessions/new", + json={"graph": _simple_graph()}, + ) + session_id = resp.json()["sessionId"] + + with client.stream( + "GET", f"/v1beta1/graphSessions/{session_id}", + ) as sse_resp: + raw = sse_resp.read().decode() + + events = _parse_sse_events(raw) + types = [e.get("type") for e in events if "type" in e] + assert "replayComplete" in types + + # replayComplete should come after all graph events and before done. + rc_idx = types.index("replayComplete") + gc_idx = types.index("graphComplete") + assert rc_idx > gc_idx + + def test_replay_complete_in_replay_with_after(self): + app, *_ = _create_app() + client = TestClient(app) + + resp = client.post( + "/v1beta1/graphSessions/new", + json={"graph": _simple_graph()}, + ) + session_id = resp.json()["sessionId"] + + with client.stream( + "GET", f"/v1beta1/graphSessions/{session_id}?after=2", + ) as sse_resp: + raw = sse_resp.read().decode() + + events = _parse_sse_events(raw) + types = [e.get("type") for e in events if "type" in e] + assert "replayComplete" in types + + +class TestGraphIdPassthrough: + def test_graph_id_accepted_on_creation(self): + """Verify POST /new accepts graphId and creates a valid session.""" + app, *_ = _create_app() + client = TestClient(app) + + resp = client.post( + "/v1beta1/graphSessions/new", + json={"graph": _simple_graph(), "graphId": "my-graph-123"}, + ) + assert resp.status_code == 200 + session_id = resp.json()["sessionId"] + + # Session exists and completed (proves graph_id didn't break anything). + status_resp = client.get( + f"/v1beta1/graphSessions/{session_id}/status", + ) + assert status_resp.status_code == 200 + assert status_resp.json()["status"] == "completed" + + def test_default_graph_id_works(self): + """Verify POST /new without graphId defaults cleanly.""" + app, *_ = _create_app() + client = TestClient(app) + + resp = client.post( + "/v1beta1/graphSessions/new", + json={"graph": _simple_graph()}, + ) + assert resp.status_code == 200 + + +class TestDeleteEndpoint: + def test_delete_completed_session(self): + app, store, bus = _create_app() + client = TestClient(app) + + resp = client.post( + "/v1beta1/graphSessions/new", + json={"graph": _simple_graph()}, + ) + session_id = resp.json()["sessionId"] + + # Session should exist. + status_resp = client.get( + f"/v1beta1/graphSessions/{session_id}/status", + ) + assert status_resp.status_code == 200 + + # Delete it. + del_resp = client.delete( + f"/v1beta1/graphSessions/{session_id}", + ) + assert del_resp.status_code == 200 + assert del_resp.json() == {"ok": True} + + # Should be gone. + status_resp = client.get( + f"/v1beta1/graphSessions/{session_id}/status", + ) + assert status_resp.status_code == 404 + + def test_delete_not_found(self): + app, *_ = _create_app() + client = TestClient(app) + + resp = client.delete("/v1beta1/graphSessions/nonexistent") + assert resp.status_code == 404 + + + + # ── helpers ── @@ -341,3 +458,4 @@ def _parse_sse_events(raw: str) -> list[dict]: pass return events + diff --git a/packages/opal-backend/tests/test_graph_session_store.py b/packages/opal-backend/tests/test_graph_session_store.py index 98245504521..96f07223dda 100644 --- a/packages/opal-backend/tests/test_graph_session_store.py +++ b/packages/opal-backend/tests/test_graph_session_store.py @@ -53,7 +53,7 @@ class TestCreateAndGetPlan: async def test_create_stores_plan(self): store = InMemoryGraphSessionStore() plan = _linear_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") result = await store.get_plan("s1") assert result is plan @@ -65,7 +65,7 @@ async def test_get_plan_missing_session(self): @pytest.mark.asyncio async def test_initial_node_status(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") # Entry node "a" should be ready (no upstream). state = store._sessions["s1"] assert state.nodes["a"].status == "ready" @@ -79,7 +79,7 @@ class TestCompleteNode: @pytest.mark.asyncio async def test_complete_returns_newly_ready(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") newly_ready = await store.complete_node( "s1", "a", {"data": "hello"}, @@ -89,7 +89,7 @@ async def test_complete_returns_newly_ready(self): @pytest.mark.asyncio async def test_complete_chain(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {"data": "hello"}) newly_ready = await store.complete_node( @@ -100,7 +100,7 @@ async def test_complete_chain(self): @pytest.mark.asyncio async def test_diamond_both_deps_needed(self): store = InMemoryGraphSessionStore() - await store.create("s1", _diamond_plan()) + await store.create("s1", _diamond_plan(), "g1") # Complete start → left and right become ready. newly = await store.complete_node("s1", "start", {"data": "x"}) @@ -119,7 +119,7 @@ class TestGetNodeInputs: @pytest.mark.asyncio async def test_gathers_upstream_outputs(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {"data": "hello"}) inputs = await store.get_node_inputs("s1", "b") @@ -128,7 +128,7 @@ async def test_gathers_upstream_outputs(self): @pytest.mark.asyncio async def test_diamond_gathers_from_both_upstreams(self): store = InMemoryGraphSessionStore() - await store.create("s1", _diamond_plan()) + await store.create("s1", _diamond_plan(), "g1") await store.complete_node("s1", "start", {"data": "x"}) await store.complete_node("s1", "left", {"r1": "left_val"}) @@ -154,7 +154,7 @@ async def test_returns_config(self): ), ]]) store = InMemoryGraphSessionStore() - await store.create("s1", plan) + await store.create("s1", plan, "g1") config = await store.get_node_config("s1", "gen") assert config == {"model": "gemini-3"} @@ -168,7 +168,7 @@ async def test_returns_empty_if_no_config(self): ), ]]) store = InMemoryGraphSessionStore() - await store.create("s1", plan) + await store.create("s1", plan, "g1") assert await store.get_node_config("s1", "n") == {} @@ -176,7 +176,7 @@ class TestSuspendResume: @pytest.mark.asyncio async def test_suspend_and_load(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") state = SuspendedNodeState( node_id="b", interaction_id="int-1", @@ -200,13 +200,13 @@ class TestGraphLifecycle: @pytest.mark.asyncio async def test_is_graph_complete_false_initially(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") assert not await store.is_graph_complete("s1") @pytest.mark.asyncio async def test_is_graph_complete_true_when_all_done(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {}) await store.complete_node("s1", "b", {}) await store.complete_node("s1", "c", {}) @@ -215,7 +215,7 @@ async def test_is_graph_complete_true_when_all_done(self): @pytest.mark.asyncio async def test_get_graph_outputs(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {"data": "x"}) await store.complete_node("s1", "b", {"context": "y"}) @@ -228,7 +228,7 @@ class TestMarkNodeFailed: @pytest.mark.asyncio async def test_marks_node_failed(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {"data": "x"}) await store.mark_node_failed("s1", "b", "some error") @@ -239,7 +239,7 @@ async def test_marks_node_failed(self): @pytest.mark.asyncio async def test_skips_dependents(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.complete_node("s1", "a", {"data": "x"}) await store.mark_node_failed("s1", "b", "error") @@ -251,7 +251,7 @@ class TestEventLog: @pytest.mark.asyncio async def test_append_and_get_events(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") idx0 = await store.append_event("s1", {"type": "graphStart"}) idx1 = await store.append_event("s1", {"type": "nodeStart", "nodeId": "a"}) @@ -265,7 +265,7 @@ async def test_append_and_get_events(self): @pytest.mark.asyncio async def test_get_events_after_index(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.append_event("s1", {"type": "a"}) await store.append_event("s1", {"type": "b"}) await store.append_event("s1", {"type": "c"}) @@ -279,13 +279,13 @@ class TestStatus: @pytest.mark.asyncio async def test_default_status(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") assert await store.get_status("s1") == "running" @pytest.mark.asyncio async def test_set_status(self): store = InMemoryGraphSessionStore() - await store.create("s1", _linear_plan()) + await store.create("s1", _linear_plan(), "g1") await store.set_status("s1", "completed") assert await store.get_status("s1") == "completed" @@ -293,3 +293,115 @@ async def test_set_status(self): async def test_missing_session_status(self): store = InMemoryGraphSessionStore() assert await store.get_status("missing") is None + + +class TestListSessions: + @pytest.mark.asyncio + async def test_list_sessions_by_graph_id(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + await store.create("s2", _linear_plan(), graph_id="g1") + await store.create("s3", _linear_plan(), graph_id="g2") + + g1_sessions = await store.list_sessions("g1") + assert len(g1_sessions) == 2 + ids = {s.session_id for s in g1_sessions} + assert ids == {"s1", "s2"} + + g2_sessions = await store.list_sessions("g2") + assert len(g2_sessions) == 1 + assert g2_sessions[0].session_id == "s3" + + @pytest.mark.asyncio + async def test_list_sessions_empty(self): + store = InMemoryGraphSessionStore() + sessions = await store.list_sessions("nonexistent") + assert sessions == [] + + @pytest.mark.asyncio + async def test_list_sessions_includes_status(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + await store.set_status("s1", "completed") + + sessions = await store.list_sessions("g1") + assert sessions[0].status == "completed" + + @pytest.mark.asyncio + async def test_list_sessions_sorted_newest_first(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + await store.create("s2", _linear_plan(), graph_id="g1") + + sessions = await store.list_sessions("g1") + # s2 created after s1, should be first. + assert sessions[0].session_id == "s2" + assert sessions[1].session_id == "s1" + assert sessions[0].created_at >= sessions[1].created_at + + +class TestDeleteSession: + @pytest.mark.asyncio + async def test_delete_existing_session(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + + result = await store.delete_session("s1") + assert result is True + assert await store.get_status("s1") is None + + @pytest.mark.asyncio + async def test_delete_missing_session(self): + store = InMemoryGraphSessionStore() + result = await store.delete_session("nonexistent") + assert result is False + + @pytest.mark.asyncio + async def test_delete_removes_from_graph_index(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + await store.create("s2", _linear_plan(), graph_id="g1") + + await store.delete_session("s1") + + sessions = await store.list_sessions("g1") + assert len(sessions) == 1 + assert sessions[0].session_id == "s2" + + @pytest.mark.asyncio + async def test_delete_all_cleans_graph_index(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="g1") + + await store.delete_session("s1") + + sessions = await store.list_sessions("g1") + assert sessions == [] + + +class TestGraphIdTracking: + @pytest.mark.asyncio + async def test_graph_id_stored(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), graph_id="my-graph") + + sessions = await store.list_sessions("my-graph") + assert sessions[0].graph_id == "my-graph" + + @pytest.mark.asyncio + async def test_graph_id_stored_from_create(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), "g1") + + sessions = await store.list_sessions("g1") + assert len(sessions) == 1 + assert sessions[0].graph_id == "g1" + + @pytest.mark.asyncio + async def test_created_at_is_set(self): + store = InMemoryGraphSessionStore() + await store.create("s1", _linear_plan(), "g1") + + sessions = await store.list_sessions("g1") + assert sessions[0].created_at > 0 + diff --git a/packages/opal-backend/tests/test_headless_mode.py b/packages/opal-backend/tests/test_headless_mode.py index 8d042393e09..1b06bb612a7 100644 --- a/packages/opal-backend/tests/test_headless_mode.py +++ b/packages/opal-backend/tests/test_headless_mode.py @@ -132,7 +132,7 @@ async def test_store_and_retrieve_headless_input(self): headless_inputs = { "inp": {"role": "user", "parts": [{"text": "Hello"}]}, } - await store.create("s1", plan, headless_inputs=headless_inputs) + await store.create("s1", plan, "g1", headless_inputs=headless_inputs) result = await store.get_headless_input("s1", "inp") assert result == {"role": "user", "parts": [{"text": "Hello"}]} @@ -141,7 +141,7 @@ async def test_store_and_retrieve_headless_input(self): async def test_get_headless_input_missing_node(self): store = InMemoryGraphSessionStore() plan = _input_gen_output_plan() - await store.create("s1", plan, headless_inputs={"other": "x"}) + await store.create("s1", plan, "g1", headless_inputs={"other": "x"}) result = await store.get_headless_input("s1", "inp") assert result is None @@ -151,7 +151,7 @@ async def test_get_headless_input_interactive_session(self): """Interactive session (no headless_inputs) returns None.""" store = InMemoryGraphSessionStore() plan = _input_gen_output_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") result = await store.get_headless_input("s1", "inp") assert result is None @@ -166,7 +166,7 @@ async def test_get_headless_input_missing_session(self): async def test_is_headless_session_true(self): store = InMemoryGraphSessionStore() plan = _input_gen_output_plan() - await store.create("s1", plan, headless_inputs={}) + await store.create("s1", plan, "g1", headless_inputs={}) assert await store.is_headless_session("s1") is True @@ -174,7 +174,7 @@ async def test_is_headless_session_true(self): async def test_is_headless_session_false(self): store = InMemoryGraphSessionStore() plan = _input_gen_output_plan() - await store.create("s1", plan) + await store.create("s1", plan, "g1") assert await store.is_headless_session("s1") is False @@ -203,7 +203,7 @@ async def test_input_auto_resolves_with_headless_value(self): headless_inputs = { "inp": {"role": "user", "parts": [{"text": "Hello headless"}]}, } - await store.create("s1", plan, headless_inputs=headless_inputs) + await store.create("s1", plan, "g1", headless_inputs=headless_inputs) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -237,7 +237,7 @@ async def test_two_inputs_both_auto_resolve(self): "inp1": {"role": "user", "parts": [{"text": "First"}]}, "inp2": {"role": "user", "parts": [{"text": "Second"}]}, } - await store.create("s1", plan, headless_inputs=headless_inputs) + await store.create("s1", plan, "g1", headless_inputs=headless_inputs) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -273,7 +273,7 @@ async def test_required_input_missing_errors(self): input_config={"p-required": True}, ) # Headless mode, but no input for "inp". - await store.create("s1", plan, headless_inputs={}) + await store.create("s1", plan, "g1", headless_inputs={}) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -318,7 +318,7 @@ async def test_optional_input_skips_with_empty_context(self): input_config={"p-required": False}, ) # Headless mode, no input for "inp". - await store.create("s1", plan, headless_inputs={}) + await store.create("s1", plan, "g1", headless_inputs={}) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -344,7 +344,7 @@ async def test_default_required_is_false(self): runner = _make_runner(store, bus) plan = _input_gen_output_plan() # No config → p-required defaults False. - await store.create("s1", plan, headless_inputs={}) + await store.create("s1", plan, "g1", headless_inputs={}) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -376,7 +376,7 @@ async def test_one_supplied_one_optional_skipped(self): headless_inputs = { "inp1": {"role": "user", "parts": [{"text": "Supplied"}]}, } - await store.create("s1", plan, headless_inputs=headless_inputs) + await store.create("s1", plan, "g1", headless_inputs=headless_inputs) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -405,7 +405,7 @@ async def test_input_still_suspends_in_interactive_mode(self): plan = _input_gen_output_plan() # Interactive mode: no headless_inputs. - await store.create("s1", plan) + await store.create("s1", plan, "g1") subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -478,7 +478,7 @@ async def _fake_resume_agent( runner._scheduler = scheduler plan = _agent_node_plan() - await store.create("s1", plan, headless_inputs={}) + await store.create("s1", plan, "g1", headless_inputs={}) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -537,7 +537,7 @@ async def _fake_resume_agent( "parts": [{"text": "Write about bananas"}], }, } - await store.create("s1", plan, headless_inputs=headless_inputs) + await store.create("s1", plan, "g1", headless_inputs=headless_inputs) subscriber = bus.subscribe("s1") await runner.start_graph("s1") @@ -576,7 +576,7 @@ async def _fake_run_agent(*, segments, backend, store, graph): plan = _agent_node_plan() # Interactive mode — no headless_inputs. - await store.create("s1", plan) + await store.create("s1", plan, "g1") subscriber = bus.subscribe("s1") await runner.start_graph("s1") diff --git a/packages/visual-editor/src/sca/actions/actions.ts b/packages/visual-editor/src/sca/actions/actions.ts index 7c12a4e8a9d..7b002e23510 100644 --- a/packages/visual-editor/src/sca/actions/actions.ts +++ b/packages/visual-editor/src/sca/actions/actions.ts @@ -22,6 +22,7 @@ import * as NotebookLmPicker from "./notebooklm/notebooklm-picker-actions.js"; import * as Router from "./router/router-actions.js"; import * as Run from "./run/run-actions.js"; import * as BackendRun from "./run/backend-run-action.js"; +import * as Session from "./session/session-actions.js"; import * as ScreenSize from "./screen-size/screen-size-actions.js"; import * as Share from "./share/share-actions.js"; import * as Shell from "./shell/shell-actions.js"; @@ -79,6 +80,7 @@ export function actions( Router.bind({ controller, services, env }); Run.bind({ controller, services, env }); BackendRun.bind({ controller, services, env }); + Session.bind({ controller, services, env }); ScreenSize.bind({ controller, services, env }); Share.bind({ controller, services, env }); Shell.bind({ controller, services, env }); diff --git a/packages/visual-editor/src/sca/actions/binder.ts b/packages/visual-editor/src/sca/actions/binder.ts index 9eb79287e52..f3fdb38fc5f 100644 --- a/packages/visual-editor/src/sca/actions/binder.ts +++ b/packages/visual-editor/src/sca/actions/binder.ts @@ -205,4 +205,11 @@ export function stopRun(controller: AppController): void { run.renderer.reset(); run.main.setStatus(STATUS.STOPPED); run.main.bumpStopVersion(); + + // Disconnect any active devtools session — the run is over, so + // the connection is stale. + const sessionHistory = controller.editor.devtools.sessionHistory; + sessionHistory.activeSessionId = null; + sessionHistory.connectionAbortController?.abort(); + sessionHistory.connectionAbortController = null; } diff --git a/packages/visual-editor/src/sca/actions/run/backend-run-action.ts b/packages/visual-editor/src/sca/actions/run/backend-run-action.ts index 90b41cf4cb4..62ead391455 100644 --- a/packages/visual-editor/src/sca/actions/run/backend-run-action.ts +++ b/packages/visual-editor/src/sca/actions/run/backend-run-action.ts @@ -29,9 +29,18 @@ import type { OutputValues, Schema, } from "@breadboard-ai/types"; -import type { GraphRunEvent } from "../../services/graph-run-service.js"; +import type { + GraphRunEvent, + GraphRunSession, +} from "../../services/graph-run-service.js"; import { AgentEventConsumer } from "../../../a2/agent/agent-event-consumer.js"; +import type { + WaitForInputPayload, + WaitForChoicePayload, +} from "../../../a2/agent/agent-event.js"; import { addChatOutput } from "../../../a2/agent/chat-output.js"; +import { ChoicePresenter } from "../../../a2/agent/choice-presenter.js"; +import { A2UIInteraction } from "../../../a2/agent/a2ui-interaction.js"; import { ConsoleProgressManager } from "../../../a2/agent/console-progress-manager.js"; import { registerProgressHandlers } from "../../../a2/agent/register-progress-handlers.js"; import { RunController } from "../../controller/subcontrollers/run/run-controller.js"; @@ -51,8 +60,8 @@ import { const LABEL = "Backend Run"; -export { startBackendRun, processEvent }; -export type { ProcessResult, NodeEventBridge }; +export { startBackendRun, connectToSession, processEvent, handleSuspend }; +export type { ProcessResult, NodeEventBridge, EventMode }; export const bind = makeAction(); @@ -80,6 +89,151 @@ function raceAbort(promise: Promise, signal: AbortSignal): Promise { }); } +// --------------------------------------------------------------------------- +// Input-node suspend helper +// --------------------------------------------------------------------------- + +/** The inputRequired event narrowed to its specific shape. */ +type InputRequiredEvent = Extract; + +/** + * Default schema used when the backend's input node has no schema. + */ +const DEFAULT_INPUT_SCHEMA: Schema = { + type: "object", + properties: { + request: { + type: "object", + title: "Please provide input", + behavior: ["transient", "llm-content"] as BehaviorSchema[], + format: "asterisk", + }, + }, +}; + +/** + * Handles the suspend flow for any `inputRequired` event. + * + * Examines `suspendEvent` to determine what kind of input is needed: + * - `inputNode`: form-based input from an input node + * - `waitForInput`: text/edit input from an agent + * - `waitForChoice`: multiple-choice selection from an agent + * + * Shows the appropriate UI in the side-nav, waits for the user's + * response, and resumes the session. + * + * @returns `true` if input was collected and the session resumed, + * `false` if the suspend type is unrecognized. + */ +async function handleSuspend( + inputEvent: InputRequiredEvent, + controller: AppController, + session: GraphRunSession, + abortController: AbortController +): Promise { + const { suspendEvent } = inputEvent; + if (!suspendEvent) return false; + + const entry = controller.run.main.console.get(inputEvent.nodeId); + if (!entry) { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.warning( + `inputRequired for ${inputEvent.nodeId} but no console entry` + ), + LABEL + ); + return false; + } + + const appScreen = controller.run.screen.screens.get(inputEvent.nodeId); + + // Determine what to show and create the input Promise. + let inputPromise: Promise | undefined; + + if ("inputNode" in suspendEvent) { + // Input node — form-based input. + const inputNode = suspendEvent.inputNode as { schema?: Schema } | undefined; + const schema: Schema = + inputNode?.schema && (inputNode.schema as Schema).properties + ? (inputNode.schema as Schema) + : DEFAULT_INPUT_SCHEMA; + + // If the schema's request title is empty, use the node's metadata title. + const requestProp = schema.properties?.request; + if (requestProp && !requestProp.title) { + const inspectable = controller.editor.graph.get()?.graphs.get(""); + const node = inspectable?.nodeById(inputEvent.nodeId); + requestProp.title = node?.title() ?? "Please provide input"; + } + + const promptTitle = + schema.properties?.request?.title || "Please provide input"; + addChatOutput( + { role: "model", parts: [{ text: promptTitle }] }, + entry, + appScreen + ); + + inputPromise = entry.requestInput(schema); + } else if ("waitForInput" in suspendEvent) { + // Agent text/edit input. + const event = suspendEvent.waitForInput as WaitForInputPayload; + addChatOutput(event.prompt, entry, appScreen); + + const behaviors: BehaviorSchema[] = ["transient", "llm-content"]; + if (!event.skipLabel) { + behaviors.push("hint-required"); + } + const schema: Schema = { + properties: { + input: { + type: "object", + behavior: behaviors, + format: event.inputType, + }, + }, + }; + + inputPromise = entry.requestInput(schema, event.skipLabel); + } else if ("waitForChoice" in suspendEvent) { + // Agent multiple-choice selection — use the same ChoicePresenter + // that the A2 agent uses to render proper choice buttons. + const event = suspendEvent.waitForChoice as WaitForChoicePayload; + addChatOutput(event.prompt, entry, appScreen); + + const interaction = new A2UIInteraction(entry, appScreen); + const choicePresenter = new ChoicePresenter( + null as never, // translator unused — presentTranslatedChoices skips it + interaction + ); + + inputPromise = choicePresenter.presentTranslatedChoices( + event.prompt, + event.choices, + undefined, // surfaceId default + event.selectionMode, + event.noneOfTheAboveLabel + ); + } + + if (!inputPromise) return false; + + controller.run.main.setStatus(STATUS.PAUSED); + const userInput = await raceAbort(inputPromise, abortController.signal); + + controller.run.main.setStatus(STATUS.RUNNING); + await session.resume(inputEvent.interactionId, userInput); + + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Resumed session after suspend on ${inputEvent.nodeId}` + ), + LABEL + ); + + return true; +} + // --------------------------------------------------------------------------- // Per-node event bridge // --------------------------------------------------------------------------- @@ -106,6 +260,152 @@ interface NodeEventBridge { pendingInput?: Promise; } +// --------------------------------------------------------------------------- +// Shared run-state setup and event-loop helpers +// --------------------------------------------------------------------------- + +/** + * Common run-state initialization shared by `startBackendRun` and + * `connectToSession`. Resets controllers, wires input lifecycle, + * and starts the screen progress ticker. + * + * @returns A handle to clear the progress ticker in `finally`. + */ +function initRunState(controller: AppController): ReturnType { + controller.run.main.reset(); + controller.run.renderer.reset(); + controller.run.screen.reset(); + controller.run.main.setStatus(STATUS.RUNNING); + + controller.run.main.onInputRequested = (id, schema, skipLabel) => + handleInputRequested(id, schema, controller.run, skipLabel); + + return setInterval(() => { + for (const screen of controller.run.screen.screens.values()) { + tickScreenProgress(screen); + } + }, 250); +} + +/** Options that differ between a fresh run and a session reconnect. */ +interface ConsumeOptions { + /** Start in replay mode — events before `replayComplete` are non-interactive. */ + replay?: boolean; +} + +/** + * Consumes events from a `GraphRunSession`, dispatching each through + * `processEvent` and handling suspend/resume. This is the shared core + * of both `startBackendRun` and `connectToSession`. + * + * The loop reconnects the SSE stream after every suspend/resume cycle + * and runs until the session completes, errors, or is aborted. + */ +async function consumeSessionEvents( + session: GraphRunSession, + controller: AppController, + abortController: AbortController, + options: ConsumeOptions = {} +): Promise { + const bridges = new Map(); + let mode: EventMode = options.replay ? "replay" : "live"; + const modeGetter = () => mode; + + // Track the last inputRequired seen during replay — if the session + // is still suspended, we need to show the input UI after replay. + let pendingReplayInput: InputRequiredEvent | null = null; + + let running = true; + while (running) { + const events = session.openStream(abortController.signal); + + for await (const event of events) { + if (abortController.signal.aborted) { + running = false; + break; + } + + // ── Replay bookkeeping ── + if (mode === "replay" && event.type === "inputRequired") { + pendingReplayInput = event as InputRequiredEvent; + } + if ( + mode === "replay" && + pendingReplayInput && + (event.type === "nodeEnd" || event.type === "nodeError") && + event.nodeId === pendingReplayInput.nodeId + ) { + pendingReplayInput = null; + } + + const result = processEvent(event, controller, bridges, mode, modeGetter); + + // ── Replay → live transition ── + if (result === "replayComplete") { + mode = "live"; + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Replay complete — switching to live mode` + ), + LABEL + ); + + if (pendingReplayInput) { + processEvent(pendingReplayInput, controller, bridges, "live", modeGetter); + const handled = await handleSuspend( + pendingReplayInput, controller, session, abortController + ); + if (!handled) running = false; + pendingReplayInput = null; + break; + } + continue; + } + + if (result === "done") { + running = false; + break; + } + + if (result === "suspend") { + const inputEvent = event as InputRequiredEvent; + const bridge = bridges.get(inputEvent.nodeId); + + if (bridge?.pendingInput) { + // Agent node path — the `waitForInput` handler already + // showed the input UI and created a Promise on the bridge. + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Waiting for user input on node ${inputEvent.nodeId}` + ), + LABEL + ); + + const userInput = await raceAbort( + bridge.pendingInput, abortController.signal + ); + bridge.pendingInput = undefined; + + controller.run.main.setStatus(STATUS.RUNNING); + await session.resume(inputEvent.interactionId, userInput); + } else { + // Input node or agent suspend without a bridge — the unified + // handler reads suspendEvent to determine what UI to show. + const handled = await handleSuspend( + inputEvent, controller, session, abortController + ); + if (!handled) running = false; + } + + // Break inner loop to reconnect stream. + break; + } + } + } + + bridges.clear(); +} + // --------------------------------------------------------------------------- // Main action // --------------------------------------------------------------------------- @@ -139,33 +439,25 @@ async function startBackendRun(): Promise { return; } - // Reset state before starting. - controller.run.main.reset(); - controller.run.renderer.reset(); - controller.run.screen.reset(); - controller.run.main.setStatus(STATUS.RUNNING); - - // Wire input lifecycle: when a console entry calls requestInputForNode, - // notify the input queue to show the UI — same wiring as run-actions.ts. - controller.run.main.onInputRequested = (id, schema, skipLabel) => - handleInputRequested(id, schema, controller.run, skipLabel); - - // Start ticking progress every 250ms — same as onRunnerStart. - const progressTickerHandle = setInterval(() => { - for (const screen of controller.run.screen.screens.values()) { - tickScreenProgress(screen); - } - }, 250); - + const progressTickerHandle = initRunState(controller); const abortController = new AbortController(); controller.run.main.abortController = abortController; - // Per-node event bridges for dispatching agent and thought events. - const bridges = new Map(); - try { + // Extract the Drive file ID from the graph URL for session scoping. + const graphUrl = controller.editor.graph.url; + const graphId = graphUrl?.startsWith("drive:/") + ? graphUrl.replace("drive:/", "") + : ""; + if (!graphId) { + throw new Error( + "Backend graph runner requires a Drive-backed Opal (no graphId)" + ); + } + const session = await graphRunService.createSession( graph, + graphId, abortController.signal ); Utils.Logging.getLogger(controller).log( @@ -175,164 +467,11 @@ async function startBackendRun(): Promise { LABEL ); - // Main stream-consume loop — reconnects on suspend. - let running = true; - while (running) { - const events = session.openStream(abortController.signal); + // Track as the active session in the devtools sessions panel. + controller.editor.devtools.sessionHistory.activeSessionId = + session.sessionId; - for await (const event of events) { - if (abortController.signal.aborted) { - running = false; - break; - } - - const result = processEvent(event, controller, bridges); - - if (result === "done") { - running = false; - break; - } - - if (result === "suspend") { - // The `waitForInput` agent event (which arrived before this - // `inputRequired` graph event) has already shown the input UI - // and created a Promise on the bridge. We now: - // 1. Wait for the user's response - // 2. POST :resume with the response - // 3. Break the inner loop so the outer loop reconnects - - const inputEvent = event as Extract< - GraphRunEvent, - { type: "inputRequired" } - >; - const bridge = bridges.get(inputEvent.nodeId); - - if (bridge?.pendingInput) { - // Agent node path — the `waitForInput` handler already - // showed the input UI and created a Promise on the bridge. - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.verbose( - `Waiting for user input on node ${inputEvent.nodeId}` - ), - LABEL - ); - - const userInput = await raceAbort( - bridge.pendingInput, abortController.signal - ); - bridge.pendingInput = undefined; - - controller.run.main.setStatus(STATUS.RUNNING); - await session.resume(inputEvent.interactionId, userInput); - - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.verbose( - `Resumed session after input on node ${inputEvent.nodeId}` - ), - LABEL - ); - } else if (inputEvent.suspendEvent?.inputNode) { - // Input node path — no agent involved. The backend built - // the schema from the node's config (description, modality, - // required). Use it directly. - const inputNode = inputEvent.suspendEvent.inputNode as { - schema?: Schema; - }; - const schema: Schema = inputNode.schema && - (inputNode.schema as Schema).properties - ? (inputNode.schema as Schema) - : { - type: "object", - properties: { - request: { - type: "object", - title: "Please provide input", - behavior: [ - "transient", - "llm-content", - ] as BehaviorSchema[], - format: "asterisk", - }, - }, - }; - - // If the schema's request title is empty or the generic - // fallback, use the node's metadata title instead - // (e.g. "Topic"). - const requestProp = schema.properties?.request; - if (requestProp && !requestProp.title) { - const inspectable = controller.editor.graph - .get() - ?.graphs.get(""); - const node = inspectable?.nodeById(inputEvent.nodeId); - requestProp.title = node?.title() ?? "Please provide input"; - } - - const entry = controller.run.main.console.get( - inputEvent.nodeId - ); - if (entry) { - // Show the prompt text as a chat bubble above the input - // form — mirrors the `report()` call in the TS ask-user - // module. The floating-input component doesn't render - // the schema title itself. - const promptTitle = - schema.properties?.request?.title || "Please provide input"; - const appScreen = controller.run.screen.screens.get( - inputEvent.nodeId - ); - addChatOutput( - { role: "model", parts: [{ text: promptTitle }] }, - entry, - appScreen - ); - - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.verbose( - `Input node ${inputEvent.nodeId} requesting input` - ), - LABEL - ); - - controller.run.main.setStatus(STATUS.PAUSED); - const userInput = await raceAbort( - entry.requestInput(schema), abortController.signal - ); - - controller.run.main.setStatus(STATUS.RUNNING); - await session.resume(inputEvent.interactionId, userInput); - - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.verbose( - `Resumed session after input node ${inputEvent.nodeId}` - ), - LABEL - ); - } else { - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.warning( - `inputRequired for input node ${inputEvent.nodeId} but no console entry` - ), - LABEL - ); - running = false; - } - } else { - // Unknown suspend type — log and stop gracefully. - Utils.Logging.getLogger(controller).log( - Utils.Logging.Formatter.warning( - `inputRequired but no handler for node ${inputEvent.nodeId}` - ), - LABEL - ); - running = false; - } - - // Break inner loop to reconnect stream. - break; - } - } - } + await consumeSessionEvents(session, controller, abortController); } catch (error) { if (abortController.signal.aborted) { Utils.Logging.getLogger(controller).log( @@ -349,19 +488,88 @@ async function startBackendRun(): Promise { }); } } finally { - // Cleanup — mirrors onRunnerEnd. clearInterval(progressTickerHandle); - bridges.clear(); controller.run.main.clearInput(); controller.run.main.setStatus(STATUS.STOPPED); } } +// --------------------------------------------------------------------------- +// Session reconnection +// --------------------------------------------------------------------------- + +/** + * Connects the graph UI to an existing backend session. + * + * The event stream is processed in two phases: + * 1. **Replay** (events before `replayComplete`): builds up console + * entries, node states, and agent outputs without blocking on input. + * 2. **Live** (events after `replayComplete`): normal interactive mode + * where `inputRequired` suspends and waits for user input. + * + * If the session has already completed, only the replay phase runs + * and no `replayComplete` marker is emitted — the stream ends with + * `graphComplete` / `graphError`. + */ +async function connectToSession(sessionId: string): Promise { + const { controller, services } = bind; + const graphRunService = services.graphRunService; + const sessionHistory = controller.editor.devtools.sessionHistory; + + // Cancel any existing connection. + sessionHistory.connectionAbortController?.abort(); + + const progressTickerHandle = initRunState(controller); + const abortController = new AbortController(); + sessionHistory.connectionAbortController = abortController; + sessionHistory.activeSessionId = sessionId; + + try { + const session = graphRunService.connectSession(sessionId); + + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Connecting to session ${sessionId} (replay mode)` + ), + LABEL + ); + + await consumeSessionEvents(session, controller, abortController, { + replay: true, + }); + } catch (error) { + if (!abortController.signal.aborted) { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.warning( + `Session connection failed: ${String(error)}` + ), + LABEL + ); + controller.run.main.setError({ + message: String(error), + }); + } + } finally { + clearInterval(progressTickerHandle); + controller.run.main.clearInput(); + controller.run.main.setStatus(STATUS.STOPPED); + sessionHistory.connectionAbortController = null; + } +} + + // --------------------------------------------------------------------------- // Event processing // --------------------------------------------------------------------------- -type ProcessResult = "continue" | "done" | "suspend"; +type ProcessResult = "continue" | "done" | "suspend" | "replayComplete"; + +/** + * Controls how events are processed. + * - `"live"`: normal interactive mode (suspend on inputRequired) + * - `"replay"`: skip inputRequired, no interactive input + */ +type EventMode = "live" | "replay"; /** * Maps a single Heartstone SSE event to controller state updates. @@ -370,7 +578,9 @@ type ProcessResult = "continue" | "done" | "suspend"; function processEvent( event: GraphRunEvent, controller: AppController, - bridges: Map + bridges: Map, + mode: EventMode = "live", + modeGetter?: () => EventMode ): ProcessResult { switch (event.type) { case "nodeStart": { @@ -398,7 +608,7 @@ function processEvent( // Create per-node agent bridge. const consoleEntry = controller.run.main.console.get(nodeId); const appScreen = controller.run.screen.screens.get(nodeId); - createBridge(nodeId, consoleEntry, appScreen, controller, bridges); + createBridge(nodeId, consoleEntry, appScreen, controller, bridges, modeGetter ?? (() => mode)); return "continue"; } @@ -516,6 +726,16 @@ function processEvent( } case "inputRequired": { + // In replay mode, don't suspend — input was already provided + // (or the session is still suspended). But do update visual state + // so the node shows as paused rather than spinning. + if (mode === "replay") { + controller.run.main.setStatus(STATUS.PAUSED); + controller.run.renderer.setNodeState(event.nodeId, { + status: "waiting", + }); + return "continue"; + } // Signal the main loop to suspend and handle input. return "suspend"; } @@ -541,6 +761,20 @@ function processEvent( return "done"; } + case "graphCancelled": { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Graph cancelled: ${event.sessionId}` + ), + LABEL + ); + return "done"; + } + + case "replayComplete": { + return "replayComplete"; + } + default: return "continue"; } @@ -566,7 +800,8 @@ function createBridge( consoleEntry: ConsoleEntry | undefined, appScreen: AppScreen | undefined, controller: AppController, - bridges: Map + bridges: Map, + getMode: () => EventMode = () => "live" ): void { const consumer = new AgentEventConsumer(); const progress = new ConsoleProgressManager(consoleEntry, appScreen); @@ -591,6 +826,13 @@ function createBridge( // Register the waitForInput handler — this bridges the UI input // flow to the backend suspend/resume protocol. consumer.on("waitForInput", (event) => { + // Show the agent's prompt text in the console and app screen + // (both replay and live — so conversation history is visible). + addChatOutput(event.prompt, consoleEntry, appScreen); + + // In replay mode, don't show the input UI — input was already provided. + if (getMode() === "replay") return undefined; + const behaviors: BehaviorSchema[] = ["transient", "llm-content"]; if (!event.skipLabel) { behaviors.push("hint-required"); @@ -605,9 +847,6 @@ function createBridge( }, }; - // Show the agent's prompt text in the console and app screen. - addChatOutput(event.prompt, consoleEntry, appScreen); - // requestInput shows the input UI and returns a Promise that // resolves when the user submits. Store it on the bridge for // the main loop to await when inputRequired arrives. @@ -622,3 +861,4 @@ function createBridge( bridges.set(nodeId, bridge); } + diff --git a/packages/visual-editor/src/sca/actions/session/session-actions.ts b/packages/visual-editor/src/sca/actions/session/session-actions.ts new file mode 100644 index 00000000000..77448b785f0 --- /dev/null +++ b/packages/visual-editor/src/sca/actions/session/session-actions.ts @@ -0,0 +1,126 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * @fileoverview + * + * Session management actions for the devtools sessions panel. + * + * - **startSessionMonitor**: opens the monitor SSE stream for the + * current graph and feeds sessionStatus events to the controller. + * - **stopSessionMonitor**: cancels the active monitor stream. + * - **deleteSession**: deletes a session on the backend. + * - **connectToSession**: connects the graph UI to an existing session + * by replaying its events and then switching to live mode. + */ + +import { makeAction } from "../binder.js"; +import { Utils } from "../../utils.js"; + +export { startSessionMonitor, stopSessionMonitor, deleteSession }; + +export const bind = makeAction(); + +const LABEL = "Session Monitor"; + +// --------------------------------------------------------------------------- +// Monitor lifecycle +// --------------------------------------------------------------------------- + +/** + * Opens the session monitor SSE stream for the current graph. + * + * Feeds sessionStatus events to the SessionHistoryController, which + * the sessions panel reads reactively. Cancels any existing monitor. + */ +async function startSessionMonitor(): Promise { + const { controller, services } = bind; + const sessionHistory = controller.editor.devtools.sessionHistory; + const graphRunService = services.graphRunService; + + // Cancel any existing monitor. + stopSessionMonitor(); + + const graphUrl = controller.editor.graph.url; + const graphId = graphUrl?.startsWith("drive:/") + ? graphUrl.replace("drive:/", "") + : ""; + + if (!graphId) { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + "No graphId — skipping session monitor" + ), + LABEL + ); + return; + } + + const abortController = new AbortController(); + sessionHistory.monitorAbortController = abortController; + + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose( + `Starting session monitor for graph ${graphId}` + ), + LABEL + ); + + try { + for await (const event of graphRunService.monitorSessions( + graphId, + abortController.signal + )) { + if (abortController.signal.aborted) break; + sessionHistory.applySessionStatus(event); + } + } catch (error) { + if (!abortController.signal.aborted) { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.warning( + `Session monitor error: ${String(error)}` + ), + LABEL + ); + } + } +} + +/** Cancels the active session monitor stream. */ +function stopSessionMonitor(): void { + const { controller } = bind; + const sessionHistory = controller.editor.devtools.sessionHistory; + sessionHistory.monitorAbortController?.abort(); + sessionHistory.monitorAbortController = null; +} + +// --------------------------------------------------------------------------- +// Session deletion +// --------------------------------------------------------------------------- + +/** Deletes a session on the backend. */ +async function deleteSession(sessionId: string): Promise { + const { controller, services } = bind; + const graphRunService = services.graphRunService; + + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.verbose(`Deleting session ${sessionId}`), + LABEL + ); + + try { + await graphRunService.deleteSession(sessionId); + // The monitor stream will receive a sessionStatus with + // status: "deleted" and remove it from the controller. + } catch (error) { + Utils.Logging.getLogger(controller).log( + Utils.Logging.Formatter.warning( + `Delete session failed: ${String(error)}` + ), + LABEL + ); + } +} diff --git a/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/devtools-controller.ts b/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/devtools-controller.ts index 07c801b8f59..304606357a6 100644 --- a/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/devtools-controller.ts +++ b/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/devtools-controller.ts @@ -7,6 +7,7 @@ import { field } from "../../../decorators/field.js"; import { RootController } from "../../root-controller.js"; import { OpieController } from "./opie-controller.js"; +import { SessionHistoryController } from "./session-history-controller.js"; export class DevToolsController extends RootController { @field({ persist: "session" }) @@ -16,6 +17,7 @@ export class DevToolsController extends RootController { accessor activeTab = "opie"; public readonly opie: OpieController; + public readonly sessionHistory: SessionHistoryController; constructor(controllerId: string, persistenceId: string) { super(controllerId, persistenceId); @@ -23,5 +25,10 @@ export class DevToolsController extends RootController { `${controllerId}_Opie`, `${persistenceId}_Opie` ); + this.sessionHistory = new SessionHistoryController( + `${controllerId}_SessionHistory`, + `${persistenceId}_SessionHistory` + ); } } + diff --git a/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/session-history-controller.ts b/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/session-history-controller.ts new file mode 100644 index 00000000000..68c60ea47d3 --- /dev/null +++ b/packages/visual-editor/src/sca/controller/subcontrollers/editor/devtools/session-history-controller.ts @@ -0,0 +1,81 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * @fileoverview + * + * Reactive state for the session history panel in devtools. + * + * Mirrors the backend's session monitor stream: each `sessionStatus` + * event either upserts or removes an entry in the sessions map. + * UI components read the signals; the monitor action writes them. + */ + +import { field } from "../../../decorators/field.js"; +import { RootController } from "../../root-controller.js"; + +export { SessionHistoryController }; +export type { SessionEntry }; + +/** Frontend representation of a single session. */ +interface SessionEntry { + sessionId: string; + status: string; + createdAt: number; +} + +class SessionHistoryController extends RootController { + /** + * Reactive session map — keyed by sessionId. + * Updated by the session monitor action. + * `deep: true` tracks mutations so `.set()` / `.delete()` trigger updates. + */ + @field({ deep: true }) + accessor sessions: Map = new Map(); + + /** + * The session currently connected to the graph UI. + * When null, no session is connected. + */ + @field() + accessor activeSessionId: string | null = null; + + /** + * Abort controller for the active monitor SSE stream. + * Actions use this to cancel the monitor when the graph changes. + */ + monitorAbortController: AbortController | null = null; + + /** + * Abort controller for the active session connection (replay/live stream). + * Used when switching sessions or stopping the current connection. + */ + connectionAbortController: AbortController | null = null; + + /** Process a sessionStatus event from the monitor stream. */ + applySessionStatus(event: SessionEntry): void { + if (event.status === "deleted") { + this.sessions.delete(event.sessionId); + + // If the deleted session was active, disconnect. + if (this.activeSessionId === event.sessionId) { + this.activeSessionId = null; + } + } else { + this.sessions.set(event.sessionId, event); + } + } + + /** Clear all sessions and disconnect. */ + reset(): void { + this.sessions = new Map(); + this.activeSessionId = null; + this.monitorAbortController?.abort(); + this.monitorAbortController = null; + this.connectionAbortController?.abort(); + this.connectionAbortController = null; + } +} diff --git a/packages/visual-editor/src/sca/services/graph-run-service.ts b/packages/visual-editor/src/sca/services/graph-run-service.ts index e717ba5b6c6..de170a11933 100644 --- a/packages/visual-editor/src/sca/services/graph-run-service.ts +++ b/packages/visual-editor/src/sca/services/graph-run-service.ts @@ -26,7 +26,7 @@ import { iteratorFromStream } from "@breadboard-ai/utils"; export { GraphRunService }; -export type { GraphRunEvent, GraphRunSession }; +export type { GraphRunEvent, GraphRunSession, SessionStatusEvent }; // --------------------------------------------------------------------------- // Event types emitted by the Heartstone backend SSE stream @@ -69,8 +69,17 @@ type GraphRunEvent = GraphRunEventBase & outputs: Record>; } | { type: "graphError"; sessionId: string; error: string } + | { type: "graphCancelled"; sessionId: string } + | { type: "replayComplete" } ); +/** Event shape emitted by the session monitor SSE stream. */ +interface SessionStatusEvent { + sessionId: string; + status: string; + createdAt: number; +} + /** * Handle for a running graph session. Separates session lifecycle * from stream lifecycle to support disconnect/reconnect on suspend. @@ -137,13 +146,14 @@ class GraphRunService { */ async createSession( graph: Record, + graphId: string, signal?: AbortSignal ): Promise { const createUrl = `${this.#baseUrl}/v1beta1/graphSessions/new`; const createResp = await this.#fetchFn(createUrl, { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ graph }), + body: JSON.stringify({ graph, graphId }), signal, }); @@ -233,4 +243,127 @@ class GraphRunService { }, }; } + + /** + * Connect to an existing session (for reconnection/replay). + * + * Returns the same GraphRunSession handle as createSession, but + * without creating a new session on the backend. The stream starts + * from index 0 (full replay) or from a cursor position. + */ + connectSession(sessionId: string): GraphRunSession { + const fetchFn = this.#fetchFn; + const baseUrl = this.#baseUrl; + let cursor = -1; + + return { + sessionId, + + openStream(signal?: AbortSignal): AsyncIterable { + const streamUrl = + `${baseUrl}/v1beta1/graphSessions/${sessionId}?after=${cursor}`; + return { + [Symbol.asyncIterator]() { + let started = false; + let iterator: AsyncIterator | null = null; + + return { + async next(): Promise> { + if (!started) { + started = true; + const resp = await fetchFn(streamUrl, { signal }); + if (!resp.ok) { + throw new Error( + `Graph SSE stream failed: ${resp.status} ${resp.statusText}` + ); + } + if (!resp.body) { + throw new Error("Graph SSE response has no body"); + } + const iterable = + iteratorFromStream(resp.body); + iterator = iterable[Symbol.asyncIterator](); + } + const result = await iterator!.next(); + if (!result.done && typeof result.value.index === "number") { + cursor = Math.max(cursor, result.value.index); + } + return result; + }, + }; + }, + }; + }, + + async resume( + interactionId: string, + response: unknown + ): Promise { + const url = `${baseUrl}/v1beta1/graphSessions/${sessionId}:resume`; + const resp = await fetchFn(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ interactionId, response }), + }); + if (!resp.ok) { + throw new Error( + `Resume failed: ${resp.status} ${resp.statusText}` + ); + } + }, + + async cancel(): Promise { + const url = `${baseUrl}/v1beta1/graphSessions/${sessionId}:cancel`; + try { + await fetchFn(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + }); + } catch { + // Best-effort. + } + }, + }; + } + + /** + * Open a monitor SSE stream for a graph's session list. + * + * Yields `sessionStatus` events — one per existing session initially, + * then live updates as sessions are created, change status, or are + * deleted (status: "deleted"). + */ + async *monitorSessions( + graphId: string, + signal?: AbortSignal + ): AsyncIterable { + const url = + `${this.#baseUrl}/v1beta1/graphSessions?graphId=${encodeURIComponent(graphId)}`; + const resp = await this.#fetchFn(url, { signal }); + if (!resp.ok) { + throw new Error( + `Session monitor failed: ${resp.status} ${resp.statusText}` + ); + } + if (!resp.body) { + throw new Error("Session monitor response has no body"); + } + const iterable = iteratorFromStream(resp.body); + yield* iterable; + } + + /** Delete a session on the backend. */ + async deleteSession(sessionId: string): Promise { + const url = + `${this.#baseUrl}/v1beta1/graphSessions/${sessionId}`; + const resp = await this.#fetchFn(url, { + method: "DELETE", + }); + if (!resp.ok) { + throw new Error( + `Delete session failed: ${resp.status} ${resp.statusText}` + ); + } + } } + diff --git a/packages/visual-editor/src/ui/elements/devtools/devtools.ts b/packages/visual-editor/src/ui/elements/devtools/devtools.ts index 0838f65a23a..87c0034433c 100644 --- a/packages/visual-editor/src/ui/elements/devtools/devtools.ts +++ b/packages/visual-editor/src/ui/elements/devtools/devtools.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { LitElement, html, css } from "lit"; +import { LitElement, html, css, nothing } from "lit"; import { customElement } from "lit/decorators.js"; import { SignalWatcher } from "@lit-labs/signals"; import { consume } from "@lit/context"; @@ -14,6 +14,7 @@ import { icons } from "../../styles/icons.js"; import "../json-tree/json-tree.js"; import "./opie/opie-panel.js"; import "./feedback/feedback-panel.js"; +import "./sessions/sessions-panel.js"; @customElement("bb-devtools") export class DevTools extends SignalWatcher(LitElement) { @@ -139,6 +140,8 @@ export class DevTools extends SignalWatcher(LitElement) { const functions = opie.functionDeclarations; const entries = opie.entries; const feedbackEntries = this.sca.controller.global.feedback.entries; + const showSessions = + this.sca.services.graphRunService.enabled; return html`
@@ -163,6 +166,19 @@ export class DevTools extends SignalWatcher(LitElement) { > Feedback + ${showSessions + ? html` + + ` + : nothing}
+ + `; + } + + #onSessionClick(sessionId: string) { + const sessionHistory = + this.sca.controller.editor.devtools.sessionHistory; + + if (sessionHistory.activeSessionId === sessionId) { + // Clicking the active session disconnects it. + sessionHistory.activeSessionId = null; + sessionHistory.connectionAbortController?.abort(); + sessionHistory.connectionAbortController = null; + + // Reset run state so side-nav / console clear. + const controller = this.sca.controller; + controller.run.main.reset(); + controller.run.renderer.reset(); + controller.run.screen.reset(); + } else { + connectToSession(sessionId); + } + } + + #onDeleteClick(sessionId: string) { + deleteSession(sessionId); + } +} + +declare global { + interface HTMLElementTagNameMap { + "bb-devtools-sessions-panel": DevToolsSessionsPanel; + } +} diff --git a/packages/visual-editor/tests/sca/actions/run/backend-run-action.test.ts b/packages/visual-editor/tests/sca/actions/run/backend-run-action.test.ts index cdf3f30c4fc..70232066b70 100644 --- a/packages/visual-editor/tests/sca/actions/run/backend-run-action.test.ts +++ b/packages/visual-editor/tests/sca/actions/run/backend-run-action.test.ts @@ -5,15 +5,19 @@ */ import assert from "node:assert"; -import { beforeEach, suite, test } from "node:test"; -import { setDOM } from "../../../fake-dom.js"; +import { afterEach, beforeEach, mock, suite, test } from "node:test"; +import { setDOM, unsetDOM } from "../../../fake-dom.js"; import { makeTestController } from "../../helpers/mock-controller.js"; import { processEvent, + handleSuspend, type NodeEventBridge, + type EventMode, } from "../../../../src/sca/actions/run/backend-run-action.js"; -import type { GraphRunEvent } from "../../../../src/sca/services/graph-run-service.js"; +import type { GraphRunEvent, GraphRunSession } from "../../../../src/sca/services/graph-run-service.js"; import type { AppController } from "../../../../src/sca/controller/controller.js"; +import { RunController } from "../../../../src/sca/controller/subcontrollers/run/run-controller.js"; +import { STATUS } from "../../../../src/sca/types.js"; // --------------------------------------------------------------------------- // Helpers @@ -80,9 +84,10 @@ function makeControllerWithGraph() { function dispatch( event: GraphRunEvent, controller: AppController, - bridges: Map + bridges: Map, + mode: EventMode = "live" ) { - return processEvent(event, controller, bridges); + return processEvent(event, controller, bridges, mode); } // --------------------------------------------------------------------------- @@ -405,7 +410,7 @@ suite("processEvent", () => { // --- inputRequired ------------------------------------------------------- suite("inputRequired", () => { - test("returns suspend", () => { + test("returns suspend in live mode", () => { const result = dispatch( { type: "inputRequired", @@ -419,6 +424,62 @@ suite("processEvent", () => { assert.strictEqual(result, "suspend"); }); + + test("returns continue and sets paused/waiting state in replay mode", () => { + const result = dispatch( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-1", + index: 0, + }, + controller, + bridges, + "replay" + ); + + assert.strictEqual(result, "continue"); + + // Node should show as "waiting" (not spinning as "working"). + const nodeState = controller.run.renderer.nodes.get("n1"); + assert.ok(nodeState); + assert.strictEqual(nodeState.status, "waiting"); + }); + }); + + // --- graphCancelled ------------------------------------------------------ + + suite("graphCancelled", () => { + test("returns done", () => { + const result = dispatch( + { + type: "graphCancelled", + sessionId: "sess-1", + index: 0, + }, + controller, + bridges + ); + + assert.strictEqual(result, "done"); + }); + }); + + // --- replayComplete ------------------------------------------------------ + + suite("replayComplete", () => { + test("returns replayComplete", () => { + const result = dispatch( + { + type: "replayComplete", + index: 0, + }, + controller, + bridges + ); + + assert.strictEqual(result, "replayComplete"); + }); }); // --- unknown event ------------------------------------------------------- @@ -436,3 +497,329 @@ suite("processEvent", () => { }); }); }); + +// --------------------------------------------------------------------------- +// handleSuspend +// --------------------------------------------------------------------------- + +/** Creates a mock GraphRunSession with a spied `resume`. */ +function makeMockSession(): GraphRunSession & { resumeCalls: unknown[][] } { + const calls: unknown[][] = []; + return { + sessionId: "test-session", + resumeCalls: calls, + async resume(interactionId: string, response: unknown) { + calls.push([interactionId, response]); + }, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + openStream(): any { + return { + [Symbol.asyncIterator]: () => ({ + next: async () => ({ done: true, value: undefined }), + }), + }; + }, + async cancel() {}, + }; +} + +suite("handleSuspend", () => { + let controller: AppController; + let addNode: (id: string, title?: string, outputSchema?: object) => void; + let bridges: Map; + + beforeEach(() => { + setDOM(); + const ctx = makeControllerWithGraph(); + controller = ctx.controller; + addNode = ctx.addNode; + bridges = new Map(); + }); + + afterEach(() => { + mock.restoreAll(); + unsetDOM(); + }); + + /** Dispatch nodeStart to create a console entry for the node. */ + function setupNode(nodeId: string, title = nodeId) { + addNode(nodeId, title); + processEvent( + { type: "nodeStart", nodeId, index: 0 }, + controller, + bridges, + "live" + ); + } + + // --- early returns ------------------------------------------------------- + + test("returns false when suspendEvent is undefined", async () => { + setupNode("n1"); + + const session = makeMockSession(); + const result = await handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-1", + index: 1, + // No suspendEvent. + }, + controller, + session, + new AbortController() + ); + + assert.strictEqual(result, false); + assert.strictEqual(session.resumeCalls.length, 0); + }); + + test("returns false when no console entry exists for nodeId", async () => { + // Don't call setupNode — no entry. + const session = makeMockSession(); + const result = await handleSuspend( + { + type: "inputRequired", + nodeId: "unknown-node", + interactionId: "int-1", + suspendEvent: { inputNode: { schema: { properties: {} } } }, + index: 1, + }, + controller, + session, + new AbortController() + ); + + assert.strictEqual(result, false); + assert.strictEqual(session.resumeCalls.length, 0); + }); + + test("returns false for unrecognized suspendEvent type", async () => { + setupNode("n1"); + + const session = makeMockSession(); + const result = await handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-1", + suspendEvent: { unknownType: { data: 123 } }, + index: 1, + }, + controller, + session, + new AbortController() + ); + + assert.strictEqual(result, false); + assert.strictEqual(session.resumeCalls.length, 0); + }); + + // --- inputNode ----------------------------------------------------------- + + test("inputNode: calls requestInput and resumes session with user input", async () => { + setupNode("n1", "Topic Input"); + + const session = makeMockSession(); + const abortController = new AbortController(); + + const suspendPromise = handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-abc", + suspendEvent: { + inputNode: { + schema: { + type: "object", + properties: { + request: { + type: "object", + title: "Enter topic", + behavior: ["transient", "llm-content"], + }, + }, + }, + }, + }, + index: 1, + }, + controller, + session, + abortController + ); + + // Status should transition to PAUSED while waiting. + const runCtrl = controller.run.main as RunController; + assert.strictEqual(runCtrl.status, STATUS.PAUSED); + + // Simulate user providing input. + runCtrl.resolveInputForNode("n1", { request: "peanuts" }); + + const result = await suspendPromise; + + assert.strictEqual(result, true); + assert.strictEqual(runCtrl.status, STATUS.RUNNING); + assert.strictEqual(session.resumeCalls.length, 1); + assert.strictEqual(session.resumeCalls[0][0], "int-abc"); + assert.deepStrictEqual(session.resumeCalls[0][1], { request: "peanuts" }); + }); + + test("inputNode: uses default schema when inputNode has no schema", async () => { + setupNode("n1", "Bare Input"); + + const session = makeMockSession(); + const abortController = new AbortController(); + + const suspendPromise = handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-def", + suspendEvent: { inputNode: {} }, + index: 1, + }, + controller, + session, + abortController + ); + + // Should still be paused — requestInput was called with default schema. + const runCtrl = controller.run.main as RunController; + assert.strictEqual(runCtrl.status, STATUS.PAUSED); + + // Resolve with some default input. + runCtrl.resolveInputForNode("n1", { request: "default value" }); + + const result = await suspendPromise; + assert.strictEqual(result, true); + assert.strictEqual(session.resumeCalls.length, 1); + }); + + // --- waitForInput -------------------------------------------------------- + + test("waitForInput: calls requestInput with agent schema and resumes", async () => { + setupNode("n1", "Agent Node"); + + const session = makeMockSession(); + const abortController = new AbortController(); + + const suspendPromise = handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-wait", + suspendEvent: { + waitForInput: { + requestId: "req-1", + prompt: { role: "model", parts: [{ text: "What topic?" }] }, + inputType: "text", + interactionId: "int-wait", + }, + }, + index: 1, + }, + controller, + session, + abortController + ); + + const runCtrl = controller.run.main as RunController; + assert.strictEqual(runCtrl.status, STATUS.PAUSED); + + // Resolve input. + runCtrl.resolveInputForNode("n1", { + input: { parts: [{ text: "haiku about cats" }], role: "user" }, + }); + + const result = await suspendPromise; + assert.strictEqual(result, true); + assert.strictEqual(runCtrl.status, STATUS.RUNNING); + assert.strictEqual(session.resumeCalls.length, 1); + assert.strictEqual(session.resumeCalls[0][0], "int-wait"); + }); + + test("waitForInput: includes hint-required behavior when skipLabel is falsy", async () => { + setupNode("n1", "Agent Node"); + + const runCtrl = controller.run.main as RunController; + + // Capture the schema passed to requestInputForNode via callback. + let capturedSchema: Record | null = null; + runCtrl.onInputRequested = (_id, schema) => { + capturedSchema = schema as Record; + }; + + const session = makeMockSession(); + handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-hint", + suspendEvent: { + waitForInput: { + requestId: "req-2", + prompt: { role: "model", parts: [{ text: "Prompt" }] }, + inputType: "edit_note", + // skipLabel is undefined → should include hint-required + interactionId: "int-hint", + }, + }, + index: 1, + }, + controller, + session, + new AbortController() + ); + + assert.ok(capturedSchema, "onInputRequested should have been called"); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const inputProp = (capturedSchema as any).properties?.input; + assert.ok(inputProp, "schema should have input property"); + const behaviors = inputProp.behavior as string[]; + assert.ok( + behaviors.includes("hint-required"), + "should include hint-required when skipLabel is falsy" + ); + + // Resolve to avoid dangling promise. + runCtrl.resolveInputForNode("n1", { input: {} }); + }); + + // --- abort --------------------------------------------------------------- + + test("rejects when aborted before user provides input", async () => { + setupNode("n1"); + + const session = makeMockSession(); + const abortController = new AbortController(); + + const suspendPromise = handleSuspend( + { + type: "inputRequired", + nodeId: "n1", + interactionId: "int-abort", + suspendEvent: { + inputNode: { + schema: { properties: { request: { type: "string" } } }, + }, + }, + index: 1, + }, + controller, + session, + abortController + ); + + // Abort before user provides input. + abortController.abort(); + + await assert.rejects(suspendPromise); + assert.strictEqual( + session.resumeCalls.length, + 0, + "should not resume when aborted" + ); + }); +}); diff --git a/projects/heartstone/PROJECT.md b/projects/heartstone/PROJECT.md index f03bc8486ed..f4379f52051 100644 --- a/projects/heartstone/PROJECT.md +++ b/projects/heartstone/PROJECT.md @@ -816,7 +816,7 @@ POST /v1beta1/graphSessions/new → {sessionId} GET /v1beta1/graphSessions/{id}?after=-1 - → SSE stream (replay + live) + → SSE stream (replay + replayComplete marker + live) POST /v1beta1/graphSessions/{id}:resume {interactionId, response, accessToken?} @@ -827,6 +827,12 @@ GET /v1beta1/graphSessions/{id}/status POST /v1beta1/graphSessions/{id}:cancel → {sessionId, status: "cancelled"} + +GET /v1beta1/graphSessions?graphId=X + → SSE stream (initial session list + live status changes) + +DELETE /v1beta1/graphSessions/{id} + → {ok: true} ``` --- @@ -860,10 +866,13 @@ POST /v1beta1/graphSessions/{id}:cancel ### New Frontend Code -| Component | Location | Purpose | -| --------------------- | ----------------------------- | ----------------------------- | -| `GraphRunService` | `visual-editor/sca/services/` | Backend graph run management | -| `GraphRunEventSource` | `visual-editor/` | SSE consumer for graph events | +| Component | Location | Purpose | +| ------------------------ | --------------------------------------- | -------------------------------------------- | +| `GraphRunService` | `visual-editor/sca/services/` | Backend graph run management | +| `GraphRunEventSource` | `visual-editor/` | SSE consumer for graph events | +| `SessionListController` | `visual-editor/sca/controller/devtools` | Session list mirror + connected session | +| `session-actions.ts` | `visual-editor/sca/actions/run` | Connect/disconnect/delete/monitor actions | +| `sessions-panel.ts` | `visual-editor/ui/elements/devtools` | Devtools "Sessions" tab UI | ### Rename @@ -962,11 +971,60 @@ defaults and `agentContext`. 🎯 **Objective:** `POST {graphId, mode: "headless", inputs: {...}}` loads and runs a graph without frontend interaction. -### Phase 11: Reconnection + History - -Event replay on reconnect. Read-only history view for completed runs. - -🎯 **Objective:** Close/reopen browser — UI reconstructs from replayed events. +### Phase 11: Session Management + Reconnection + +Add a "Sessions" tab to `bb-devtools` (behind `enableBackendGraphRunner` flag) +showing per-graph session history. Allow connecting, disconnecting, and +reconnecting the UI to running or completed sessions. Refactor +`backend-run-action` to support soft disconnect (switch session) vs. hard +cancel (restart). + +🎯 **Objective:** Start a graph run, disconnect (switch to another session or +just disconnect), then reconnect — UI reconstructs from replayed events and +picks up live events if the session is still running. Suspended sessions +allow providing input on reconnect. + +#### Decisions + +| # | Decision | Resolution | +| --- | --------------------------------- | ------------------------------------------------------------------------- | +| 1 | Session list source of truth | Server (`GraphSessionStore`), not frontend. Lightweight monitor SSE. | +| 2 | Session scoping | Per-graph, indexed by caller-provided `graphId` (typically Drive file ID) | +| 3 | Replay→live boundary | Explicit `replayComplete` marker event emitted by server | +| 4 | Suspended session reconnect | Allow — show input UI after replay | +| 5 | Cancel vs. disconnect | Two-tier: Restart = `POST :cancel`. Switch session = soft disconnect. | +| 6 | Session cleanup | `DELETE /v1beta1/graphSessions/{id}` endpoint + trash icon in UI | +| 7 | Tab naming | "Sessions" (devtools view) | + +#### Sub-phases + +**11a — Backend: Session monitor + delete.** +`GraphSessionStore` extended with `graph_id` tracking, `list_sessions()`, +`delete_session()`. `InMemoryGraphSessionStore` publishes status changes to a +graph-level `EventBus` channel (`graph:{graphId}`). New `GET /` monitor SSE +endpoint streams initial session list + status updates. New `DELETE /{id}` +endpoint. `replayComplete` marker added to per-session SSE stream. + +**11b — Frontend: Refactor `backend-run-action`.** +Extract `consumeSessionEvents()` with two modes: +- **Live** — same as current loop: incremental node states, input handling. +- **Replay** — accumulate node states silently, populate console from event + history, skip interactive handlers. At `replayComplete`, apply cumulative + state and switch to live mode. Special case: if session is suspended, + show input UI after replay. + +Introduce soft disconnect mechanism (separate `disconnectController` that +stops the SSE reader without canceling the backend session). + +**11c — Frontend: Session list controller + devtools tab.** +`SessionListController` under `editor.devtools.sessions` — thin reactive +mirror of the monitor stream. Sessions panel UI in devtools with status +icons, timestamps, connect/disconnect/delete. + +**11d — Frontend: Session actions + monitor lifecycle.** +`session-actions.ts`: `connectSession()` (replay + live), `disconnectSession()` +(soft), `deleteSession()`. Monitor action subscribes to the session list SSE +when the graph changes. ---