From 75cf3f1261b4279c418cf45fd55107a1e1af9455 Mon Sep 17 00:00:00 2001 From: Nick Sweeting Date: Wed, 2 Sep 2026 22:02:03 -0700 Subject: [PATCH] Keep OpenCode SSE startup off ASGI request threads --- archivebox/core/asgi.py | 59 ------------------------- archivebox/opencode/views.py | 35 +++++++++++---- archivebox/tests/test_opencode_agent.py | 48 ++++++++++++++++++++ 3 files changed, 74 insertions(+), 68 deletions(-) diff --git a/archivebox/core/asgi.py b/archivebox/core/asgi.py index 46c19223..05f43287 100644 --- a/archivebox/core/asgi.py +++ b/archivebox/core/asgi.py @@ -13,64 +13,5 @@ from archivebox.config.django import setup_django setup_django(check_db=True) - -def _patch_thread_sensitive_context_shutdown() -> None: - """Stop ``ThreadSensitiveContext.__aexit__`` from blocking the daphne loop. - - Django 6.0's ASGIHandler wraps every request in ``async with - ThreadSensitiveContext():`` (django/core/handlers/asgi.py:169). On exit - asgiref calls ``executor.shutdown()`` with the default ``wait=True`` - (asgiref/sync.py:148), which is a *synchronous* ``Thread.join()`` inside - an async function — so it blocks the daphne event loop until the - executor's worker thread exits. - - That's normally fine because the request handler has already awaited - every ``sync_to_async`` it submitted, so the worker is idle and dies as - soon as the shutdown sentinel reaches it. The blocking turns into a - problem when a client disconnects mid-request: - - * ``SyncToAsync.__call__`` shields the executor work with - ``await asyncio.shield(exec_coro)`` (asgiref/sync.py:506) so that the - sync DB call doesn't get torn down halfway through. - * On cancellation it calls ``exec_coro.cancel()`` (line 522) which only - flips the asyncio ``Future`` to cancelled — the underlying thread - keeps running the SQL query. - * Control unwinds to ``__aexit__`` while the orphaned thread is still - mid-query. ``shutdown(wait=True)`` then blocks the event loop until - that orphan finishes. - - Under heavy SQLite contention (the load-test scenario that surfaced - this on cabbage) those orphan threads can take 30 seconds each waiting - for write locks, and the daphne loop is single-threaded — so every - such orphan stalls every other in-flight request, healthchecks time - out, and the container goes ``unhealthy``. - - Switching to ``shutdown(wait=False)`` queues the sentinel and returns - immediately; the worker thread still exits cleanly once its current - task finishes, and asgiref's ``WeakKeyDictionary`` releases the - executor as soon as the request's context is GC'd. No per-request - teardown guarantee is lost — there was no caller relying on it. - """ - from asgiref import sync as _asgiref_sync - - original_aexit = _asgiref_sync.ThreadSensitiveContext.__aexit__ - - async def __aexit__(self, exc, value, tb): # type: ignore[no-redef] - if not self.token: - return - executor = _asgiref_sync.SyncToAsync.context_to_thread_executor.pop(self, None) - if executor is not None: - executor.shutdown(wait=False) - _asgiref_sync.SyncToAsync.thread_sensitive_context.reset(self.token) - - # Idempotent: only patch once even if asgi.py is reloaded. - if getattr(original_aexit, "_archivebox_patched", False): - return - __aexit__._archivebox_patched = True # type: ignore[attr-defined] - _asgiref_sync.ThreadSensitiveContext.__aexit__ = __aexit__ - - -_patch_thread_sensitive_context_shutdown() - # Standard Django ASGI application (no websockets/channels needed) application = get_asgi_application() diff --git a/archivebox/opencode/views.py b/archivebox/opencode/views.py index 2a541a3e..8367d58d 100644 --- a/archivebox/opencode/views.py +++ b/archivebox/opencode/views.py @@ -20,6 +20,7 @@ from abx_plugins.plugins import opencode as opencode_plugin from archivebox.config import CONSTANTS from archivebox.config.common import get_config from archivebox.core.routes_util import build_admin_url, get_api_base_url, get_base_url +from asgiref.sync import sync_to_async from django.http import ( Http404, HttpRequest, @@ -555,17 +556,27 @@ def _request_params(request: HttpRequest) -> tuple[tuple[str, str], ...]: return tuple((key, str(value)) for key, value in dict(request.GET).items()) -async def _event_chunks(request: HttpRequest, settings: dict, path: str | None): +async def _event_chunks( + settings: dict, + path: str | None, + method: str, + params: tuple[tuple[str, str], ...], + headers: dict[str, str], +): + ok, error = await sync_to_async(_ensure_opencode, thread_sensitive=False)(settings) + if not ok: + _LOGGER.warning("OpenCode event stream unavailable: %s", error) + return + timeout = httpx.Timeout(settings["timeout"], read=None) url = _proxy_url(settings, path) - method = request.method or "GET" try: async with httpx.AsyncClient(timeout=timeout, follow_redirects=False) as client: async with client.stream( method, url, - params=_request_params(request), - headers=_request_headers(request, settings), + params=params, + headers=headers, ) as upstream: async for chunk in upstream.aiter_raw(chunk_size=512): yield chunk @@ -662,19 +673,25 @@ def opencode_proxy_view(request: HttpRequest, path: str | None = None): settings["archivebox_admin_url"] = admin_url settings["archivebox_api_url"] = api_url - ok, error = _ensure_opencode(settings) - if not ok: - return _proxy_error_response(error) - if request.method == "GET" and (path or "").endswith("/event"): response = StreamingHttpResponse( - _event_chunks(request, settings, path), + _event_chunks( + settings, + path, + request.method or "GET", + _request_params(request), + _request_headers(request, settings), + ), content_type="text/event-stream", ) response.headers["Cache-Control"] = "no-store" response.headers["X-Accel-Buffering"] = "no" return response + ok, error = _ensure_opencode(settings) + if not ok: + return _proxy_error_response(error) + try: method = request.method or "GET" upstream = requests.request( diff --git a/archivebox/tests/test_opencode_agent.py b/archivebox/tests/test_opencode_agent.py index 2ead6d81..316876d8 100644 --- a/archivebox/tests/test_opencode_agent.py +++ b/archivebox/tests/test_opencode_agent.py @@ -1,3 +1,4 @@ +import asyncio import os import socket import subprocess @@ -8,6 +9,7 @@ from urllib.parse import quote import pytest import requests +from asgiref.testing import ApplicationCommunicator from archivebox.tests.conftest import ADMIN_TEST_HOST, run_archivebox_cmd @@ -268,10 +270,56 @@ def test_opencode_proxy_sse_response_is_unbuffered(admin_client, live_opencode): assert response.status_code == 200 assert response.streaming + assert response.is_async assert response.headers["X-Accel-Buffering"] == "no" assert response.headers["Cache-Control"] == "no-store" +def test_opencode_proxy_sse_returns_headers_before_restart_finishes(admin_client, live_opencode): + from archivebox.core.asgi import application + from archivebox.opencode import views + from django.conf import settings as django_settings + + views._stop_owned_process() + session_cookie_name = django_settings.SESSION_COOKIE_NAME + session_cookie = admin_client.cookies[session_cookie_name].value + path = "/admin/agent/opencode/global/event" + + async def request_event_stream(): + communicator = ApplicationCommunicator( + application, + { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [ + (b"host", ADMIN_TEST_HOST.encode()), + (b"cookie", f"{session_cookie_name}={session_cookie}".encode()), + (b"sec-fetch-site", b"same-origin"), + ], + "client": ("127.0.0.1", 12345), + "server": ("127.0.0.1", 8000), + }, + ) + views._PROCESS_LOCK.acquire() + try: + await communicator.send_input({"type": "http.request", "body": b"", "more_body": False}) + response_start = await communicator.receive_output(timeout=2) + assert response_start["type"] == "http.response.start" + assert response_start["status"] == 200 + finally: + views._PROCESS_LOCK.release() + await communicator.send_input({"type": "http.disconnect"}) + await communicator.wait(timeout=5) + + asyncio.run(request_event_stream()) + + def test_opencode_starts_with_isolated_state(live_opencode): workdir = str(live_opencode.config.data_dir.resolve()) state_dir = live_opencode.config.state_dir