diff --git a/CHANGELOG.md b/CHANGELOG.md index 7bf72a31a..52cb1e83b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,20 @@ Two things are versioned separately from this file and worth knowing about: ## [Unreleased] +### Fixed + +- **The Workspaces page no longer sits on its loading state before it draws.** + Its landing view, every file at once, read up to twenty-five container-backed + workspaces one after another and then fetched their image thumbnails one after + another, so the page waited for the sum of all those round trips. Workspaces + and thumbnails are now read side by side, on eight threads of their own shared + by the whole process, so several people opening the page against a host that + has stopped answering cannot hold up sign-in or anything else on the default + thread pool. A listing waits ten seconds for each call to a host instead of the + archive's sixty-second default, so a host that does not answer at all is + reported as unreadable after ten seconds; one that answers slowly can still + take longer over a deep walk. The "Count files" switch reads hosts the same way. + ## [0.0.521] - 2026-10-06 ### Fixed diff --git a/backend/app/services/sandbox_workspace.py b/backend/app/services/sandbox_workspace.py index 797e31393..3196f32ff 100644 --- a/backend/app/services/sandbox_workspace.py +++ b/backend/app/services/sandbox_workspace.py @@ -25,12 +25,15 @@ import asyncio import base64 import contextlib +import contextvars import logging import shlex from binascii import Error as BinasciiError -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable, Coroutine, Sequence +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from dataclasses import dataclass, field, replace +from functools import partial from io import BytesIO from pathlib import PurePosixPath from typing import Any @@ -38,6 +41,7 @@ from PIL import Image, ImageOps from pydantic_ai_backends import FileData, FileInfo +from pydantic_ai_backends.remote.client import DEFAULT_TIMEOUT_SECONDS from sqlalchemy.ext.asyncio import AsyncSession from app.agents.capabilities.sandbox import SandboxConfig @@ -75,6 +79,58 @@ """Where a listing stops. A host holding a `node_modules` must not turn one workspace into ten thousand rows on a page about twenty-five of them.""" +_HOST_CALLS_AT_ONCE = 8 +"""How many calls to sandbox hosts the whole process has in flight - one `ls` of +a walk or one `read_bytes` each count as one. The listings that read more than one +workspace used to read them one after another, so the page waited for the sum of +twenty-five walks; with this it waits for the slowest few.""" + +_HOST_CALLS = ThreadPoolExecutor( + max_workers=_HOST_CALLS_AT_ONCE, thread_name_prefix="workspace-host" +) +"""The threads host calls run on, and the bound on them. + +Not the loop's default executor, which `to_thread` uses: that one also carries +`bcrypt` and pinned-host DNS, and it has `min(32, cpus + 4)` threads - six on a +two-core host. A host that stops answering holds a thread for the browse timeout, +or sixty seconds for a read, so a few people opening the Workspaces page at once +would leave sign-in queueing behind them. A per-request limit does not prevent +that, since every request brings its own; a pool of the process's own does, and +a dead host then slows only other host calls. + +Unlike `app.core.blocking`, which carries file work, there is no admission gate: +what waits in this queue is a bound method and two strings rather than an upload +buffer, and a cancelled caller's call that has not started is dropped from it. +Threads start on first use and are joined at interpreter exit.""" + + +async def _on_host_thread[T](call: Callable[..., T], *args: object) -> T: + """Run a blocking host call on `_HOST_CALLS`, carrying the caller's context as + `to_thread` would, so its spans stay under the request's trace.""" + context = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor( + _HOST_CALLS, partial(context.run, call, *args) + ) + + +_BROWSE_TIMEOUT_SECONDS = 10.0 +"""How long a listing waits on a host. The archive's own default is sixty seconds, +and a host that has gone away held the whole Workspaces page for that long before +it could say "this host did not answer". A download keeps the default.""" + + +async def _concurrently[T](calls: Sequence[Callable[[], Coroutine[object, object, T]]]) -> list[T]: + """Every call's result in the order given, run side by side. + + A `TaskGroup` rather than `gather`: a call that raises cancels the others and + the group waits for them, so none keeps running after the request has failed + and reaches its database session afterwards. How many reach a host at once is + `_HOST_CALLS`'s to bound, not this. + """ + async with asyncio.TaskGroup() as group: + tasks = [group.create_task(call()) for call in calls] + return [task.result() for task in tasks] + @dataclass(frozen=True) class WorkspaceContents: @@ -362,6 +418,10 @@ def __init__(self, db: AsyncSession) -> None: self.db = db self.connections = SandboxConnectionService(db) self._connections_by_id: dict[UUID, ResolvedConnection] = {} + # The listings read hosts concurrently, and resolving a connection is a + # query on the one session this service holds - which must not be used by + # two coroutines at once. + self._resolving = asyncio.Lock() async def open( self, @@ -950,6 +1010,7 @@ async def measured( unreadable = 0 read = 0 remote = 0 + rows: list[AgentWorkspace] = [] for overview in overviews: row = overview.row if row.backend != "state": @@ -958,7 +1019,8 @@ async def measured( if remote >= limit: continue remote += 1 - contents = await self._entries(ctx, row) + rows.append(row) + for row, contents in zip(rows, await self._entries_of(ctx, rows), strict=True): if contents.unreadable_reason is not None: unreadable += 1 continue @@ -996,20 +1058,38 @@ async def flat_files(self, ctx: AuthContext, *, limit: int = 25) -> FlatFileList """ overviews = await self.visible_to(ctx) files: list[FlatEntry] = [] - unreadable = 0 # A budget across the whole request, not per workspace: a host's thumbnail # is a `read_bytes` for that file, and twenty-five workspaces of photographs # would be a page that fetches two hundred images to draw them 64 pixels # wide. The tiles past it fall back to the grey glyph, which is what every # container-backed image looked like before. budget = HOST_THUMBNAIL_BUDGET - for overview in overviews[:limit]: - contents = await self._entries(ctx, overview.row) - if contents.unreadable_reason is not None: - unreadable += 1 + opened = overviews[:limit] + listed = await self._entries_of(ctx, [overview.row for overview in opened]) + readable: list[tuple[WorkspaceOverview, list[FileInfo]]] = [ + ( + overview, + browsable([entry for entry in contents.entries if not entry.get("is_dir")]), + ) + for overview, contents in zip(opened, listed, strict=True) + if contents.unreadable_reason is None + ] + # Shared out in listing order before anything is fetched, so which tiles are + # drawn does not depend on which host happened to answer first. + shares: list[int] = [] + for overview, entries in readable: + if overview.row.backend == "state": + shares.append(0) continue - entries = browsable([entry for entry in contents.entries if not entry.get("is_dir")]) - previews, thumbnails, budget = await self._tiles(ctx, overview.row, entries, budget) + picked, budget = thumbnail_picks(entries, budget) + shares.append(len(picked)) + tiled = await _concurrently( + [ + partial(self._tiles, ctx, overview.row, entries, share) + for (overview, entries), share in zip(readable, shares, strict=True) + ] + ) + for (overview, entries), (previews, thumbnails, _left) in zip(readable, tiled, strict=True): files.extend( FlatEntry( overview=overview, @@ -1021,11 +1101,24 @@ async def flat_files(self, ctx: AuthContext, *, limit: int = 25) -> FlatFileList ) return FlatFileListing( files=files, - workspaces_read=min(len(overviews), limit) - unreadable, - unreadable=unreadable, + workspaces_read=len(readable), + unreadable=len(opened) - len(readable), truncated=len(overviews) > limit, ) + async def _entries_of( + self, ctx: AuthContext, rows: list[AgentWorkspace] + ) -> list[WorkspaceContents]: + """What each of several workspaces holds, in the order they were given. + + Concurrently, because a container-backed workspace is a walk of round trips + to its host, and these listings used to make them one workspace after + another - twenty-five walks end to end before the Workspaces page could draw + anything. A host that will not answer still comes back as + `unreadable_reason`, never as a raise, so one of them cannot fail the rest. + """ + return await _concurrently([partial(self._entries, ctx, row) for row in rows]) + async def _host_thumbnails( self, ctx: AuthContext, @@ -1045,29 +1138,24 @@ async def _host_thumbnails( one, or a host that stops answering mid-grid, leaves that tile with the glyph every container-backed image had before this existed. """ - drawn: dict[str, str] = {} - for entry in entries: - if budget <= 0: - break - path = str(entry.get("path")) - if PurePosixPath(path).suffix.lower() not in THUMBNAIL_SUFFIXES: - continue - size = entry.get("size") - if size is None or int(size) > THUMBNAIL_SOURCE_LIMIT: - continue - budget -= 1 - try: - raw = await self._read_bytes_from(ctx, row, path) - except Exception: - logger.warning("workspace_thumbnail_unreadable", extra={"path": path}) - continue - if raw is None: - continue - scaled = thumbnail_of(path, raw) - if scaled is not None: - drawn[path] = scaled + paths, budget = thumbnail_picks(entries, budget) + tiles = await _concurrently( + [partial(self._host_thumbnail, ctx, row, path) for path in paths] + ) + drawn = {path: tile for path, tile in zip(paths, tiles, strict=True) if tile is not None} return drawn, budget + async def _host_thumbnail(self, ctx: AuthContext, row: AgentWorkspace, path: str) -> str | None: + """One host image drawn as a tile, or `None` for the glyph.""" + try: + raw = await self._read_bytes_from(ctx, row, path) + except Exception: + logger.warning("workspace_thumbnail_unreadable", extra={"path": path}) + return None + if raw is None: + return None + return thumbnail_of(path, raw) + async def files_of( self, ctx: AuthContext, workspace_id: UUID ) -> tuple[AgentWorkspace, WorkspaceContents]: @@ -1274,7 +1362,7 @@ async def _read_bytes_from( if archive is None: return None try: - return await asyncio.to_thread( + return await _on_host_thread( archive.read_bytes, row.session_id or row.scope_key, path ) except Exception as exc: @@ -1334,9 +1422,11 @@ async def _walk(archive: Any, session: str) -> tuple[list[FileInfo], bool]: not available here: it needs a running session, and browsing deliberately never starts one. - `to_thread`, because `WorkspaceArchive` is a synchronous `httpx.Client`. + Off the loop, because `WorkspaceArchive` is a synchronous `httpx.Client`: `flat_files` runs this for up to 25 workspaces in one request, so a blocking - call would hold the loop for all 25 - and not only for that request. + call would hold the loop for all 25 - and not only for that request. On + `_HOST_CALLS` rather than `to_thread`, so a host that stops answering cannot + take the default executor's threads; see there. Returns: What was found, and whether either bound stopped it - because a @@ -1349,7 +1439,7 @@ async def _walk(archive: Any, session: str) -> tuple[list[FileInfo], bool]: while queue: path, depth = queue.pop(0) try: - entries = await asyncio.to_thread(archive.ls, session, path) + entries = await _on_host_thread(archive.ls, session, path) except Exception: # The root's failure is the host's and belongs to the caller, which # turns it into `unreadable_reason`. One directory below it refusing @@ -1379,7 +1469,7 @@ async def _walk(archive: Any, session: str) -> tuple[list[FileInfo], bool]: async def _remote_entries(self, ctx: AuthContext, row: AgentWorkspace) -> WorkspaceContents: try: - async with self._archive(ctx, row) as archive: + async with self._archive(ctx, row, timeout=_BROWSE_TIMEOUT_SECONDS) as archive: if archive is None: return WorkspaceContents(entries=[]) session = row.session_id or row.scope_key @@ -1464,7 +1554,9 @@ def _state_bytes(row: AgentWorkspace, path: str) -> bytes | None: return backend.read_bytes(path) @asynccontextmanager - async def _archive(self, ctx: AuthContext, row: AgentWorkspace) -> AsyncIterator[Any | None]: + async def _archive( + self, ctx: AuthContext, row: AgentWorkspace, *, timeout: float = DEFAULT_TIMEOUT_SECONDS + ) -> AsyncIterator[Any | None]: """A reader for the host volume behind a container-backed workspace. `None` when the workspace has no connection left to ask - the host was @@ -1488,7 +1580,7 @@ async def _archive(self, ctx: AuthContext, row: AgentWorkspace) -> AsyncIterator if resolved.kind != "docker" or not resolved.row.base_url: yield None return - archive = WorkspaceArchive(resolved.row.base_url, token=resolved.token) + archive = WorkspaceArchive(resolved.row.base_url, token=resolved.token, timeout=timeout) try: yield archive finally: @@ -1510,11 +1602,12 @@ async def _connection(self, ctx: AuthContext, connection_id: UUID) -> ResolvedCo process. Keyed on the id rather than the row, so two workspaces on one connection share the answer and two connections do not. """ - cached = self._connections_by_id.get(connection_id) - if cached is None: - cached = await self.connections.resolve(ctx, connection_id) - self._connections_by_id[connection_id] = cached - return cached + async with self._resolving: + cached = self._connections_by_id.get(connection_id) + if cached is None: + cached = await self.connections.resolve(ctx, connection_id) + self._connections_by_id[connection_id] = cached + return cached def _absent(row: AgentWorkspace, contents: WorkspaceContents, path: str) -> bool: @@ -1658,6 +1751,30 @@ def stored_thumbnail(path: str, data: FileData | None) -> str | None: return thumbnail_of(path, raw) +def thumbnail_picks(entries: list[FileInfo], budget: int) -> tuple[list[str], int]: + """Which of a host's images a listing will fetch, and the budget left. + + A stored image arrives base64 in the document the listing already read, so + drawing it costs nothing; a host's is a `read_bytes` for that one file. So this + is bounded and the size is checked *before* fetching, off the listing entry - a + 20 MB photograph would be read in full to be thrown away by the scaler's own + ceiling. + """ + picked: list[str] = [] + for entry in entries: + if budget <= 0: + break + path = str(entry.get("path")) + if PurePosixPath(path).suffix.lower() not in THUMBNAIL_SUFFIXES: + continue + size = entry.get("size") + if size is None or int(size) > THUMBNAIL_SOURCE_LIMIT: + continue + budget -= 1 + picked.append(path) + return picked, budget + + def thumbnail_of(path: str, raw: bytes) -> str | None: """An image's bytes scaled to a data URI, wherever they came from. diff --git a/backend/tests/test_sandbox_workspace.py b/backend/tests/test_sandbox_workspace.py index 01fef256c..c33e7bb1f 100644 --- a/backend/tests/test_sandbox_workspace.py +++ b/backend/tests/test_sandbox_workspace.py @@ -15,6 +15,7 @@ import logging import os import subprocess +import threading from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock from uuid import UUID, uuid4 @@ -1003,7 +1004,7 @@ async def test_a_text_file_on_a_container_host_is_served(self, monkeypatch, mock from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1036,7 +1037,7 @@ async def test_a_binary_on_a_container_host_is_served_byte_for_byte( png = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR" class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1066,7 +1067,7 @@ async def test_reading_one_file_draws_no_thumbnails(self, monkeypatch, mock_db_s reads: list[str] = [] class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1123,7 +1124,7 @@ async def test_a_path_the_listing_does_not_have_is_missing_rather_than_a_refusal from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1150,7 +1151,7 @@ async def test_a_listing_that_could_not_be_read_does_not_claim_a_file_is_missing from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1197,7 +1198,7 @@ async def test_a_file_with_no_suffix_is_served_like_any_other( from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1261,7 +1262,7 @@ async def test_a_container_binary_is_served_this_way_too(self, monkeypatch, mock png = b"\x89PNG\r\n\x1a\n" class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -1373,6 +1374,128 @@ async def test_one_unreadable_workspace_does_not_empty_the_list( assert listing.unreadable == 1 assert listing.workspaces_read == 1 + async def test_host_workspaces_are_read_side_by_side(self, monkeypatch, mock_db_session): + """The Workspaces page opens on this listing, and it read each host in turn - + so the page waited for twenty-five walks end to end. Each `ls` here waits at + a barrier for the other two; read one at a time, all three time out there + and the listing comes back unreadable.""" + from pydantic_ai_backends import remote as remote_module + + from app.repositories import agent as agent_repo + + meeting = threading.Barrier(3, timeout=5) + + class _Archive(_ClosesItsClient): + def __init__(self, url, token="", timeout=None): + pass + + def ls(self, session_id, path="."): + if path != ".": + return [] + meeting.wait() + return [{"path": f"/{session_id}.csv", "is_dir": False, "size": 1}] + + monkeypatch.setattr(remote_module, "WorkspaceArchive", _Archive, raising=False) + _serve(monkeypatch, _resolved()) + rows = [ + _row(backend="service", session_id=f"xc-{n}", connection_id=uuid4()) for n in range(3) + ] + monkeypatch.setattr(workspace_repo, "list_for_reader", AsyncMock(return_value=rows)) + monkeypatch.setattr(agent_repo, "get_many", AsyncMock(return_value={})) + _no_conversations(monkeypatch) + + listing = await SandboxWorkspaceService(mock_db_session).flat_files(_ctx()) + + assert listing.unreadable == 0 + assert [str(file.info.get("path")) for file in listing.files] == [ + "/xc-0.csv", + "/xc-1.csv", + "/xc-2.csv", + ] + + async def test_a_silent_host_is_given_up_on_before_the_archive_default( + self, monkeypatch, mock_db_session + ): + """The archive waits sixty seconds by default, and a host that had gone away + held the whole page that long before it could say so.""" + from pydantic_ai_backends import remote as remote_module + + from app.repositories import agent as agent_repo + from app.services.sandbox_workspace import _BROWSE_TIMEOUT_SECONDS + + timeouts: list[float | None] = [] + + class _Archive(_ClosesItsClient): + def __init__(self, url, token="", timeout=None): + timeouts.append(timeout) + + def ls(self, session_id, path="."): + return [] + + monkeypatch.setattr(remote_module, "WorkspaceArchive", _Archive, raising=False) + _serve(monkeypatch, _resolved()) + row = _row(backend="service", session_id="xc-1", connection_id=uuid4()) + monkeypatch.setattr(workspace_repo, "list_for_reader", AsyncMock(return_value=[row])) + monkeypatch.setattr(agent_repo, "get_many", AsyncMock(return_value={})) + _no_conversations(monkeypatch) + + await SandboxWorkspaceService(mock_db_session).flat_files(_ctx()) + + assert timeouts == [_BROWSE_TIMEOUT_SECONDS] + assert _BROWSE_TIMEOUT_SECONDS < 60 + + async def test_workspaces_read_side_by_side_resolve_their_connection_one_at_a_time( + self, monkeypatch, mock_db_session + ): + """Resolving a connection is a query on the request's one session, and the + listings now read hosts concurrently. Two workspaces on one connection + resolve it once, and no two resolves are ever in progress together.""" + import asyncio + + from pydantic_ai_backends import remote as remote_module + + from app.repositories import agent as agent_repo + + class _Archive(_ClosesItsClient): + def __init__(self, url, token="", timeout=None): + pass + + def ls(self, session_id, path="."): + return [] + + shared = uuid4() + resolved: list[UUID] = [] + in_progress = 0 + most = 0 + + async def _resolve(self, ctx, connection_id): + nonlocal in_progress, most + in_progress += 1 + most = max(most, in_progress) + # Yield to the loop so an unserialised second resolve would start here. + await asyncio.sleep(0.01) + in_progress -= 1 + resolved.append(connection_id) + return _resolved() + + monkeypatch.setattr(remote_module, "WorkspaceArchive", _Archive, raising=False) + monkeypatch.setattr(SandboxConnectionService, "resolve", _resolve) + rows = [ + _row(backend="service", session_id="xc-1", connection_id=shared), + _row(backend="service", session_id="xc-2", connection_id=shared), + _row(backend="service", session_id="xc-3", connection_id=uuid4()), + ] + monkeypatch.setattr(workspace_repo, "list_for_reader", AsyncMock(return_value=rows)) + monkeypatch.setattr(agent_repo, "get_many", AsyncMock(return_value={})) + _no_conversations(monkeypatch) + + listing = await SandboxWorkspaceService(mock_db_session).flat_files(_ctx()) + + assert listing.unreadable == 0 + assert most == 1 + assert resolved.count(shared) == 1 + assert len(resolved) == 2 + async def test_each_file_names_the_workspace_it_came_from(self, monkeypatch, mock_db_session): """`/report.csv` exists in several workspaces, so a path on its own is ambiguous - and who can see it is the whole point of the column.""" @@ -1676,7 +1799,7 @@ def _archive_recording(monkeypatch, *, ls_fails: bool = False) -> list: made: list = [] class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): made.append(self) def ls(self, session_id, path="."): @@ -2801,7 +2924,7 @@ async def test_reading_by_id_a_path_a_host_did_not_list_is_nothing( from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -2852,7 +2975,7 @@ async def test_a_container_workspace_is_read_off_the_host_volume( from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -2892,7 +3015,7 @@ async def test_a_host_that_cannot_be_read_says_so_instead_of_looking_empty( from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -2926,7 +3049,7 @@ async def test_a_file_on_a_host_that_cannot_be_read_is_refused_not_reported_miss from pydantic_ai_backends import remote as remote_module class _Archive(_ClosesItsClient): - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3447,7 +3570,7 @@ def _archive(monkeypatch, tree: dict[str, list[dict[str, object]]], asked: list[ from pydantic_ai_backends import remote as remote_module class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3567,7 +3690,7 @@ async def test_one_unreadable_folder_does_not_lose_the_workspace( } class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3599,7 +3722,7 @@ async def test_a_root_that_refuses_is_still_the_hosts_failure( from pydantic_ai_backends import remote as remote_module class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3648,7 +3771,7 @@ def _host(self, monkeypatch, entries, bytes_by_path, read: list[str]): png = self._png() class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3759,6 +3882,135 @@ async def test_the_budget_bounds_a_page_of_photographs(self, monkeypatch, mock_d HOST_THUMBNAIL_BUDGET ) + async def test_the_budget_goes_to_the_first_workspaces_listed( + self, monkeypatch, mock_db_session + ): + """Hosts are read side by side, so the one listed first may answer last. The + budget is still shared out in listing order: which tiles are drawn must not + change from one page load to the next with whichever host was quicker. The + first host's listing waits until the second has answered, so it lands last.""" + from pydantic_ai_backends import remote as remote_module + + from app.repositories import agent as agent_repo + from app.services.sandbox_workspace import HOST_THUMBNAIL_BUDGET + + png = self._png() + first = HOST_THUMBNAIL_BUDGET - 4 + quick_listed = threading.Event() + + class _Archive(_ClosesItsClient): + def __init__(self, url, token="", timeout=None): + pass + + def ls(self, session_id, path="."): + if path != ".": + return [] + if session_id == "slow": + assert quick_listed.wait(timeout=5) + else: + quick_listed.set() + count = first if session_id == "slow" else 10 + return [ + {"path": f"{session_id}-{n}.png", "is_dir": False, "size": 120} + for n in range(count) + ] + + def read_bytes(self, session_id, file_path): + return png + + monkeypatch.setattr(remote_module, "WorkspaceArchive", _Archive, raising=False) + _serve(monkeypatch, _resolved()) + rows = [ + _row(backend="service", session_id="slow", connection_id=uuid4()), + _row(backend="service", session_id="quick", connection_id=uuid4()), + ] + monkeypatch.setattr(workspace_repo, "list_for_reader", AsyncMock(return_value=rows)) + monkeypatch.setattr(agent_repo, "get_many", AsyncMock(return_value={})) + _no_conversations(monkeypatch) + + listing = await SandboxWorkspaceService(mock_db_session).flat_files(_ctx()) + + drawn = [str(file.info.get("path")) for file in listing.files if file.thumbnail] + assert sum(path.startswith("slow-") for path in drawn) == first + assert sum(path.startswith("quick-") for path in drawn) == 4 + + async def test_host_calls_are_bounded_across_requests_on_their_own_threads( + self, monkeypatch, mock_db_session + ): + """Every host call holds a thread until the host answers, and a request + builds its own service - so a per-request limit lets a few people opening + the page at once take every thread of the default executor, which sign-in + hashes passwords on. Two page loads at once still make at most + `_HOST_CALLS_AT_ONCE` calls together, and none of them on that executor.""" + import asyncio + import time + + from pydantic_ai_backends import remote as remote_module + + from app.repositories import agent as agent_repo + from app.services.sandbox_workspace import _HOST_CALLS_AT_ONCE, HOST_THUMBNAIL_BUDGET + + png = self._png() + counter = threading.Lock() + in_flight = 0 + most = 0 + threads: set[str] = set() + + def _enter() -> None: + nonlocal in_flight, most + with counter: + in_flight += 1 + most = max(most, in_flight) + threads.add(threading.current_thread().name) + + def _leave() -> None: + nonlocal in_flight + with counter: + in_flight -= 1 + + class _Archive(_ClosesItsClient): + def __init__(self, url, token="", timeout=None): + pass + + def ls(self, session_id, path="."): + if path != ".": + return [] + _enter() + time.sleep(0.02) + _leave() + return [ + {"path": f"{session_id}-{n}.png", "is_dir": False, "size": 120} + for n in range(HOST_THUMBNAIL_BUDGET // 2) + ] + + def read_bytes(self, session_id, file_path): + _enter() + time.sleep(0.02) + _leave() + return png + + monkeypatch.setattr(remote_module, "WorkspaceArchive", _Archive, raising=False) + _serve(monkeypatch, _resolved()) + rows = [ + _row(backend="service", session_id=name, connection_id=uuid4()) + for name in ("one", "two", "three", "four") + ] + monkeypatch.setattr(workspace_repo, "list_for_reader", AsyncMock(return_value=rows)) + monkeypatch.setattr(agent_repo, "get_many", AsyncMock(return_value={})) + _no_conversations(monkeypatch) + + listings = await asyncio.gather( + SandboxWorkspaceService(mock_db_session).flat_files(_ctx()), + SandboxWorkspaceService(mock_db_session).flat_files(_ctx()), + ) + + for listing in listings: + assert sum(file.thumbnail is not None for file in listing.files) == ( + HOST_THUMBNAIL_BUDGET + ) + assert 1 < most <= _HOST_CALLS_AT_ONCE + assert threads and all(name.startswith("workspace-host") for name in threads) + async def test_a_host_that_stops_answering_leaves_the_glyph(self, monkeypatch, mock_db_session): """Failure is silence, and the rest of the grid still draws. A thumbnail is decoration: refusing the listing over one unreadable image would take the @@ -3768,7 +4020,7 @@ async def test_a_host_that_stops_answering_leaves_the_glyph(self, monkeypatch, m from app.repositories import agent as agent_repo class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): @@ -3807,7 +4059,7 @@ async def test_a_file_the_host_no_longer_holds_is_still_listed( from app.repositories import agent as agent_repo class _Archive: - def __init__(self, url, token=""): + def __init__(self, url, token="", timeout=None): pass def ls(self, session_id, path="."): diff --git a/backend/tests/test_security_marker.py b/backend/tests/test_security_marker.py index 06ad9817d..e64216995 100644 --- a/backend/tests/test_security_marker.py +++ b/backend/tests/test_security_marker.py @@ -110,6 +110,7 @@ "tests/test_rag_parent_context.py::TestExpandingResults::test_the_per_search_budget_is_shared_and_stops_expansion": "'budget' names the per-search parent-context character cap, not a spend budget; a size bound, no refusal", "tests/test_sandbox_workspace.py::TestContainerBackedWorkspaces::test_a_docker_workspace_labels_its_tenant_and_reattaches": "a docker workspace labels its tenant for accounting and reattaches, feature, no isolation refusal", "tests/test_sandbox_workspace.py::TestDrawingAHostsImages::test_the_budget_bounds_a_page_of_photographs": "a thumbnail budget bounds image reads, resource limit, not a spend budget", + "tests/test_sandbox_workspace.py::TestDrawingAHostsImages::test_the_budget_goes_to_the_first_workspaces_listed": "thumbnail fetch budget ordering, not a spend budget", "tests/test_services_organizations.py::TestOrganizationService::test_a_new_team_org_starts_with_the_default_monthly_budget": "a new team org gets the default monthly budget, default config, no refusal", "tests/test_services_organizations.py::TestOrganizationService::test_a_personal_org_starts_with_the_default_monthly_budget": "a personal org gets the default monthly budget, default config, no refusal", "tests/test_services_organizations.py::TestOrganizationService::test_the_default_budget_can_be_disabled": "the default budget can be disabled, config, no refusal",