[heartstone] Phase 11 — Session Management + Reconnection

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,

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,

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(

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