WebSocket reconnection fails during approval blocks

[wu-changxing]

Bug

When a client refreshes while the agent is blocked waiting for approval (e.g., bash tool), reconnection fails with "No active connection" error.

Reproduction

1. Start a hosted agent with tool approval enabled
2. Send a prompt that triggers a bash tool
3. Agent sends approval_needed, blocks on io.receive()
4. Refresh the browser page
5. Frontend shows "No active connection" error
6. Agent logs: ✗ bash - connection closed

Root Cause

Three bugs compound:

1. run_agent() has no error handling

If the agent crashes (e.g., on io_closed sentinel), agent_finished.set() never fires. _pipe_ws_io hangs forever waiting.

File: connectonion/network/asgi/websocket.py — run_agent() function

2. Reattach uses closed IO

On reconnect, the server reattaches to the old io object which has _closed = True (set by io.close() on disconnect). io.send() silently drops all events — agent can't communicate with the new client.

File: connectonion/network/io/websocket.py — WebSocketIO._closed flag

3. Two _pipe_ws_io loops compete

The old loop (stuck waiting for agent_finished) and the new loop (from reattach) both reference the same agent_finished event, causing race conditions on completion.

File: connectonion/network/asgi/websocket.py — _pipe_ws_io()

Fix Plan

1. run_agent(): wrap in try/finally — always set agent_finished, capture error in error_holder
2. Reattach: reopen IO — reset io._closed = False so agent can send events through new WebSocket
3. Old _pipe_ws_io: detect superseded — when new connection reattaches, old pipe should exit cleanly

Documentation

• Source docs: docs/network/session-reconnect.md (Known Issue section)
• Docs site: https://docs.connectonion.com/session-reconnect (Known Issue section)
• Cross-references: docs/network/io.md, docs/network/session-websocket-lifecycle.md

[wu-changxing]

Added TODO comments in code referencing this issue:

• connectonion/network/asgi/websocket.py:261 — run_agent() needs try/finally
• connectonion/network/asgi/websocket.py:191 — reattach needs to reopen closed IO

Commits: dcca81a, a539b81
Docs updated: session-reconnect, websocket-protocol, io, session-websocket-lifecycle
Docs site deployed: https://docs.connectonion.com/session-reconnect

[wu-changxing]

🔍 Summary

WebSocket reconnection fails when a client refreshes during an approval block because three bugs compound: (1) run_agent() lacks try/finally so agent_finished never fires if the agent crashes, hanging _pipe_ws_io forever; (2) reattach reuses the old io object whose _closed=True flag silently drops all new events; (3) the old _pipe_ws_io loop and the new one race on the same agent_finished event.

📊 Current vs Proposed

CURRENT (Bug):

T+0  Agent blocks on io.receive() waiting for approval
T+5  Client refreshes → WebSocket disconnects
     │
     ▼
┌──────────────┐     ┌──────────────────┐
│ _pipe_ws_io  │────▶│ io.close()       │
│ (old loop)   │     │ _closed = True   │
└──────┬───────┘     │ puts io_closed   │
       │             └──────────────────┘
       │                     │
       │                     ▼
       │             ┌──────────────────┐
       │             │ Agent crashes on │
       │             │ io_closed sentinel│
       │             └──────────────────┘
       │                     │
       │                     ▼
       │             ┌──────────────────┐
       │             │ run_agent() dies │
       │             │ agent_finished   │
       │             │ NEVER set ❌     │
       │             └──────────────────┘
       │
       ▼
┌──────────────────────────┐
│ _pipe_ws_io hangs forever│  ← Bug 1: no try/finally
│ waiting agent_finished   │
└──────────────────────────┘

T+10 New WebSocket connects → CONNECT { session_id }
     │
     ▼
┌──────────────┐     ┌──────────────────┐
│ Reattach     │────▶│ io._closed=True  │ ← Bug 2: closed IO reused
│ uses old io  │     │ send() drops all │
└──────┬───────┘     └──────────────────┘
       │
       ▼
