diff --git a/src/smallestai/atoms/crew/server.py b/src/smallestai/atoms/crew/server.py index 54e8833d..227ae2a5 100644 --- a/src/smallestai/atoms/crew/server.py +++ b/src/smallestai/atoms/crew/server.py @@ -6,7 +6,7 @@ from fastapi import FastAPI, WebSocket from loguru import logger -from smallestai.atoms.crew.session import CrewSession +from smallestai.atoms.crew.session import CrewSession, _StartupProbeComplete async def _dry_run_setup_handler( @@ -19,8 +19,12 @@ async def _dry_run_setup_handler( an env var that isn't set, or imports something that isn't installed) *before* the pod accepts a real WebSocket connection. - Doesn't connect to any external services — node `start()` is never - called. Only constructor / `add_node()` / `add_edge()` logic runs. + Runs the session in dry-run mode: node `start()` is never reached in a + live sense, no external services are contacted. The canonical handler's + closing `await session.start()` builds the graph and then raises + `_StartupProbeComplete` (caught below), so validation no longer reports a + spurious "Session not initialized" for healthy code. Only constructor / + `add_node()` / `add_edge()` / graph-build logic runs. """ class _NullWebSocket: @@ -42,10 +46,16 @@ async def close(self, code=1000): session_id="startup-validation", setup_handler=setup_handler, ) + session._dry_run = True # Don't `await session.initialize()` — that would actually wait on # the init handshake and start the receive loop. We only want to # exercise the user's setup_handler enough to surface __init__ errors. - await setup_handler(session) + # A canonical handler ends with `await session.start()`, which in dry-run + # builds the graph then raises `_StartupProbeComplete` to halt cleanly. + try: + await setup_handler(session) + except _StartupProbeComplete: + pass class SessionHandler: diff --git a/src/smallestai/atoms/crew/session.py b/src/smallestai/atoms/crew/session.py index 7629cdb5..9d4bbd0d 100644 --- a/src/smallestai/atoms/crew/session.py +++ b/src/smallestai/atoms/crew/session.py @@ -27,6 +27,19 @@ from smallestai.atoms.crew.task_manager import TaskManager, TaskManagerParams +class _StartupProbeComplete(Exception): + """Internal signal raised by ``CrewSession.start()`` during the startup + dry-run once the graph has been built. + + The server's startup validation runs the user's ``setup_handler`` without a + live init handshake, purely to surface node ``__init__`` / import / env / + graph errors before traffic arrives. The canonical handler ends with + ``await session.start()``; in the dry-run that call builds the graph and then + raises this to halt cleanly (no init required, no nodes started, no external + connections). It is caught by the validator and never surfaces to users. + """ + + @dataclass class EventHandler: name: str @@ -131,6 +144,9 @@ def __init__( self._init_event: Optional[SDKSystemInitEvent] = None + # Set only by the server's startup dry-run (never on a real session). + self._dry_run = False + self.task_manager = TaskManager() self.loop = loop or asyncio.get_event_loop() @@ -205,6 +221,17 @@ def add_edge(self, parent: CrewNode, child: CrewNode): async def start(self) -> None: """Start the session""" logger.info(f"[{self.name}] Starting session") + + if self._dry_run: + # Startup validation: build the graph to surface node/edge/cycle + # errors, then halt. No init handshake is required and no nodes are + # started, so nothing connects to external services. Raising here + # keeps the canonical `await session.start()` handler from blocking + # on the (never-arriving) init event during validation. + logger.info(f"[{self.name}] Startup dry-run: building graph with {len(self.nodes)} nodes") + self._build_graph() + raise _StartupProbeComplete() + if not self._init_event: logger.error( "This should not happen because this method should always be called after the init event is received which will set the init event" diff --git a/tests/custom/test_crew_startup_validation.py b/tests/custom/test_crew_startup_validation.py new file mode 100644 index 00000000..c8a6bace --- /dev/null +++ b/tests/custom/test_crew_startup_validation.py @@ -0,0 +1,76 @@ +"""Startup validation dry-run. + +A healthy `setup_handler` that ends with `await session.start()` (the canonical +pattern) must pass validation without logging an error. Only handlers that raise +during node construction / graph build (missing env, bad import, cycles) should +fail validation and leave the pod not-ready. +""" + +import unittest + +from smallestai.atoms.crew.nodes import OutputCrewNode +from smallestai.atoms.crew.server import AtomsCrewApp, _dry_run_setup_handler +from smallestai.atoms.crew.session import CrewSession, _StartupProbeComplete + + +class _Assistant(OutputCrewNode): + def __init__(self): + super().__init__(name="assistant") + + async def generate_response(self): + if False: + yield "" # never runs; satisfies the abstract async-generator + + +class StartupValidationTest(unittest.IsolatedAsyncioTestCase): + async def test_canonical_handler_passes_validation(self): + # Canonical handler: build a node, then `await session.start()`. + async def setup(session: CrewSession): + session.add_node(_Assistant()) + await session.start() + # Real handlers continue here (event handlers, wait_until_complete); + # the dry-run halts at start(), so this is never reached. + raise AssertionError("dry-run should have halted at start()") + + # Should NOT raise: start() raises _StartupProbeComplete internally, + # which _dry_run_setup_handler swallows. + await _dry_run_setup_handler(setup) + + app = AtomsCrewApp(setup_handler=setup) + await app._validate_startup() + self.assertTrue(app._ready) + self.assertIsNone(app._not_ready_reason) + + async def test_broken_handler_fails_validation(self): + # A node whose __init__ raises (e.g. missing env var) must fail. + class _Broken(OutputCrewNode): + def __init__(self): + super().__init__(name="broken") + raise ValueError("MISSING_API_KEY not set") + + async def generate_response(self): + if False: + yield "" + + async def setup(session: CrewSession): + session.add_node(_Broken()) + await session.start() + + app = AtomsCrewApp(setup_handler=setup) + await app._validate_startup() + self.assertFalse(app._ready) + self.assertIn("MISSING_API_KEY", app._not_ready_reason or "") + + async def test_dry_run_start_raises_probe_complete_and_builds_graph(self): + session = CrewSession(websocket=None, session_id="t", setup_handler=None) # type: ignore[arg-type] + session._dry_run = True + session.add_node(_Assistant()) + with self.assertRaises(_StartupProbeComplete): + await session.start() + # graph was built (root + node + sink), and no init event was required + self.assertGreaterEqual(len(session.nodes), 3) + self.assertIsNone(session._init_event) + + +if __name__ == "__main__": + unittest.main()