daniel156161
0c005bc119
Requiring an explicit tab_id forced every navigate or close through a preceding tabs_list call, which costs an MCP client a full round trip just to learn the ID the browser already considers current. navigate and tabs_close now resolve the active tab when tab_id is omitted, matching the screenshot tool. Resolution is explicit rather than forwarding None into the SDK, so the acting tool knows which tab it touched; tabs_close reports it, since closing the wrong tab is not recoverable. This stays in the MCP layer: the SDK and CLI signatures are unchanged.
127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
"""Per-process pool of authenticated remote connections for reuse.
|
|
|
|
A ``browser-cli serve`` connection stays open after its first (encrypted)
|
|
command, so the client can send further commands over it without re-running the
|
|
TCP/TLS/challenge/auth handshake (~hundreds of ms each). Only encrypted (PQ)
|
|
sessions are pooled — plaintext/legacy sessions stay one-shot, matching the
|
|
server, which only loops for encrypted sessions.
|
|
|
|
Connections are checked out exclusively (never shared between threads at once),
|
|
returned on success, and dropped on any I/O error or once older than an idle
|
|
bound (kept below the server's idle timeout so we don't reuse a connection the
|
|
server has already closed).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import json
|
|
import socket
|
|
import threading
|
|
import time
|
|
|
|
from browser_cli.constants import REMOTE_SESSION_IDLE_TIMEOUT
|
|
from browser_cli.framing import frame
|
|
|
|
# Retire a pooled connection a few seconds before the server would, so we never
|
|
# hand back one the server has just timed out and closed.
|
|
_MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5)
|
|
_MAX_PER_ENDPOINT = 8
|
|
_MAX_ENDPOINTS = 64
|
|
|
|
class PooledConnection:
|
|
__slots__ = ("sock", "secret", "last_used")
|
|
|
|
def __init__(self, sock: socket.socket, secret: bytes) -> None:
|
|
self.sock = sock
|
|
self.secret = secret
|
|
self.last_used = time.monotonic()
|
|
|
|
_POOL: dict[str, list[PooledConnection]] = {}
|
|
_LOCK = threading.Lock()
|
|
|
|
def _close(sock: socket.socket) -> None:
|
|
try:
|
|
sock.close()
|
|
except OSError:
|
|
pass
|
|
|
|
def checkout(endpoint: str) -> PooledConnection | None:
|
|
"""Take an idle authenticated connection for *endpoint*, or None."""
|
|
now = time.monotonic()
|
|
with _LOCK:
|
|
conns = _POOL.get(endpoint)
|
|
while conns:
|
|
conn = conns.pop()
|
|
if now - conn.last_used <= _MAX_IDLE_SECONDS:
|
|
return conn
|
|
_close(conn.sock) # too old — assume the server has dropped it
|
|
return None
|
|
|
|
def _prune_endpoints_locked(now: float) -> None:
|
|
"""Keep the number of endpoint buckets bounded for long-running SDK users."""
|
|
for endpoint, bucket in list(_POOL.items()):
|
|
fresh = [conn for conn in bucket if now - conn.last_used <= _MAX_IDLE_SECONDS]
|
|
if fresh:
|
|
_POOL[endpoint] = fresh
|
|
else:
|
|
for conn in bucket:
|
|
_close(conn.sock)
|
|
_POOL.pop(endpoint, None)
|
|
|
|
while len(_POOL) >= _MAX_ENDPOINTS:
|
|
oldest_endpoint, bucket = min(
|
|
_POOL.items(),
|
|
key=lambda item: min(conn.last_used for conn in item[1]) if item[1] else 0.0,
|
|
)
|
|
for conn in bucket:
|
|
_close(conn.sock)
|
|
_POOL.pop(oldest_endpoint, None)
|
|
|
|
def checkin(endpoint: str, conn: PooledConnection) -> None:
|
|
"""Return a still-healthy connection to the pool for reuse."""
|
|
conn.last_used = time.monotonic()
|
|
with _LOCK:
|
|
if endpoint not in _POOL and len(_POOL) >= _MAX_ENDPOINTS:
|
|
_prune_endpoints_locked(conn.last_used)
|
|
bucket = _POOL.setdefault(endpoint, [])
|
|
if len(bucket) >= _MAX_PER_ENDPOINT:
|
|
_close(conn.sock)
|
|
return
|
|
bucket.append(conn)
|
|
|
|
def discard(conn: PooledConnection) -> None:
|
|
"""Drop a connection that errored or is no longer usable."""
|
|
_close(conn.sock)
|
|
|
|
def close_all() -> None:
|
|
"""Close every pooled connection (process exit / test isolation)."""
|
|
with _LOCK:
|
|
for bucket in _POOL.values():
|
|
for conn in bucket:
|
|
_close(conn.sock)
|
|
_POOL.clear()
|
|
|
|
def session_inner_message(msg: dict) -> dict:
|
|
"""Strip auth/transport fields, leaving the command for an established session."""
|
|
keep = {"id", "command", "args", "user_agent", "accept_encoding", "_route", "_suppress_pq_warning"}
|
|
return {k: v for k, v in msg.items() if k in keep}
|
|
|
|
def send_over(conn: PooledConnection, msg: dict) -> bytes | None:
|
|
"""Send one command over an existing encrypted session. Raises on I/O error."""
|
|
from browser_cli.auth import pq_encrypt
|
|
from browser_cli.remote.socket import recv_all
|
|
from browser_cli.remote.transport import _decode_pq_response
|
|
|
|
inner = json.dumps(session_inner_message(msg)).encode("utf-8")
|
|
envelope = json.dumps({"encrypted": pq_encrypt(conn.secret, "request", inner)}).encode("utf-8")
|
|
conn.sock.sendall(frame(envelope))
|
|
response = recv_all(conn.sock)
|
|
if not response:
|
|
# EOF — an older server (no session loop) closed after one command. Treat as
|
|
# a transport failure so the caller re-handshakes; never as an app error,
|
|
# which could double-execute a non-idempotent command on retry.
|
|
raise EOFError("remote closed the pooled connection")
|
|
return _decode_pq_response(response, conn.secret)
|
|
|
|
atexit.register(close_all)
|