┌──────────────┐     ┌──────────────────┐
│ _pipe_ws_io  │     │ _pipe_ws_io      │ ← Bug 3: two loops
│ (NEW loop)   │     │ (OLD loop stuck) │    race on agent_finished
└──────────────┘     └──────────────────┘
       │
       ▼
  Client sees "No active connection" ❌

PROPOSED (Fix):

T+0  Agent blocks on io.receive() waiting for approval
T+5  Client refreshes → WebSocket disconnects
     │
     ▼
┌──────────────┐     ┌──────────────────┐
│ _pipe_ws_io  │────▶│ io.close()       │
│ (old loop)   │     │ _closed = True   │
└──────┬───────┘     └──────────────────┘
       │                     │
       │                     ▼
       │             ┌──────────────────┐
       │             │ Agent crashes on │
       │             │ io_closed sentinel│
       │             └──────────────────┘
       │                     │
       │                     ▼
       │             ┌──────────────────────┐
       │             │ run_agent() finally: │
       │             │ agent_finished.set() │ ✅ Fix 1
       │             │ error → error_holder │
       │             └──────────────────────┘
       │                     │
       ▼                     ▼
┌──────────────────────────┐
│ _pipe_ws_io exits cleanly│ ✅ agent_finished is set
└──────────────────────────┘

T+10 New WebSocket connects → CONNECT { session_id }
     │
     ▼
┌──────────────┐     ┌──────────────────────┐
│ Reattach     │────▶│ io.reopen()          │ ✅ Fix 2
│ reopens IO   │     │ _closed = False      │
└──────┬───────┘     │ drain stale sentinel │
       │             └──────────────────────┘
       │
       ▼
┌──────────────┐     ┌──────────────────┐
│ _pipe_ws_io  │     │ old loop already │ ✅ Fix 3
│ (NEW, only)  │     │ exited (fix 1)  │
└──────┬───────┘     └──────────────────┘
       │
       ▼
  Agent resends approval_needed → client approves ✅

📁 Files to Change

connectonion/
├── network/
│   ├── asgi/
│   │   └── websocket.py        ← Fix 1: try/finally in run_agent() (line 264)
│   │                             ← Fix 3: old _pipe_ws_io exits (line 378)
│   │                             ← Reattach calls io.reopen() (line 194)
│   └── io/
│       └── websocket.py        ← Fix 2: add reopen() method (new, ~8 lines)
tests/
└── unit/
    └── test_host_session.py    ← Add reconnection-during-approval test

🏗️ Fix Architecture

Fix 1: run_agent() try/finally
─────────────────────────────────
BEFORE (line 264-267):
    def run_agent():
        result_holder[0] = route_handlers["ws_input"](...)
        registry.mark_connected(session_id)
        agent_finished.set()          ← skipped on crash

AFTER:
    def run_agent():
        try:
            result_holder[0] = route_handlers["ws_input"](...)
            registry.mark_connected(session_id)
        except Exception as e:
            error_holder[0] = str(e)
        finally:
            agent_finished.set()      ← always fires


Fix 2: WebSocketIO.reopen()
─────────────────────────────────
NEW METHOD in websocket.py (after close()):

    def reopen(self):
        """Reopen IO after reconnection."""
        self._closed = False
        # Drain any stale io_closed sentinels
        while True:
            try:
                msg = self._incoming.get_nowait()
                if msg.get("type") != "io_closed":
                    self._incoming.put(msg)  # keep non-sentinel msgs
                    break
            except queue.Empty:
                break

CALL SITE in asgi/websocket.py (line 194):
    if status == "executing":
        io = active.io
        io.reopen()                   ← reset _closed flag
        agent_finished = active.agent_finished
        ...


