diff --git a/README.md b/README.md index 77a1bd2..0d7b804 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,11 @@ The server is stateless at the MCP layer. Every tool call creates a fresh `browser`, `remote`, and `key` on a tool call when a specific local profile or authenticated browser-cli remote is required. +`browser_navigate`, `browser_tabs_close`, and `browser_screenshot` take an +optional `tab_id` and act on the active tab when it is omitted, so a caller +does not need a preceding `browser_tabs_list` round trip. `browser_tabs_close` +reports the tab it closed. + Available tools: - `browser_tabs_list`, `browser_tabs_open`, `browser_tabs_close` - `browser_navigate`, `browser_page_info` diff --git a/browser_cli/mcp/server.py b/browser_cli/mcp/server.py index 64f6f72..48ba0df 100644 --- a/browser_cli/mcp/server.py +++ b/browser_cli/mcp/server.py @@ -16,6 +16,7 @@ from urllib.parse import urlsplit from browser_cli import BrowserCLI from browser_cli.mcp.naming import resolve_tool_prefix, tool_name from browser_cli.mcp.serialization import structured +from browser_cli.mcp.targets import resolve_tab_id ClientFactory = Callable[..., BrowserCLI] @@ -93,27 +94,29 @@ def create_server(*, client_factory: ClientFactory = BrowserCLI, tool_prefix: st @mcp.tool(name=tool_name("tabs_close", prefix)) def tabs_close( - tab_id: int, + tab_id: int | None = None, browser: str | None = None, remote: str | None = None, key: str | None = None, ) -> dict[str, int]: - """Close one tab by its current ID. This changes the user's real browser.""" - closed = _client(client_factory, browser, remote, key).tabs.close(tab_id) - return {"closed": closed} + """Close a tab, defaulting to the active tab. This changes the real browser.""" + client = _client(client_factory, browser, remote, key) + target = resolve_tab_id(client, tab_id) + return {"closed": client.tabs.close(target), "tab_id": target} @mcp.tool(name=tool_name("navigate", prefix)) def navigate( - tab_id: int, url: str, + tab_id: int | None = None, browser: str | None = None, remote: str | None = None, key: str | None = None, ) -> dict[str, Any]: - """Navigate an existing tab to a URL and return a fresh tab snapshot.""" + """Navigate a tab to a URL, defaulting to the active tab, and return it.""" client = _client(client_factory, browser, remote, key) - client.nav.to(tab_id, url) - return structured(client.tabs.status(tab_id)) + target = resolve_tab_id(client, tab_id) + client.nav.to(target, url) + return structured(client.tabs.status(target)) @mcp.tool(name=tool_name("page_info", prefix)) def page_info( diff --git a/browser_cli/mcp/targets.py b/browser_cli/mcp/targets.py new file mode 100644 index 0000000..bbbc5e8 --- /dev/null +++ b/browser_cli/mcp/targets.py @@ -0,0 +1,16 @@ +"""Tab targeting for the MCP surface. + +MCP callers pay a full round trip for every extra tool call, so tools that act +on a tab accept an optional ``tab_id`` and fall back to the browser's current +active tab. Resolution happens here rather than by forwarding ``None`` into the +SDK, so the acting tool always knows which tab it touched and can report it. +""" +from __future__ import annotations + +from browser_cli import BrowserCLI + +def resolve_tab_id(client: BrowserCLI, tab_id: int | None) -> int: + """Return *tab_id*, or the ID of the currently active tab when it is ``None``.""" + if tab_id is not None: + return tab_id + return client.tabs.active().id diff --git a/browser_cli/remote/pool.py b/browser_cli/remote/pool.py index e00bbec..f60bc38 100644 --- a/browser_cli/remote/pool.py +++ b/browser_cli/remote/pool.py @@ -26,6 +26,7 @@ from browser_cli.framing import frame # 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") @@ -56,10 +57,32 @@ def checkout(endpoint: str) -> PooledConnection | None: _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) diff --git a/browser_cli/serve/security.py b/browser_cli/serve/security.py index 65c181d..5ab67ad 100644 --- a/browser_cli/serve/security.py +++ b/browser_cli/serve/security.py @@ -70,19 +70,48 @@ class RateLimiter: ``rate`` is the sustained refill in tokens/second; ``burst`` is the bucket capacity (defaults to ``rate``). ``rate <= 0`` disables limiting entirely. Thread-safe so it can be shared across all connections of one serve process. + + The bucket table is capped. Without that bound, a long-running public server + could retain one entry per ever-seen identity/IP forever; GC cannot reclaim + those entries because the limiter still references them. """ - def __init__(self, rate: float, burst: float | None = None) -> None: + def __init__(self, rate: float, burst: float | None = None, max_buckets: int = 4096) -> None: self.rate = float(rate) self.capacity = float(burst) if burst is not None else max(float(rate), 1.0) + self.max_buckets = max(1, int(max_buckets)) self._buckets: dict[str, tuple[float, float]] = {} self._lock = threading.Lock() + def _prune_locked(self, now: float) -> None: + """Drop idle/full buckets, then oldest buckets, until the table is bounded.""" + if len(self._buckets) < self.max_buckets or self.rate <= 0: + return + + # Once a bucket has fully refilled, keeping it around carries no useful + # throttling state. Use at least 60s so normal active identities are not + # churned out aggressively on high-rate configs. + idle_seconds = max(60.0, (self.capacity / self.rate) * 2) + full_epsilon = 1e-9 + for bucket_key, (tokens, last) in list(self._buckets.items()): + refilled = min(self.capacity, tokens + (now - last) * self.rate) + if refilled >= self.capacity - full_epsilon and now - last >= idle_seconds: + self._buckets.pop(bucket_key, None) + + # If an attacker keeps creating fresh identities faster than they go idle, + # still keep memory bounded. Evict the oldest identity state; that may reset + # throttling for that identity, but bounded memory is more important here. + while len(self._buckets) >= self.max_buckets: + oldest_key = min(self._buckets, key=lambda k: self._buckets[k][1]) + self._buckets.pop(oldest_key, None) + def allow(self, key: str) -> bool: if self.rate <= 0: return True now = time.monotonic() with self._lock: + if key not in self._buckets and len(self._buckets) >= self.max_buckets: + self._prune_locked(now) tokens, last = self._buckets.get(key, (self.capacity, now)) tokens = min(self.capacity, tokens + (now - last) * self.rate) if tokens < 1.0: diff --git a/extension/src/classes/JobManager.ts b/extension/src/classes/JobManager.ts index b4517c0..bc0a162 100644 --- a/extension/src/classes/JobManager.ts +++ b/extension/src/classes/JobManager.ts @@ -15,6 +15,11 @@ import type { Job, Serializable, ErrorLike, DispatchArgs } from '../types'; // jobs only need to survive long enough for the CLI to poll their result. export const MAX_FINISHED_JOBS = 20; +// Cap simultaneously running background jobs. A hung job has a watchdog, but a +// command flood could still pin many timers/results for up to JOB_TIMEOUT_MS. +// Rejecting above this bound keeps service-worker memory predictable. +export const MAX_RUNNING_JOBS = 32; + // Watchdog: if a runner never resolves/rejects (e.g. executeScript against a // dead tab), finalize the job as an error so its persist interval stops instead // of writing to api.storage.local every second forever. @@ -77,6 +82,11 @@ export class JobManager { } async start(command: string, args: DispatchArgs, runner: JobRunner) { + const runningCount = [...this.jobs.values()].filter(job => job.status === "running").length; + if (runningCount >= MAX_RUNNING_JOBS) { + throw new Error(`too many background jobs running (${runningCount}); wait for jobs to finish or cancel one`); + } + const jobId = `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`; const job: Job = { id: jobId, diff --git a/extension/src/commands/perf.ts b/extension/src/commands/perf.ts index ea86e84..ca4fe45 100644 --- a/extension/src/commands/perf.ts +++ b/extension/src/commands/perf.ts @@ -1,7 +1,7 @@ import { getLargeOperationThrottle, getPerformanceProfile, hasAudibleTabs, setPerformanceProfile } from '../core'; import { CommandGroup } from '../classes/CommandGroup'; import type { CommandEntry } from '../classes/CommandGroup'; -import type { PerfSetProfileArgs, JobIdArgs } from '../types'; +import type { Job, PerfSetProfileArgs, JobIdArgs } from '../types'; // PerfCommands also owns the jobs.* status/cancel queries: they read the same // JobManager (ctx.jobs) that perf.status reports, and there is no dedicated @@ -15,6 +15,19 @@ export class PerfCommands extends CommandGroup { "jobs.cancel": (a: JobIdArgs) => this.ctx.jobs.cancel(a), }; + private jobSummary(job: Job) { + return { + id: job.id, + command: job.command, + status: job.status, + phase: job.phase, + current: job.current, + total: job.total, + percent: job.percent, + cancelRequested: job.cancelRequested, + }; + } + private async perfStatus() { const profile = await getPerformanceProfile(); const audible = await hasAudibleTabs(); @@ -23,16 +36,7 @@ export class PerfCommands extends CommandGroup { performanceProfile: profile, audible, throttle, - jobs: this.ctx.jobs.list().map(job => ({ - id: job.id, - command: job.command, - status: job.status, - phase: job.phase, - current: job.current, - total: job.total, - percent: job.percent, - cancelRequested: job.cancelRequested, - })), + jobs: this.ctx.jobs.list().map(job => this.jobSummary(job)), }; } } diff --git a/extension/src/commands/windows.ts b/extension/src/commands/windows.ts index 761b66d..3c414e3 100644 --- a/extension/src/commands/windows.ts +++ b/extension/src/commands/windows.ts @@ -14,9 +14,31 @@ export class WindowsCommands extends CommandGroup { "windows.open": (a: WindowsOpenArgs) => this.windowsOpen(a), }; + private async activeWindowIds(): Promise> { + const windows = await api.windows.getAll({}); + return new Set(windows.map(w => w.id).filter(id => typeof id === "number")); + } + + private async pruneAliases(activeIds?: Set): Promise> { + const aliases = await getAliases(); + const liveIds = activeIds || await this.activeWindowIds(); + const pruned: Record = {}; + let changed = false; + for (const [id, alias] of Object.entries(aliases)) { + if (liveIds.has(Number(id))) { + pruned[id] = alias; + } else { + changed = true; + } + } + if (changed) await api.storage.local.set({ windowAliases: pruned }); + return pruned; + } + private async windowsList() { const windows = await api.windows.getAll({ populate: true }); - const aliases = await getAliases(); + const activeIds = new Set(windows.map(w => w.id).filter(id => typeof id === "number")); + const aliases = await this.pruneAliases(activeIds); return windows.map(w => ({ id: w.id, alias: aliases[w.id] || null, @@ -27,7 +49,7 @@ export class WindowsCommands extends CommandGroup { } private async windowsRename({ windowId, name }: WindowsRenameArgs) { - const aliases = await getAliases(); + const aliases = await this.pruneAliases(); aliases[windowId] = name; await api.storage.local.set({ windowAliases: aliases }); return { windowId, name }; @@ -35,6 +57,11 @@ export class WindowsCommands extends CommandGroup { private async windowsClose({ windowId }: WindowsCloseArgs) { await api.windows.remove(windowId); + const aliases = await this.pruneAliases(); + if (windowId in aliases) { + delete aliases[windowId]; + await api.storage.local.set({ windowAliases: aliases }); + } return { windowId }; } diff --git a/extension/test/jobs.test.ts b/extension/test/jobs.test.ts index f5b8dac..3bc0fa9 100644 --- a/extension/test/jobs.test.ts +++ b/extension/test/jobs.test.ts @@ -1,7 +1,7 @@ // @ts-nocheck import { test, mock } from 'node:test'; import assert from 'node:assert/strict'; -import { JobManager, JOB_TIMEOUT_MS, MAX_FINISHED_JOBS, pruneFinishedJobs } from '../src/classes/JobManager'; +import { JobManager, JOB_TIMEOUT_MS, MAX_FINISHED_JOBS, MAX_RUNNING_JOBS, pruneFinishedJobs } from '../src/classes/JobManager'; import { makeChromeMock } from './chrome-mock'; // Drain pending microtasks (finalize() chains several awaits). setImmediate is @@ -129,6 +129,21 @@ test('JobManager: a runner that settles after the watchdog cannot resurrect the mock.timers.reset(); }); +test('JobManager: rejects new background jobs above the running-job cap', async () => { + mock.timers.enable({ apis: ['setInterval', 'setTimeout'] }); + globalThis.chrome = makeChromeMock(); + const mgr = new JobManager(); + for (let i = 0; i < MAX_RUNNING_JOBS; i++) { + await mgr.start(`running${i}`, {}, () => new Promise(() => {})); + } + await assert.rejects( + () => mgr.start('overflow', {}, async () => 'nope'), + /too many background jobs running/, + ); + assert.equal(mgr.list().filter(job => job.status === 'running').length, MAX_RUNNING_JOBS); + mock.timers.reset(); +}); + test('JobManager: persisted set keeps running jobs even past the finished cap', async () => { mock.timers.enable({ apis: ['setInterval', 'setTimeout'] }); globalThis.chrome = makeChromeMock(); diff --git a/extension/test/windows.test.ts b/extension/test/windows.test.ts new file mode 100644 index 0000000..1ef41e6 --- /dev/null +++ b/extension/test/windows.test.ts @@ -0,0 +1,69 @@ +// @ts-nocheck +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import { WindowsCommands } from '../src/commands/windows'; +import { makeChromeMock } from './chrome-mock'; + +function makeWindowsChromeMock(windows) { + const chrome = makeChromeMock(); + chrome.windows = { + getAll: async () => windows, + remove: async () => {}, + create: async () => ({ id: 99 }), + }; + return chrome; +} + +test('windows.list prunes aliases for closed windows', async () => { + globalThis.chrome = makeWindowsChromeMock([ + { id: 1, focused: true, state: 'normal', tabs: [{ id: 10 }] }, + { id: 2, focused: false, state: 'minimized', tabs: [] }, + ]); + globalThis.chrome.storage.local._store.windowAliases = { + 1: 'main', + 2: 'side', + 999: 'closed', + }; + + const commands = new WindowsCommands({ jobs: {} }); + const result = await commands.commands['windows.list']({}); + + assert.deepEqual(result.map(w => [w.id, w.alias]), [[1, 'main'], [2, 'side']]); + assert.deepEqual(globalThis.chrome.storage.local._store.windowAliases, { 1: 'main', 2: 'side' }); +}); + +test('windows.rename prunes stale aliases before saving the new name', async () => { + globalThis.chrome = makeWindowsChromeMock([ + { id: 1, focused: true, state: 'normal', tabs: [] }, + { id: 2, focused: false, state: 'normal', tabs: [] }, + ]); + globalThis.chrome.storage.local._store.windowAliases = { + 1: 'main', + 999: 'closed', + }; + + const commands = new WindowsCommands({ jobs: {} }); + await commands.commands['windows.rename']({ windowId: 2, name: 'work' }); + + assert.deepEqual(globalThis.chrome.storage.local._store.windowAliases, { 1: 'main', 2: 'work' }); +}); + +test('windows.close removes the closed window alias immediately', async () => { + let removed = null; + globalThis.chrome = makeWindowsChromeMock([ + { id: 1, focused: true, state: 'normal', tabs: [] }, + { id: 2, focused: false, state: 'normal', tabs: [] }, + ]); + globalThis.chrome.windows.remove = async id => { removed = id; }; + globalThis.chrome.storage.local._store.windowAliases = { + 1: 'main', + 2: 'side', + 999: 'closed', + }; + + const commands = new WindowsCommands({ jobs: {} }); + await commands.commands['windows.close']({ windowId: 2 }); + + assert.equal(removed, 2); + assert.deepEqual(globalThis.chrome.storage.local._store.windowAliases, { 1: 'main' }); +}); diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 6185da8..5e0fea0 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -26,6 +26,7 @@ class FakeClient: list=self.tabs_list, open=self.tabs_open, close=self.tabs_close, + active=self.tabs_active, status=self.tabs_status, screenshot=self.tabs_screenshot, ) @@ -46,6 +47,10 @@ class FakeClient: self.calls.append(("close", tab_id)) return 1 + def tabs_active(self): + self.calls.append(("active",)) + return SimpleNamespace(id=7) + def tabs_status(self, tab_id): self.calls.append(("status", tab_id)) return {"id": tab_id, "title": "Navigated", "url": "https://example.com/next"} @@ -178,6 +183,18 @@ async def test_mutating_tools_use_sdk_and_return_fresh_state(client): ("navigate", 8, "https://example.com/next"), ("status", 8) ] +@pytest.mark.anyio +async def test_tab_tools_default_to_the_active_tab(client): + navigated = await client.call_tool("browser_navigate", {"url": "https://example.com/next"}) + closed = await client.call_tool("browser_tabs_close", {}) + + assert navigated.structured_content["id"] == 7 + assert closed.structured_content == {"closed": 1, "tab_id": 7} + assert FakeClient.instances[0].calls == [ + ("active",), ("navigate", 7, "https://example.com/next"), ("status", 7) + ] + assert FakeClient.instances[1].calls == [("active",), ("close", 7)] + @pytest.mark.anyio async def test_screenshot_returns_image_content(client): result = await client.call_tool("browser_screenshot", {"tab_id": 7, "format": "png"}) diff --git a/tests/test_remote_pool.py b/tests/test_remote_pool.py index 5c9e776..220e226 100644 --- a/tests/test_remote_pool.py +++ b/tests/test_remote_pool.py @@ -43,6 +43,39 @@ def test_checkin_caps_pool_size(): b.close() pool.close_all() +def test_checkin_caps_endpoint_buckets(): + pool.close_all() + peers = [] + try: + for i in range(pool._MAX_ENDPOINTS + 5): + a, b = _socketpair() + peers.append(b) + pool.checkin(f"host-{i}:443", pool.PooledConnection(a, b"secret")) + assert len(pool._POOL) <= pool._MAX_ENDPOINTS + finally: + for peer in peers: + peer.close() + pool.close_all() + +def test_checkin_prunes_stale_endpoint_buckets(): + pool.close_all() + old_a, old_b = _socketpair() + old = pool.PooledConnection(old_a, b"secret") + pool.checkin("old:443", old) + old.last_used -= pool._MAX_IDLE_SECONDS + 1 + peers = [old_b] + try: + for i in range(pool._MAX_ENDPOINTS): + a, b = _socketpair() + peers.append(b) + pool.checkin(f"new-{i}:443", pool.PooledConnection(a, b"secret")) + assert "old:443" not in pool._POOL + assert len(pool._POOL) <= pool._MAX_ENDPOINTS + finally: + for peer in peers: + peer.close() + pool.close_all() + def test_session_inner_message_strips_auth_fields(): msg = { "id": "1", "command": "tabs.list", "args": {}, "user_agent": "browser-cli/1", diff --git a/tests/test_serve_security.py b/tests/test_serve_security.py index 89384bd..b3d2330 100644 --- a/tests/test_serve_security.py +++ b/tests/test_serve_security.py @@ -181,6 +181,23 @@ def test_rate_limiter_is_per_key(): assert limiter.allow("a") is False assert limiter.allow("b") is False +def test_rate_limiter_caps_identity_buckets(): + limiter = RateLimiter(rate=0.0001, burst=1, max_buckets=3) + for i in range(10): + assert limiter.allow(f"key-{i}") is True + assert len(limiter._buckets) <= 3 + +def test_rate_limiter_prunes_refilled_idle_buckets(monkeypatch): + current = 1000.0 + monkeypatch.setattr("browser_cli.serve.security.time.monotonic", lambda: current) + limiter = RateLimiter(rate=1, burst=2, max_buckets=2) + assert limiter.allow("old") is True + current += 120.0 + assert limiter.allow("a") is True + assert limiter.allow("b") is True + assert "old" not in limiter._buckets + assert len(limiter._buckets) <= 2 + # ── ServeSecurity ──────────────────────────────────────────────────────────────── def test_effective_policy_prefers_per_key_override():