Skip to content

feat(flow): support multiple extension profiles with failover routing #22

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
KumarLakshmanan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(flow): support multiple extension profiles with failover routing #22

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 8 additions & 5 deletions agent/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,15 +47,15 @@ async def ws_handler(websocket):
async for raw in websocket:
try:
data = json.loads(raw)
await client.handle_message(data)
await client.handle_message(data, websocket)
except json.JSONDecodeError:
logger.warning("Invalid JSON from extension")
except Exception as e:
logger.exception("Error handling extension message: %s", e)
except websockets.ConnectionClosed:
pass
finally:
client.clear_extension()
client.clear_extension(websocket)
logger.info("Extension disconnected")


Expand Down Expand Up @@ -90,9 +90,12 @@ async def lifespan(app: FastAPI):

controller = get_worker_controller()

# SIGTERM handler for graceful shutdown
loop = asyncio.get_event_loop()
loop.add_signal_handler(signal.SIGTERM, controller.request_shutdown)
# SIGTERM handler for graceful shutdown (Unix only)
try:
loop = asyncio.get_event_loop()
loop.add_signal_handler(signal.SIGTERM, controller.request_shutdown)
except (NotImplementedError, AttributeError):
pass

# Start background tasks
ws_task = asyncio.create_task(run_ws_server())
Expand Down
206 changes: 170 additions & 36 deletions agent/services/flow_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,10 @@ class FlowClient:
"""Sends commands to Chrome extension via WebSocket."""

def __init__(self):
self._extension_ws = None # Set by WS server when extension connects
self._extension_ws = None # Active authenticated extension connection
self._extensions: dict[object, dict] = {}
self._pending: dict[str, asyncio.Future] = {}
self._pending_ws: dict[str, object] = {}
self._flow_key: Optional[str] = None
# WS stats
self._ws_connect_count = 0
Expand All @@ -35,31 +37,122 @@ def __init__(self):

def set_extension(self, ws):
"""Called when extension connects via WS."""
self._extension_ws = ws
self._extensions[ws] = {
"connected_at": time.time(),
"flow_key": None,
"token_captured_at": None,
"unavailable_until": 0,
}
# A new unauthenticated profile must not displace an already
# authenticated extension. It becomes active after token_captured.
if self._extension_ws is None:
self._extension_ws = ws
self._ws_connect_count += 1
self._ws_connected_at = time.time()
logger.info("Extension connected #%d (waiting for extension_ready/token_captured to sync)", self._ws_connect_count)

def clear_extension(self):
logger.info(
"Extension connected #%d (%d active connection(s)); "
"waiting for extension_ready/token_captured to sync",
self._ws_connect_count,
len(self._extensions),
)

def clear_extension(self, ws=None):
"""Called when extension disconnects."""
self._extension_ws = None
disconnected_ws = ws or self._extension_ws
if disconnected_ws is None:
return

self._extensions.pop(disconnected_ws, None)
self._ws_disconnect_count += 1
self._ws_last_disconnect_at = time.time()
# Cancel all pending futures (copy to avoid RuntimeError on concurrent modification)
pending_copy = list(self._pending.items())
count = len(pending_copy)
for req_id, future in pending_copy:
if not future.done():

# Only cancel requests that were sent through the disconnected socket.
# Requests owned by other Chrome profiles are still valid.
disconnected_pending = [
(req_id, self._pending.get(req_id))
for req_id, pending_ws in list(self._pending_ws.items())
if pending_ws is disconnected_ws
]
for req_id, future in disconnected_pending:
if future is not None and not future.done():
future.set_exception(ConnectionError("Extension disconnected"))
self._pending.clear()
logger.warning("Extension disconnected, cleared %d pending requests", count)
self._pending_ws.pop(req_id, None)

if self._extension_ws is disconnected_ws:
self._extension_ws = self._select_extension(require_token=True)
if self._extension_ws is None:
self._extension_ws = self._select_extension(require_token=False)

active_session = self._extensions.get(self._extension_ws, {})
self._flow_key = active_session.get("flow_key")
logger.warning(
"Extension disconnected, cancelled %d owned request(s); "
"%d extension connection(s) remain",
len(disconnected_pending),
len(self._extensions),
)

def _extension_candidates(self, require_token: bool):
"""Return usable extensions in preferred routing order."""
now = time.time()
candidates = []
for ws, session in self._extensions.items():
if require_token and not session.get("flow_key"):
continue
recency = (
session.get("token_captured_at")
if require_token
else session.get("connected_at")
)
candidates.append({
"ws": ws,
"available": session.get("unavailable_until", 0) <= now,
"active": ws is self._extension_ws,
"recency": recency or 0,
})

# Prefer an available active session, then the most recently
# authenticated alternatives. Temporarily unavailable sessions remain
# last-resort candidates so a single-profile setup can still recover.
candidates.sort(
key=lambda item: (
item["available"],
item["active"] and item["available"],
item["recency"],
),
reverse=True,
)
return [item["ws"] for item in candidates]