Fix 3: _pipe_ws_io handles already-finished agent
─────────────────────────────────────────────────
The main polling loop (line 378-379) already checks
agent_finished.is_set(). With Fix 1 guaranteeing
agent_finished fires on crash, the old loop exits
naturally. No code change needed here — Fix 1 resolves it.

However, we should add a guard at _pipe_ws_io entry
for the edge case where agent finished between
reattach and pipe start:

    # At start of _pipe_ws_io (after line 316)
    if agent_finished.is_set():
        return False   # Agent already done, skip piping

🔄 Data Flow: Reconnection After Fix

Client refreshes during approval
         │
         ▼
[websocket.disconnect]
         │
         ▼
io.close() → _closed=True, sentinel in _incoming
         │
         ▼
Agent's io.receive() → gets io_closed sentinel
         │
         ▼
Agent raises ValueError("Connection closed...")
         │
         ▼
run_agent() catches in except → error_holder[0] = str(e)
         │
         ▼
finally: agent_finished.set() ✅
         │
         ▼
Old _pipe_ws_io sees agent_finished → exits cleanly
         │
         ▼
registry status = 'executing' (never got to mark_connected)
         │
         ▼ (seconds later)
New CONNECT arrives with same session_id
         │
         ▼
registry.get() → status='executing', agent_finished.is_set()=True
         │
         ▼
io.reopen() → _closed=False, drain sentinel
         │
         ▼
_pipe_ws_io → agent_finished already set → return False
         │
         ▼
Stored result from SessionStorage → send OUTPUT to client ✅

📝 Implementation Steps

1. Add reopen() to WebSocketIO (network/io/websocket.py:77)

   • Add reopen() method after close()
   • Reset _closed = False
   • Drain stale io_closed sentinels from _incoming queue
   • Keep any non-sentinel messages

2. Wrap run_agent() in try/finally (network/asgi/websocket.py:264)

   • Move registry.mark_connected() into try block
   • Capture exception to error_holder[0] in except
   • Move agent_finished.set() to finally block
   • Remove the TODO(WebSocket reconnection fails during approval blocks #119) comment

3. Call io.reopen() on reattach (network/asgi/websocket.py:194)

   • After io = active.io, call io.reopen()
   • Remove the TODO(WebSocket reconnection fails during approval blocks #119) comment

4. Add early-exit guard in _pipe_ws_io (network/asgi/websocket.py:316)

   • If agent_finished.is_set() at entry, return False immediately
   • Prevents spawning tasks for an already-finished agent

5. Add test for reconnection during approval (tests/unit/test_host_session.py)

   • Simulate: register session → close IO → set agent_finished → reopen → verify send works
   • Verify reopen() clears _closed and drains sentinel
   • Verify _pipe_ws_io early-exits when agent already finished

⚠️ Risks & Considerations

• Agent mid-tool-call: If agent crashes mid-tool (not at io.receive()), the tool execution state is lost. This is acceptable — the fix ensures clean recovery, not resumption of partial work.
• Queue ordering: reopen() drains only io_closed sentinels. Any real messages queued before disconnect are preserved. Order is maintained by FIFO queue semantics.
• Thread safety: _closed is a simple boolean read/written from one thread (agent thread for send(), websocket handler for close()/reopen()). Python's GIL makes this safe without explicit locking.
• No behavior change for non-reconnect flows: The try/finally only adds guaranteed cleanup — normal flows where agent completes successfully are unchanged.

🧪 Testing Strategy

• Unit test reopen(): Create WebSocketIO → close() → verify send() drops → reopen() → verify send() works again
• Unit test sentinel drain: close() puts sentinel → reopen() drains it → receive() doesn't get stale sentinel
• Unit test run_agent() crash path: Mock route_handler to raise → verify agent_finished is set and error_holder captured
• Unit test _pipe_ws_io early exit: Pass already-set agent_finished Event → verify returns False immediately
• Integration test: Full reconnect flow — register session, simulate disconnect during approval, reconnect, verify stored result delivered