codex/supervaizer-lifespan-cleanup

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:start -->
 # 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.
 

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

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,

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)