def _select_extension(self, require_token: bool):
"""Choose the preferred authenticated or connected extension."""
candidates = self._extension_candidates(require_token)
return candidates[0] if candidates else None

@staticmethod
def _should_failover(result: dict) -> bool:
"""Return true for profile-local failures that another tab can solve."""
message = str(result.get("error") or result.get("data") or "").lower()
return any(marker in message for marker in (
"no_flow_key",
"no_flow_tab",
"no current window",
"extension not connected",
"extension disconnected",
"extension_switched",
"public_error_per_model_daily_quota_reached",
"public_error_user_quota_reached",
))

def set_flow_key(self, key: str):
self._flow_key = key
if self._extension_ws in self._extensions:
self._extensions[self._extension_ws]["flow_key"] = key
self._extensions[self._extension_ws]["token_captured_at"] = time.time()

@property
def connected(self) -> bool:
return self._extension_ws is not None
return bool(self._extensions)

@property
def ws_stats(self) -> dict:
Expand All @@ -68,15 +161,26 @@ def ws_stats(self) -> dict:
uptime = int(time.time() - self._ws_connected_at)
return {
"connected": self.connected,
"active_connections": len(self._extensions),
"authenticated_connections": sum(
1 for session in self._extensions.values()
if session.get("flow_key")
),
"connects": self._ws_connect_count,
"disconnects": self._ws_disconnect_count,
"uptime_s": uptime,
}

async def handle_message(self, data: dict):
async def handle_message(self, data: dict, websocket=None):
"""Handle incoming message from extension."""
if data.get("type") == "token_captured":
self._flow_key = data.get("flowKey")
key = data.get("flowKey")
source_ws = websocket or self._extension_ws
if source_ws is not None and source_ws in self._extensions:
self._extensions[source_ws]["flow_key"] = key
self._extensions[source_ws]["token_captured_at"] = time.time()
self._extension_ws = source_ws
self._flow_key = key
logger.info("Flow key captured from extension")
asyncio.create_task(self._sync_tier())
return
Expand All @@ -95,8 +199,9 @@ async def handle_message(self, data: dict):

if data.get("type") == "ping":
# Respond to keepalive
if self._extension_ws:
await self._extension_ws.send(json.dumps({"type": "pong"}))
target_ws = websocket or self._extension_ws
if target_ws:
await target_ws.send(json.dumps({"type": "pong"}))
return

# Response to a pending request
Expand Down Expand Up @@ -212,27 +317,56 @@ async def _send(self, method: str, params: dict, timeout: float = 300) -> dict:
must check result.get("error") or use _is_ws_error() before reading data.
Never raises; exceptions are caught and returned as error dicts.
"""
if not self._extension_ws:
if not self.connected:
return {"error": "Extension not connected"}

req_id = str(uuid.uuid4())
future = asyncio.get_running_loop().create_future()
self._pending[req_id] = future
extension_candidates = self._extension_candidates(require_token=True)
if not extension_candidates:
return {"error": "NO_FLOW_KEY"}

try:
await self._extension_ws.send(json.dumps({
"id": req_id,
"method": method,
"params": params,
}))
result = await asyncio.wait_for(future, timeout=timeout)
return result
except asyncio.TimeoutError:
return {"error": f"Timeout ({timeout}s) waiting for {method}"}
except Exception as e:
return {"error": str(e)}
finally:
self._pending.pop(req_id, None)
last_result = {"error": "Extension not connected"}
for index, extension_ws in enumerate(extension_candidates):
if extension_ws not in self._extensions:
continue

self._extension_ws = extension_ws
self._flow_key = self._extensions[extension_ws].get("flow_key")
req_id = str(uuid.uuid4())
future = asyncio.get_running_loop().create_future()
self._pending[req_id] = future
self._pending_ws[req_id] = extension_ws

try:
await extension_ws.send(json.dumps({
"id": req_id,
"method": method,
"params": params,
}))
last_result = await asyncio.wait_for(future, timeout=timeout)
except asyncio.TimeoutError:
last_result = {"error": f"Timeout ({timeout}s) waiting for {method}"}
except Exception as e:
last_result = {"error": str(e)}
finally:
self._pending.pop(req_id, None)
self._pending_ws.pop(req_id, None)

has_alternative = index + 1 < len(extension_candidates)
if self._should_failover(last_result) and has_alternative:
if extension_ws in self._extensions:
self._extensions[extension_ws]["unavailable_until"] = (
time.time() + 60
)
logger.warning(
"Extension profile unavailable for %s; retrying through "
"another authenticated profile",
method,
)
continue

return last_result

return last_result

def _build_url(self, endpoint_key: str, **kwargs) -> str:
"""Build full API URL."""
Expand Down