diff --git a/libs/code/deepagents_code/app.py b/libs/code/deepagents_code/app.py index 944a2197cd..10db81f914 100644 --- a/libs/code/deepagents_code/app.py +++ b/libs/code/deepagents_code/app.py @@ -15254,15 +15254,29 @@ async def _send_to_agent( if self._chat_input: self._chat_input.set_cursor_active(active=False) + from functools import partial + # Use run_worker to avoid blocking the main event loop # This allows the UI to remain responsive during agent execution - self._agent_worker = self.run_worker( - self._run_agent_task( - message, - message_kwargs=message_kwargs, - goal_notice_current=resuming_blocked, - ), - exclusive=False, + # + # Passed as a callable rather than a coroutine: a worker + # cancelled before its first event-loop step never runs its + # work, and an already-built coroutine would then be finalized + # unawaited (`RuntimeWarning`, plus an unraisable exception when + # the interpreter is far enough into teardown). Building it + # inside the worker means nothing exists to strand. + turn: Callable[[], Coroutine[None, None, None]] = partial( + self._run_agent_task, + message, + message_kwargs=message_kwargs, + goal_notice_current=resuming_blocked, + ) + # Cast because Textual's `WorkType` alias admits both a + # coroutine factory and a plain callable, so the result type + # infers as the union rather than the `None` the turn returns. + self._agent_worker = cast( + "Worker[None]", + self.run_worker(turn, name="_run_agent_task", exclusive=False), ) worker_started = True finally: diff --git a/libs/code/deepagents_code/sessions.py b/libs/code/deepagents_code/sessions.py index 6f8b952199..37659bf3b8 100644 --- a/libs/code/deepagents_code/sessions.py +++ b/libs/code/deepagents_code/sessions.py @@ -32,6 +32,8 @@ _MAX_INITIAL_PROMPT_CACHE = 4096 _recent_threads_cache: dict[tuple[str | None, int], list[ThreadInfo]] = {} _MAX_RECENT_THREADS_CACHE_KEYS = 16 +_DEFAULT_SQLITE_TIMEOUT = 5.0 +"""Seconds to wait out a locked database; matches the `sqlite3` default.""" def _patch_aiosqlite() -> None: @@ -93,6 +95,80 @@ async def _drain_aiosqlite_worker(conn: aiosqlite.Connection) -> None: await asyncio.to_thread(worker.join, 5.0) +def _guard_sqlite_handle(conn: aiosqlite.Connection) -> None: + """Keep the sqlite handle closable when the opening task is cancelled. + + `aiosqlite` opens the database on its worker thread and delivers the raw + `sqlite3.Connection` back through a future, recording it on the + `Connection` only once the awaiting coroutine resumes. Background workers + are routinely cancelled at app exit, and a cancel landing anywhere in that + window leaves the handle unreachable from the cleanup that follows: + + - Cancelled while the worker is still opening, the library has no handle + recorded yet, so the cleanup it queues closes nothing. + - Cancelled after the handle is delivered but before the coroutine resumes, + the library clears its own record before that queued cleanup can run, so + again it closes nothing. + + Either way the garbage collector is left to report `ResourceWarning: + unclosed database`. Recording the handle from the worker thread covers the + first case; queueing an explicit close ahead of the library's own cleanup + covers the second. Both run on the thread that opened the handle, and + closing twice is a no-op, so neither disturbs a normal shutdown. + + Args: + conn: A connection that has not been opened yet. + """ + # No public hooks for any of this, so tolerate it moving: the leak avoided + # here is a warning at teardown, not something worth failing a query for. + connector = getattr(conn, "_connector", None) + queue = getattr(conn, "_tx", None) + stop = getattr(conn, "stop", None) + if connector is None or queue is None or stop is None: + logger.debug("aiosqlite internals moved; cannot guard the sqlite handle") + return + + def open_and_record() -> sqlite3.Connection: + handle = connector() + # The assignment aiosqlite makes once the awaiting coroutine resumes, + # made early enough that a cancel cannot get in front of it. + conn._connection = handle + return handle + + def stop_and_close() -> asyncio.Future[Any] | None: + # Runs before aiosqlite drops its own reference, so the handle is still + # here to queue a close for -- ahead of the stop sentinel, which ends + # the worker loop. A `None` future keeps the worker from reaching for an + # event loop that may already be gone. + handle = conn._connection + if handle is not None: + queue.put_nowait((None, handle.close)) + return stop() + + conn._connector = open_and_record + # Shadows the bound method on this one instance; the declared type is the + # unbound `stop(self)`, which a zero-argument replacement cannot match. + conn.stop = stop_and_close # ty: ignore[invalid-assignment] + + +def _new_connection(timeout: float = _DEFAULT_SQLITE_TIMEOUT) -> aiosqlite.Connection: + """Build an unopened connection to the sessions database. + + Args: + timeout: Seconds to wait out a locked database before giving up. + + Returns: + A connection that closes its sqlite handle even when interrupted. + """ + import aiosqlite as _aiosqlite + + _patch_aiosqlite() + + conn = _aiosqlite.connect(str(get_db_path()), timeout=timeout) + _guard_sqlite_handle(conn) + return conn + + @asynccontextmanager async def _connect() -> AsyncIterator[aiosqlite.Connection]: """Import aiosqlite, apply the compatibility patch, and connect. @@ -103,18 +179,12 @@ async def _connect() -> AsyncIterator[aiosqlite.Connection]: Yields: An open aiosqlite connection to the sessions database. """ - import aiosqlite as _aiosqlite - - _patch_aiosqlite() - - conn: aiosqlite.Connection | None = None + conn = _new_connection(timeout=30.0) try: - async with _aiosqlite.connect(str(get_db_path()), timeout=30.0) as opened: - conn = opened + async with conn as opened: yield opened finally: - if conn is not None: - await _drain_aiosqlite_worker(conn) + await _drain_aiosqlite_worker(conn) class ThreadInfo(TypedDict): @@ -1468,20 +1538,15 @@ async def get_checkpointer() -> AsyncIterator[AsyncSqliteSaver]: """ from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver - _patch_aiosqlite() - - saver: AsyncSqliteSaver | None = None + # Built here rather than through `AsyncSqliteSaver.from_conn_string` so the + # connection is one this module owns and can clean up after an interrupted + # connect; see `_guard_sqlite_handle`. + conn = _new_connection() try: - async with AsyncSqliteSaver.from_conn_string( - str(get_db_path()) - ) as checkpointer: - saver = checkpointer - yield checkpointer + async with conn as opened: + yield AsyncSqliteSaver(opened) finally: - if saver is not None: - conn = getattr(saver, "conn", None) - if conn is not None: - await _drain_aiosqlite_worker(conn) + await _drain_aiosqlite_worker(conn) _DEFAULT_THREAD_LIMIT = 20 diff --git a/libs/code/pyproject.toml b/libs/code/pyproject.toml index 8a823c6442..7ebd8dde99 100644 --- a/libs/code/pyproject.toml +++ b/libs/code/pyproject.toml @@ -341,6 +341,10 @@ filterwarnings = [ # uses `stacklevel=2`, so it is attributed to this package, not to the # module that raised it. "error:Unsupported usage key for standard pricing:UserWarning", + # `google.genai.types` builds a union alias out of `typing._UnionGenericAlias` + # at import time, which Python 3.14 deprecates. Nothing here can act on it -- + # it fires before any of our code runs and is fixed by upgrading the SDK. + "ignore:'_UnionGenericAlias' is deprecated:DeprecationWarning:google\\.genai\\.types", ] addopts = "--strict-markers --strict-config --durations=5" diff --git a/libs/code/tests/unit_tests/test_app.py b/libs/code/tests/unit_tests/test_app.py index 8f7fe225a3..86a3e9c93d 100644 --- a/libs/code/tests/unit_tests/test_app.py +++ b/libs/code/tests/unit_tests/test_app.py @@ -5290,8 +5290,8 @@ async def test_send_to_agent_resets_visible_output_started_flag(self) -> None: Without this reset the gate would be sticky: once any turn produced output, every later turn's Esc-interrupt would stop restoring the - prompt. Closing the worker coroutine leaves the flag as `_send_to_agent` - set it, without running the turn. + prompt. Stubbing the worker leaves the flag as `_send_to_agent` set it, + without running the turn. """ app = DeepAgentsApp() app._agent = MagicMock() @@ -5308,8 +5308,6 @@ async def test_send_to_agent_resets_visible_output_started_flag(self) -> None: with patch.object(app, "run_worker") as mock_rw: mock_rw.return_value = MagicMock() await app._send_to_agent("next question") - coro = mock_rw.call_args[0][0] - coro.close() assert app._active_turn_visible_output_started is False @@ -5553,6 +5551,27 @@ async def test_interrupt_before_worker_starts_releases_turn(self) -> None: await pilot.pause() assert not app._pending_messages + async def test_turn_is_spawned_as_a_callable(self) -> None: + """The worker is handed a callable, not an already-built coroutine. + + Textual never runs the work of a worker cancelled before its first + event-loop step, so a pre-built coroutine would be left to be finalized + unawaited — a `RuntimeWarning`, and an unraisable exception once + interpreter teardown has gone far enough to break the import machinery + the coroutine's cleanup relies on. + """ + app = self._configured_app() + async with app.run_test() as pilot: + await pilot.pause() + + with patch.object(app, "run_worker") as mock_run_worker: + mock_run_worker.return_value = MagicMock() + await app._send_to_agent("hello") + + work = mock_run_worker.call_args[0][0] + assert not inspect.iscoroutine(work) + assert callable(work) + async def test_later_turns_can_still_recover(self) -> None: """The started-marker resets at turn end, so turn 2 recovers as well. @@ -15604,8 +15623,6 @@ async def test_pending_shell_flushed_on_next_user_send(self) -> None: with patch.object(app, "run_worker") as mock_rw: mock_rw.return_value = MagicMock() await app._send_to_agent("what did that print?") - coro = mock_rw.call_args[0][0] - coro.close() app._agent.aupdate_state.assert_awaited_once() call = app._agent.aupdate_state.await_args @@ -15635,8 +15652,6 @@ async def test_pending_shell_first_message_uses_session_thread(self) -> None: with patch.object(app, "run_worker") as mock_rw: mock_rw.return_value = MagicMock() await app._send_to_agent("what did that print?") - coro = mock_rw.call_args[0][0] - coro.close() app._agent.aupdate_state.assert_awaited_once() call = app._agent.aupdate_state.await_args @@ -15853,8 +15868,6 @@ async def test_pending_shell_flush_precedes_agent_worker(self) -> None: manager.attach_mock(app._agent.aupdate_state, "flush") manager.attach_mock(mock_rw, "spawn") await app._send_to_agent("what did that print?") - coro = mock_rw.call_args[0][0] - coro.close() # Flush the `!` pair into state before spawning the turn, so the model # sees the shell output ahead of this turn's user message. diff --git a/libs/code/tests/unit_tests/test_sessions.py b/libs/code/tests/unit_tests/test_sessions.py index e26525d877..0825e24a58 100644 --- a/libs/code/tests/unit_tests/test_sessions.py +++ b/libs/code/tests/unit_tests/test_sessions.py @@ -1,12 +1,15 @@ """Tests for session/thread management.""" import asyncio +import gc import json import sqlite3 +import threading import uuid +import warnings from datetime import UTC, datetime, timedelta from pathlib import Path -from typing import TYPE_CHECKING, ClassVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -472,6 +475,113 @@ async def _test() -> None: "aiosqlite worker thread should be joined after _connect exit" ) + def test_cancelled_open_closes_the_sqlite_handle(self, tmp_path, monkeypatch): + """A connect cancelled mid-open still closes the sqlite handle. + + Background workers are routinely cancelled at app exit, and aiosqlite + opens the database on a worker thread. Without the handle being + recorded as soon as it exists, a cancel landing in that window strands + it for the garbage collector to report as an unclosed database. + """ + db_path = tmp_path / "cancelled.db" + opening = threading.Event() + finish_open = threading.Event() + handles: list[sqlite3.Connection] = [] + connections: list[aiosqlite.Connection] = [] + real_sqlite_connect = sqlite3.connect + real_new_connection = sessions._new_connection + + def blocking_connect(database: str, **kwargs: Any) -> sqlite3.Connection: + opening.set() + finish_open.wait(10) + handle = real_sqlite_connect(database, **kwargs) + handles.append(handle) + return handle + + def capturing_new_connection(timeout: float) -> "aiosqlite.Connection": + conn = real_new_connection(timeout) + connections.append(conn) + return conn + + async def _test() -> None: + async def use_connection() -> None: + async with sessions._connect(): + pass + + task = asyncio.create_task(use_connection()) + await asyncio.to_thread(opening.wait, 10) + task.cancel() + finish_open.set() + with pytest.raises(asyncio.CancelledError): + await task + + monkeypatch.setattr(sessions, "get_db_path", lambda: db_path) + monkeypatch.setattr(sessions, "_new_connection", capturing_new_connection) + monkeypatch.setattr(sqlite3, "connect", blocking_connect) + asyncio.run(_test()) + + assert handles, "the worker thread should have opened a handle" + assert not connections[0]._thread.is_alive() + handles.clear() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + gc.collect() + unclosed = [w for w in caught if "unclosed database" in str(w.message)] + assert not unclosed, [str(w.message) for w in unclosed] + + def test_cancel_after_open_closes_the_sqlite_handle(self, tmp_path, monkeypatch): + """A cancel landing after the open, before the resume, also closes it. + + The other half of the same window: aiosqlite has the handle by now, but + its cancellation path drops that reference before the cleanup it queued + gets a chance to run, so the handle would again be left to the garbage + collector. + """ + db_path = tmp_path / "raced.db" + handles: list[sqlite3.Connection] = [] + connections: list[aiosqlite.Connection] = [] + pending: dict[str, Any] = {} + real_sqlite_connect = sqlite3.connect + real_new_connection = sessions._new_connection + + def cancelling_connect(database: str, **kwargs: Any) -> sqlite3.Connection: + handle = real_sqlite_connect(database, **kwargs) + handles.append(handle) + # Queued from the worker thread before aiosqlite delivers the + # handle to the awaiting coroutine, so the cancel deterministically + # lands in the gap between the open and the resume. + pending["loop"].call_soon_threadsafe(pending["task"].cancel) + return handle + + def capturing_new_connection(timeout: float) -> "aiosqlite.Connection": + conn = real_new_connection(timeout) + connections.append(conn) + return conn + + async def _test() -> None: + async def use_connection() -> None: + async with sessions._connect(): + pass + + pending["loop"] = asyncio.get_running_loop() + pending["task"] = asyncio.create_task(use_connection()) + with pytest.raises(asyncio.CancelledError): + await pending["task"] + + monkeypatch.setattr(sessions, "get_db_path", lambda: db_path) + monkeypatch.setattr(sessions, "_new_connection", capturing_new_connection) + monkeypatch.setattr(sqlite3, "connect", cancelling_connect) + asyncio.run(_test()) + + assert handles, "the worker thread should have opened a handle" + assert not connections[0]._thread.is_alive() + handles.clear() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + gc.collect() + unclosed = [w for w in caught if "unclosed database" in str(w.message)] + assert not unclosed, [str(w.message) for w in unclosed] + class TestFormatTimestamp: """Tests for format_timestamp helper."""