diff --git a/server.py b/server.py index e8e387e..340f1dd 100755 --- a/server.py +++ b/server.py @@ -655,14 +655,31 @@ def register_builtin_tools() -> None: # setup fails (e.g. a build command exits non-zero) are still registered if # register_provider succeeded; broken tool invocations will surface the error # at call time rather than preventing the whole server from coming up. -for provider_spec in load_provider_specs(CONFIG_DIR): + +def bootstrap_provider(provider_spec: dict[str, Any]) -> None: + """Register one provider and run its setup, in whichever order works. + + Registration normally runs first so a slow or failing build never blocks + the tools from being advertised. But a code provider whose code block + imports its declared ``requirements`` at module level cannot exec until + they are pip-installed — so when registration fails we run setup and + retry once before giving up. Never raises. + """ source_path = provider_spec.get("_config_path", "") try: register_provider(provider_spec) except Exception as exc: - print(f"Skipping provider {source_path} — register_provider failed: {exc}") - traceback.print_exc() - continue + print( + f"register_provider failed for {source_path} ({exc}) — " + "running requirements/setup and retrying once" + ) + try: + run_provider_setup(provider_spec) + register_provider(provider_spec) + except Exception as exc2: + print(f"Skipping provider {source_path} — register_provider failed: {exc2}") + traceback.print_exc() + return try: run_provider_setup(provider_spec) except Exception as exc: @@ -674,6 +691,10 @@ def register_builtin_tools() -> None: traceback.print_exc() +for provider_spec in load_provider_specs(CONFIG_DIR): + bootstrap_provider(provider_spec) + + # --------------------------------------------------------------------------- # Remote OAuth-bridge warm-up # --------------------------------------------------------------------------- diff --git a/tests/test_server.py b/tests/test_server.py index 4ff46fd..50ed863 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1142,3 +1142,81 @@ def test_timeout_constant_overridable_via_env(self, monkeypatch): import server as srv assert isinstance(srv.BUILD_COMMAND_TIMEOUT, int) assert srv.BUILD_COMMAND_TIMEOUT > 0 + + +# --------------------------------------------------------------------------- +# bootstrap_provider — register/setup ordering and the missing-import retry +# --------------------------------------------------------------------------- + +from server import bootstrap_provider + + +class TestBootstrapProvider: + SPEC = {"_config_path": "/app/tools/demo.yaml"} + + def test_happy_path_registers_then_runs_setup(self): + with patch("server.register_provider") as mock_reg, \ + patch("server.run_provider_setup") as mock_setup: + bootstrap_provider(self.SPEC) + mock_reg.assert_called_once_with(self.SPEC) + mock_setup.assert_called_once_with(self.SPEC) + + def test_setup_failure_does_not_raise_after_registration(self): + with patch("server.register_provider") as mock_reg, \ + patch("server.run_provider_setup", side_effect=RuntimeError("build failed")): + bootstrap_provider(self.SPEC) # must not raise + mock_reg.assert_called_once() + + def test_register_failure_runs_setup_and_retries(self): + # A code provider importing its declared requirements at module level + # fails to exec before pip install — setup must run, then retry. + with patch("server.register_provider", + side_effect=[ModuleNotFoundError("No module named 'google'"), None]) as mock_reg, \ + patch("server.run_provider_setup") as mock_setup: + bootstrap_provider(self.SPEC) + assert mock_reg.call_count == 2 + mock_setup.assert_called_once_with(self.SPEC) + + def test_register_failure_twice_skips_without_raising(self): + with patch("server.register_provider", + side_effect=ModuleNotFoundError("No module named 'google'")) as mock_reg, \ + patch("server.run_provider_setup") as mock_setup: + bootstrap_provider(self.SPEC) # must not raise + assert mock_reg.call_count == 2 + mock_setup.assert_called_once() + + def test_register_and_setup_failure_skips_without_raising(self): + with patch("server.register_provider", + side_effect=ModuleNotFoundError("No module named 'google'")) as mock_reg, \ + patch("server.run_provider_setup", side_effect=RuntimeError("pip failed")): + bootstrap_provider(self.SPEC) # must not raise + mock_reg.assert_called_once() # retry never reached when setup fails + + def test_end_to_end_code_provider_with_missing_module(self, tmp_path): + """Real exec path: code imports a module that 'pip install' makes available.""" + fake_pkg = tmp_path / "fake_google_pkg" + fake_pkg.mkdir() + (fake_pkg / "__init__.py").write_text("VALUE = 42\n") + + spec = { + "_config_path": str(tmp_path / "demo.yaml"), + "code": "import fake_google_pkg\nasync def ping(context):\n return fake_google_pkg.VALUE\n", + "requirements": ["fake-google-pkg"], + "tools": [{"name": "ping", "function": "ping", + "description": "", "input_schema": {"type": "object", "properties": {}}}], + } + + import sys + def fake_pip(*args, **kwargs): + sys.path.insert(0, str(tmp_path)) + return MagicMock(returncode=0) + + tool_registry.clear() + try: + with patch("server.subprocess.run", side_effect=fake_pip): + bootstrap_provider(spec) + assert tool_registry.get("demo__ping") is not None + finally: + tool_registry.clear() + sys.path.remove(str(tmp_path)) + sys.modules.pop("fake_google_pkg", None)