diff --git a/AGENTS.md b/AGENTS.md index 2f9a9ad..b50034b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,6 +5,7 @@ This is the canonical agent guide for the Supervaizer controller repo. Supervaiz ## Working Rules - Prefer simple, typed Python changes that match existing FastAPI/Pydantic patterns. +- Do not import from inside functions, methods, or local scopes unless it is absolutely required to avoid a concrete circular import, optional dependency, or startup-cost problem. Prefer module-level imports by default, and document the reason when a local import is unavoidable. - **No guessing / no implicit fallbacks:** when protocol versions, workspace identity, action/resource contracts, authentication, or transport configuration are missing or inconsistent, fail with a clear error that names the missing configuration. Do not infer another context, broaden scope, or silently fall back. - Use `just` recipes from this repo for local commands. - Use `uv` for Python environment and package operations. @@ -80,7 +81,7 @@ Ask. Refusing to act is always safer than taking an action that bypasses these r # GitNexus — Code Intelligence -This project is indexed by GitNexus as **supervaizer** (6117 symbols, 11434 relationships, 278 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. +This project is indexed by GitNexus as **supervaizer** (6225 symbols, 11375 relationships, 276 execution flows). Use the GitNexus MCP tools to understand code, assess impact, and navigate safely. > If any GitNexus tool warns the index is stale, run `npx gitnexus analyze` in terminal first. diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 8bcd04b..b670e27 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -12,6 +12,14 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +### Changed + +- **FastAPI lifespan cleanup** — Controller shutdown now cancels the scheduled-step background loop and waits briefly for it to stop. + +### Tests + +- `tests/test_server.py` — scheduler task cancellation and bounded shutdown waiting during FastAPI lifespan shutdown. + ## [1.1.1] - 2026-05-20 ### Changed diff --git a/src/supervaizer/server.py b/src/supervaizer/server.py index 6097eb5..a6f3fea 100644 --- a/src/supervaizer/server.py +++ b/src/supervaizer/server.py @@ -16,7 +16,7 @@ import time import uuid from collections.abc import AsyncIterator, Callable -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, suppress from datetime import datetime # <-- REMOVED: Path (no longer needed) from hashlib import sha256 from typing import Any, ClassVar, TypeVar, cast @@ -76,6 +76,7 @@ insp = inspect T = TypeVar("T") +SCHEDULED_STEP_SHUTDOWN_TIMEOUT_SECONDS = 5.0 # Additional imports for server persistence @@ -520,9 +521,29 @@ def __init__( openapi_url = "/openapi.json" @asynccontextmanager - async def _lifespan(app: FastAPI) -> AsyncIterator[None]: - asyncio.create_task(_run_scheduled_step_loop(self)) - yield + async def _lifespan(_app: FastAPI) -> AsyncIterator[None]: + # Keep a task handle so shutdown can stop the scheduler cleanly. + scheduled_step_task = asyncio.create_task( + _run_scheduled_step_loop(self), + name="supervaizer-scheduled-step-loop", + ) + try: + yield + finally: + # Give the scheduler a bounded chance to observe cancellation. + scheduled_step_task.cancel() + done, pending = await asyncio.wait( + {scheduled_step_task}, + timeout=SCHEDULED_STEP_SHUTDOWN_TIMEOUT_SECONDS, + ) + if pending: + log.warning( + "[Scheduled step] Shutdown timed out while waiting for " + "the scheduler task to stop" + ) + if done: + with suppress(asyncio.CancelledError): + await scheduled_step_task app = FastAPI( lifespan=_lifespan, diff --git a/tests/test_server.py b/tests/test_server.py index 6185ce2..4a0f4d8 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -10,6 +10,7 @@ # If a copy of the MPL was not distributed with this file, you can obtain one at # https://mozilla.org/MPL/2.0/. +import asyncio import base64 import json import os @@ -24,6 +25,7 @@ from fastapi.testclient import TestClient from rich import inspect +import supervaizer.server as server_module from supervaizer import Server from supervaizer.__version__ import VERSION from supervaizer.agent import Agent @@ -371,6 +373,144 @@ def test_server_generated_api_key_is_exported_for_reload( assert os.environ["SUPERVAIZER_API_KEY"] == server.api_key +@pytest.mark.asyncio +async def test_server_lifespan_cancels_scheduled_step_task( + agent_fixture: Agent, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SUPERVAIZER_LOCAL_MODE", "false") + created_task_names: list[str | None] = [] + waited_tasks: list[set[object]] = [] + waited_timeouts: list[float | None] = [] + + class FakeScheduledStepTask: + def __init__(self) -> None: + self.cancelled = False + self.awaited = False + + def cancel(self) -> None: + self.cancelled = True + + def __await__(self) -> Any: + async def _cancelled() -> None: + self.awaited = True + raise asyncio.CancelledError + + return _cancelled().__await__() + + scheduled_step_task = FakeScheduledStepTask() + + def fake_create_task( + coro: object, *, name: str | None = None, **_kwargs: Any + ) -> FakeScheduledStepTask: + if hasattr(coro, "close"): + coro.close() + created_task_names.append(name) + return scheduled_step_task + + async def fake_wait( + tasks: set[object], *, timeout: float | None = None + ) -> tuple[set[object], set[object]]: + waited_tasks.append(set(tasks)) + waited_timeouts.append(timeout) + return set(tasks), set() + + monkeypatch.setattr( + Server.__init__.__globals__["asyncio"], + "create_task", + fake_create_task, + ) + monkeypatch.setattr( + Server.__init__.__globals__["asyncio"], + "wait", + fake_wait, + ) + + server = Server( + agents=[agent_fixture], + supervisor_account=None, + admin_interface=False, + host="localhost", + port=8001, + environment="test", + api_key="test-key", + ) + + async with server.app.router.lifespan_context(server.app): + assert created_task_names == ["supervaizer-scheduled-step-loop"] + assert scheduled_step_task.cancelled is False + + assert scheduled_step_task.cancelled is True + assert scheduled_step_task.awaited is True + assert waited_tasks == [{scheduled_step_task}] + assert waited_timeouts == [server_module.SCHEDULED_STEP_SHUTDOWN_TIMEOUT_SECONDS] + + +@pytest.mark.asyncio +async def test_server_lifespan_leaves_shutdown_after_scheduler_timeout( + agent_fixture: Agent, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SUPERVAIZER_LOCAL_MODE", "false") + + class PendingScheduledStepTask: + def __init__(self) -> None: + self.cancelled = False + self.awaited = False + + def cancel(self) -> None: + self.cancelled = True + + def __await__(self) -> Any: + async def _unexpected() -> None: + self.awaited = True + raise AssertionError("pending scheduler task should not be awaited") + + return _unexpected().__await__() + + pending_task = PendingScheduledStepTask() + + def fake_create_task( + coro: object, *, name: str | None = None, **_kwargs: Any + ) -> PendingScheduledStepTask: + if hasattr(coro, "close"): + coro.close() + return pending_task + + async def fake_wait( + tasks: set[object], *, timeout: float | None = None + ) -> tuple[set[object], set[object]]: + assert timeout == server_module.SCHEDULED_STEP_SHUTDOWN_TIMEOUT_SECONDS + return set(), set(tasks) + + monkeypatch.setattr( + Server.__init__.__globals__["asyncio"], + "create_task", + fake_create_task, + ) + monkeypatch.setattr( + Server.__init__.__globals__["asyncio"], + "wait", + fake_wait, + ) + + server = Server( + agents=[agent_fixture], + supervisor_account=None, + admin_interface=False, + host="localhost", + port=8001, + environment="test", + api_key="test-key", + ) + + async with server.app.router.lifespan_context(server.app): + assert pending_task.cancelled is False + + assert pending_task.cancelled is True + assert pending_task.awaited is False + + def test_server_decrypt(server_fixture: Server) -> None: unencrypted_parameters = str({"KEY": "VALUE"}) encrypted_parameters = server_fixture.encrypt(unencrypted_parameters)