talkops-opscloud 0.3.0b5__tar.gz → 0.3.0b7__tar.gz

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: talkops-opscloud
-Version: 0.3.0b5
+Version: 0.3.0b7
 Summary: Extensible Terminal Multi-Agent Framework for Cloud Operations & DevOps Coding
 License-File: LICENSE
 Requires-Python: >=3.12

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "talkops-opscloud"
-version = "0.3.0b5"
+version = "0.3.0b7"
 description = "Extensible Terminal Multi-Agent Framework for Cloud Operations & DevOps Coding"
 readme = "README.md"
 requires-python = ">=3.12"

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/cli/interactive.py

@@ -120,6 +120,21 @@ def run_interactive(
         "trust_project_mcp": trust_project_mcp,
     }
 
+    server_kwargs: dict[str, Any] = {
+        "assistant_id": "opscloud",
+        "model_name": resolved_model,
+        "interactive": True,
+        "auto_approve": (approval_mode == ApprovalMode.AUTO or approval_mode == "auto"),
+        "cwd": resolved_cwd,
+        "read_only": read_only,
+        "aws_profile": aws_profile,
+        "aws_region": aws_region,
+        "shell_allow_list": tuple(shell_allow_list) if shell_allow_list else None,
+        "mcp_config_path": mcp_config_path,
+        "no_mcp": no_mcp,
+        "trust_project_mcp": trust_project_mcp,
+    }
+
     try:
         app = OpsCloudApp(
             server_url=url,
@@ -129,6 +144,7 @@ def run_interactive(
             model=resolved_model,
             approval_mode=approval_mode,
             defer_server_start=defer_server_start,
+            server_kwargs=server_kwargs,
             goal=goal,
             cwd=resolved_cwd,
             mcp_preload_kwargs=mcp_preload_kwargs,

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/cli/server_manager.py

@@ -36,7 +36,7 @@ async def start_server_and_get_agent(
     port: int = 0,
     **kwargs: Any,
 ) -> tuple[RemoteAgent, ServerProcess]:
-    """Start a LangGraph dev server and return a RemoteAgent client."""
+    """Start a LangGraph dev server asynchronously and return a RemoteAgent client."""
     from opscloud.project_utils import ProjectContext
 
     project_context = (
@@ -44,15 +44,24 @@ async def start_server_and_get_agent(
         if cwd is not None
         else ProjectContext.from_user_cwd(Path.cwd())
     )
+    resolved_model = model_name or kwargs.pop("model", None)
+    if not resolved_model:
+        try:
+            from opscloud.model.factory import _get_default_model_spec
+
+            resolved_model = _get_default_model_spec()
+        except Exception:
+            resolved_model = None
+
     config = ServerConfig.from_cli_args(
         project_context=project_context,
         assistant_id=assistant_id,
-        model=model_name or "anthropic:claude-3-5-sonnet-latest",
+        model=resolved_model,
         interactive=interactive,
         approval_mode="never" if auto_approve else "auto",
         **kwargs,
     )
     server = ServerProcess(config, host=host, port=port)
-    url = server.start()
+    url = await server.astart()
     agent = RemoteAgent(url=url, graph_name="agent")
     return agent, server

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/commands/core/pool.py

@@ -41,6 +41,9 @@ class PoolHandler(BaseCommandHandler):
         # 1. Clear pool configuration
         if args in ("--clear", "clear"):
             clear_agent_pool()
+            pool_mgr = get_model_pool_manager()
+            with pool_mgr._lock:
+                pool_mgr._instances.clear()
             return CommandResult(
                 success=True,
                 message="**Agent Model Pool Cleared:** Reverted to automatic dynamic discovery from provider capabilities.",
@@ -84,7 +87,7 @@ class PoolHandler(BaseCommandHandler):
                 current = load_agent_pool() or {}
                 merged = {**current, **updates}
 
-                from opscloud.model.config import detect_provider
+                from opscloud.model.config import detect_provider, apply_stored_credentials
                 provs = set()
                 for v in merged.values():
                     if ":" in v:
@@ -97,6 +100,13 @@ class PoolHandler(BaseCommandHandler):
                 provider_to_save = next(iter(provs)) if len(provs) == 1 else ("multi" if len(provs) > 1 else None)
                 save_agent_pool(merged, provider=provider_to_save)
 
+                pool_mgr = get_model_pool_manager()
+                with pool_mgr._lock:
+                    pool_mgr._instances.clear()
+
+                for p in provs:
+                    apply_stored_credentials(p)
+
                 from opscloud.model.config import has_provider_credentials
                 warnings = []
                 for p in provs:

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/middleware/skills.py

@@ -222,6 +222,7 @@ class PluginSkillsMiddleware(SkillsMiddleware):
             source[2] if len(source) == _PLUGIN_SKILL_SOURCE_LENGTH else None for source in sources
         )
         self._allowed_skills = tuple(allowed_skills) if allowed_skills is not None else None
+        self._cached_live_skills: list[sdk_skills.SkillMetadata] | None = None
 
     def _format_skills_locations(self) -> str:
         """Format skills locations for display in system prompt."""
@@ -322,6 +323,7 @@ class PluginSkillsMiddleware(SkillsMiddleware):
             SkillsStateUpdate if changes were detected, or None.
         """
         live_skills, errors = self._get_live_skills()
+        self._cached_live_skills = live_skills
         current_skills = state.get("skills_metadata") or []
         current_names = [s["name"] for s in current_skills]
         live_names = [s["name"] for s in live_skills]
@@ -362,7 +364,14 @@ class PluginSkillsMiddleware(SkillsMiddleware):
         Returns:
             Modified request with pruned skill metadata.
         """
-        live_skills, _ = self._get_live_skills()
+        live_skills = getattr(self, "_cached_live_skills", None)
+        if live_skills is None:
+            if hasattr(request, "state") and isinstance(request.state, dict) and request.state.get("skills_metadata"):
+                live_skills = request.state["skills_metadata"]
+            else:
+                live_skills, _ = self._get_live_skills()
+            self._cached_live_skills = live_skills
+
         live_names = {s["name"] for s in live_skills}
 
         if hasattr(request, "state") and isinstance(request.state, dict):

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/model/pool.py

@@ -169,70 +169,77 @@ class DynamicModelPoolManager:
         self,
         provider: str | None = None,
         base_spec: str | None = None,
+        *,
+        ignore_custom_pool: bool = False,
     ) -> dict[int, tuple[str, str]]:
         """Dynamically build Tier 0, 1, and 2 definitions from registry metadata.
 
         Args:
             provider: Target provider name, or None to infer.
             base_spec: Baseline model spec selected by the session/thread.
+            ignore_custom_pool: If True, bypass user-configured [agent_pool] and
+                perform raw registry discovery for the requested provider/baseline.
 
         Returns:
             Dictionary mapping tier level (0, 1, 2) to (model_spec, reasoning_effort).
         """
         # 1. Prioritize user-configured pool from config.toml [agent_pool]
-        from opscloud.config.toml_config import load_agent_pool
+        if not ignore_custom_pool:
+            from opscloud.config.toml_config import load_agent_pool
+
+            user_pool = load_agent_pool()
+            if user_pool:
+                resolved_provider = self._resolve_provider(provider, base_spec)
+                pool_prov = user_pool.get("provider")
+                t0_raw = user_pool.get("fast")
+                t1_raw = user_pool.get("standard")
+                t2_raw = user_pool.get("powerful")
+
+                if t0_raw and t1_raw and t2_raw:
+                    def _norm(raw_spec: str) -> str:
+                        s = raw_spec.strip()
+                        if ":" not in s:
+                            inferred = detect_provider(s)
+                            if inferred and inferred in AVAILABLE_MODELS:
+                                return f"{inferred}:{s}"
+                            if pool_prov and pool_prov in AVAILABLE_MODELS:
+                                return f"{pool_prov}:{s}"
+                            return f"{resolved_provider}:{s}"
+                        return normalize_model_spec(s)
+
+                    t0 = _norm(t0_raw)
+                    t1 = _norm(t1_raw)
+                    t2 = _norm(t2_raw)
+
+                    prov0 = t0.split(":", 1)[0] if ":" in t0 else resolved_provider
+                    prov1 = t1.split(":", 1)[0] if ":" in t1 else resolved_provider
+                    prov2 = t2.split(":", 1)[0] if ":" in t2 else resolved_provider
+                    distinct_provs = {prov0, prov1, prov2}
+                    is_multi = len(distinct_provs) > 1 or pool_prov == "multi"
+
+                    should_use_custom = (
+                        not provider
+                        or provider in ("dynamic", "auto")
+                        or is_multi
+                        or provider == pool_prov
+                        or provider in distinct_provs
+                    )
 
-        user_pool = load_agent_pool()
-        if user_pool:
-            resolved_provider = self._resolve_provider(provider, base_spec)
-            pool_prov = user_pool.get("provider")
-            t0_raw = user_pool.get("fast")
-            t1_raw = user_pool.get("standard")
-            t2_raw = user_pool.get("powerful")
-
-            if t0_raw and t1_raw and t2_raw:
-                def _norm(raw_spec: str) -> str:
-                    s = raw_spec.strip()
-                    if ":" not in s:
-                        inferred = detect_provider(s)
-                        if inferred and inferred in AVAILABLE_MODELS:
-                            return f"{inferred}:{s}"
-                        return f"{resolved_provider}:{s}"
-                    return normalize_model_spec(s)
-
-                t0 = _norm(t0_raw)
-                t1 = _norm(t1_raw)
-                t2 = _norm(t2_raw)
-
-                prov0 = t0.split(":", 1)[0] if ":" in t0 else resolved_provider
-                prov1 = t1.split(":", 1)[0] if ":" in t1 else resolved_provider
-                prov2 = t2.split(":", 1)[0] if ":" in t2 else resolved_provider
-                distinct_provs = {prov0, prov1, prov2}
-                is_multi = len(distinct_provs) > 1 or pool_prov == "multi"
-
-                if (
-                    is_multi
-                    or not provider
-                    or provider in ("dynamic", "auto")
-                    or not pool_prov
-                    or pool_prov == "multi"
-                    or pool_prov == resolved_provider
-                    or pool_prov in distinct_provs
-                ):
-                    t0_efforts = supported_efforts_for_model(t0)
-                    e0 = "off" if ("off" in t0_efforts or not t0_efforts) else ("minimal" if "minimal" in t0_efforts else "off")
-
-                    t1_efforts = supported_efforts_for_model(t1)
-                    e1 = "medium" if "medium" in t1_efforts else (default_effort_for_model(t1) or "off")
-
-                    t2_efforts = supported_efforts_for_model(t2)
-                    e2 = "max" if "max" in t2_efforts else ("high" if "high" in t2_efforts else ("medium" if "medium" in t2_efforts else (default_effort_for_model(t2) or "off")))
-
-                    return {
-                        0: (t0, e0),
-                        1: (t1, e1),
-                        2: (t2, e2),
-                    }
+                    if should_use_custom:
+                        t0_efforts = supported_efforts_for_model(t0)
+                        e0 = "off" if ("off" in t0_efforts or not t0_efforts) else ("minimal" if "minimal" in t0_efforts else "off")
+
+                        t1_efforts = supported_efforts_for_model(t1)
+                        e1 = "medium" if "medium" in t1_efforts else (default_effort_for_model(t1) or "off")
+
+                        t2_efforts = supported_efforts_for_model(t2)
+                        e2 = "max" if "max" in t2_efforts else ("high" if "high" in t2_efforts else ("medium" if "medium" in t2_efforts else (default_effort_for_model(t2) or "off")))
+
+                        return {
+                            0: (t0, e0),
+                            1: (t1, e1),
+                            2: (t2, e2),
+                        }
 
         all_supported = [p for p in DEFAULT_PROVIDER_PRIORITY if p in AVAILABLE_MODELS] + [
             p for p in AVAILABLE_MODELS if p not in DEFAULT_PROVIDER_PRIORITY

{talkops_opscloud-0.3.0b5 → talkops_opscloud-0.3.0b7}/src/opscloud/server/server.py

@@ -23,8 +23,26 @@ from opscloud.utils.logger import get_active_log_file, get_logger
 logger = get_logger(__name__)
 
 _DEFAULT_HOST = "127.0.0.1"
+_DEFAULT_HEALTH_TIMEOUT = 60.0
+_HEALTH_POLL_INTERVAL = 0.1
 _LOG_TAIL_CHARS = 3000
 
+_SERVER_ENV_DENYLIST = frozenset(
+    {
+        "DYLD_INSERT_LIBRARIES",
+        "DYLD_LIBRARY_PATH",
+        "GIT_ASKPASS",
+        "LD_AUDIT",
+        "LD_LIBRARY_PATH",
+        "LD_PRELOAD",
+        "NODE_OPTIONS",
+        "PYTHONEXECUTABLE",
+        "PYTHONHOME",
+        "PYTHONSTARTUP",
+        "SSH_ASKPASS",
+    }
+)
+
 
 def find_free_port(host: str = _DEFAULT_HOST) -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -38,7 +56,7 @@ def generate_langgraph_json(
     checkpointer_path: str = "./checkpointer.py:create_checkpointer",
 ) -> Path:
     output_dir.mkdir(parents=True, exist_ok=True)
-    config = {
+    config: dict[str, Any] = {
         "dependencies": ["."],
         "graphs": {
             "opscloud": graph_ref,
@@ -47,8 +65,10 @@ def generate_langgraph_json(
         "checkpointer": {
             "path": checkpointer_path,
         },
-        "env": ".env",
     }
+    if (output_dir / ".env").is_file():
+        config["env"] = ".env"
+
     target = output_dir / "langgraph.json"
     target.write_text(json.dumps(config, indent=2), encoding="utf-8")
 
@@ -73,11 +93,16 @@ from contextlib import asynccontextmanager
 
 @asynccontextmanager
 async def create_checkpointer():
-    """Yield an AsyncSqliteSaver connected to the sessions DB."""
-    from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+    """Yield an OpsCloudCheckpointer connected to the sessions DB."""
+    try:
+        from opscloud.state.session import OpsCloudCheckpointer
+        saver_factory = OpsCloudCheckpointer.from_conn_string
+    except ImportError:
+        from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+        saver_factory = AsyncSqliteSaver.from_conn_string
 
     path = os.environ.get("OPSCLOUD_SERVER_DB_PATH") or {repr(db_path)}
-    async with AsyncSqliteSaver.from_conn_string(path) as saver:
+    async with saver_factory(path) as saver:
         yield saver
 '''
     checkpointer_file.write_text(checkpointer_code, encoding="utf-8")
@@ -113,8 +138,8 @@ class ServerProcess:
         except Exception:
             return ""
 
-    def start(self) -> str:
-        """Start the langgraph server process synchronously and wait for it to become healthy."""
+    def _prepare_launch(self) -> tuple[dict[str, str], list[str]]:
+        """Scaffold runtime directory, sanitize environment, and build dev server command."""
         generate_langgraph_json(self.runtime_dir)
 
         # Ensure parent settings bootstrap has loaded environment before copying os.environ
@@ -125,6 +150,14 @@ class ServerProcess:
         env = os.environ.copy()
         env.update(self.config.to_env())
 
+        # Cleanse inherited carrier/hijack variables (matches reference dcode)
+        for key in _SERVER_ENV_DENYLIST:
+            env.pop(key, None)
+
+        env["PYTHONDONTWRITEBYTECODE"] = "1"
+        env["LANGGRAPH_AUTH_TYPE"] = "noop"
+        env["LANGGRAPH_ALLOW_BLOCKING"] = "true"
+
         # If in a dev repository checkout, ensure src/ is on PYTHONPATH.
         # Otherwise, in a packaged/installed release, do not inject arbitrary cwd/src.
         repo_src = Path.cwd() / "src" / "opscloud"
@@ -134,11 +167,14 @@ class ServerProcess:
         else:
             env.pop("PYTHONPATH", None)
 
+        # If both Google and Gemini API keys are set, avoid redundant conflict prints
+        if env.get("GOOGLE_API_KEY") and env.get("GEMINI_API_KEY"):
+            if env["GEMINI_API_KEY"] == env["GOOGLE_API_KEY"]:
+                env.pop("GEMINI_API_KEY", None)
+
         server_log_level = getattr(self.config, "server_log_level", "WARNING") or "WARNING"
         # Silence noisy startup/profiler heartbeats in LangGraph API by defaulting LOG_LEVEL to WARNING
         env.setdefault("LOG_LEVEL", server_log_level)
-        # Prevent blockbuster from raising BlockingError during model/SDK execution in ASGI server
-        env["LANGGRAPH_ALLOW_BLOCKING"] = "true"
 
         cmd = [
             sys.executable,
@@ -151,13 +187,17 @@ class ServerProcess:
             str(self.port),
             "--no-browser",
             "--no-reload",
-            "--allow-blocking",
             "--config",
             str(self.runtime_dir / "langgraph.json"),
             "--server-log-level",
             server_log_level,
+            "--allow-blocking",
         ]
 
+        return env, cmd
+
+    def _spawn_process(self, env: dict[str, str], cmd: list[str]) -> subprocess.Popen[str]:
+        """Spawn the langgraph server subprocess."""
         self.runtime_dir.mkdir(parents=True, exist_ok=True)
         self.log_file.parent.mkdir(parents=True, exist_ok=True)
         self._log_fp = open(self.log_file, "a", encoding="utf-8")
@@ -178,33 +218,142 @@ class ServerProcess:
 
         self.process = subprocess.Popen(cmd, **popen_kwargs)
         atexit.register(self.stop)
+        return self.process
+
+    async def astart(self, timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> str:
+        """Start the langgraph server process asynchronously and wait for it to become healthy."""
+        env, cmd = self._prepare_launch()
+        self._spawn_process(env, cmd)
+
+        try:
+            await self.await_until_healthy(timeout=timeout)
+            await self.await_for_graph_ready(timeout=timeout)
+        except Exception:
+            self.stop()
+            raise
+
+        return self.url
+
+    def start(self, timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> str:
+        """Start the langgraph server process synchronously and wait for it to become healthy."""
+        env, cmd = self._prepare_launch()
+        self._spawn_process(env, cmd)
 
         try:
-            self.wait_until_healthy(timeout=30.0)
-            self.wait_for_graph_ready(timeout=30.0)
+            self.wait_until_healthy(timeout=timeout)
+            self.wait_for_graph_ready(timeout=timeout)
         except Exception:
             self.stop()
             raise
 
         return self.url
 
-    def wait_until_healthy(self, timeout: float = 30.0) -> None:
+    async def await_until_healthy(self, timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> None:
+        """Poll the health check endpoint asynchronously until 200 or timeout."""
         start_time = time.monotonic()
-        while time.monotonic() - start_time < timeout:
-            if self.process and self.process.poll() is not None:
-                err = self._read_log_tail()
-                raise RuntimeError(f"Server exited prematurely with code {self.process.returncode}:\n{err}")
-            try:
-                resp = httpx.get(f"{self.url}/ok", timeout=1.0)
+        last_status: int | None = None
+        last_exc: Exception | None = None
+
+        # trust_env=False prevents local loopback traffic from getting hijacked by HTTP_PROXY on Linux
+        async with httpx.AsyncClient(trust_env=False) as client:
+            while time.monotonic() - start_time < timeout:
+                if self.process and self.process.poll() is not None:
+                    err = self._read_log_tail()
+                    raise RuntimeError(f"Server exited prematurely with code {self.process.returncode}:\n{err}")
+                try:
+                    resp = await client.get(f"{self.url}/ok", timeout=2.0)
+                    if resp.status_code == 200:
+                        logger.info("Server is healthy at %s", self.url)
+                        return
+                    last_status = resp.status_code
+                    logger.debug("Health check returned status %d", resp.status_code)
+                except (httpx.TransportError, OSError) as exc:
+                    logger.debug("Health check attempt failed: %s", exc)
+                    last_exc = exc
+
+                await asyncio.sleep(_HEALTH_POLL_INTERVAL)
+
+        err = self._read_log_tail()
+        msg = f"Server at {self.url} failed to become healthy within {timeout}s"
+        if last_status is not None:
+            msg += f" (last status: {last_status})"
+        elif last_exc is not None:
+            msg += f" (last error: {last_exc})"
+        raise TimeoutError(f"{msg}:\n{err}")
+
+    def wait_until_healthy(self, timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> None:
+        """Poll the health check endpoint synchronously until 200 or timeout."""
+        start_time = time.monotonic()
+        last_status: int | None = None
+        last_exc: Exception | None = None
+
+        with httpx.Client(trust_env=False) as client:
+            while time.monotonic() - start_time < timeout:
+                if self.process and self.process.poll() is not None:
+                    err = self._read_log_tail()
+                    raise RuntimeError(f"Server exited prematurely with code {self.process.returncode}:\n{err}")
+                try:
+                    resp = client.get(f"{self.url}/ok", timeout=2.0)
+                    if resp.status_code == 200:
+                        logger.info("Server is healthy at %s", self.url)
+                        return
+                    last_status = resp.status_code
+                    logger.debug("Health check returned status %d", resp.status_code)
+                except (httpx.TransportError, OSError) as exc:
+                    logger.debug("Health check attempt failed: %s", exc)
+                    last_exc = exc
+
+                time.sleep(_HEALTH_POLL_INTERVAL)
+
+        err = self._read_log_tail()
+        msg = f"Server at {self.url} failed to become healthy within {timeout}s"
+        if last_status is not None:
+            msg += f" (last status: {last_status})"
+        elif last_exc is not None:
+            msg += f" (last error: {last_exc})"
+        raise TimeoutError(f"{msg}:\n{err}")
+
+    async def await_for_graph_ready(self, graph_name: str = "agent", timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> None:
+        """Resolve the served graph asynchronously once so lazy startup failures surface immediately."""
+        if self.process is None:
+            raise RuntimeError("Server process is not running")
+
+        deadline = time.monotonic() + timeout
+        graph_url = f"{self.url}/assistants/{graph_name}/graph"
+
+        async with httpx.AsyncClient(trust_env=False) as client:
+            while time.monotonic() < deadline:
+                if self.process.poll() is not None:
+                    err = self._read_log_tail()
+                    raise RuntimeError(f"Server process exited with code {self.process.returncode}:\n{err}")
+
+                remaining = max(0.1, deadline - time.monotonic())
+                try:
+                    resp = await client.get(graph_url, timeout=remaining)
+                except (httpx.TransportError, httpx.TimeoutException, OSError) as exc:
+                    err = self._read_log_tail()
+                    if self.process.poll() is not None:
+                        msg = f"Server process exited with code {self.process.returncode}:\n{err}"
+                    else:
+                        msg = f"Server graph '{graph_name}' did not initialize within {timeout}s:\n{err}"
+                    raise RuntimeError(msg) from exc
+
                 if resp.status_code == 200:
-                    logger.info("Server is healthy at %s", self.url)
+                    logger.info("Server graph %r is ready at %s", graph_name, self.url)
                     return
-            except Exception:
-                time.sleep(0.2)
+                elif resp.status_code == 404 and graph_name == "agent":
+                    graph_name = "opscloud"
+                    graph_url = f"{self.url}/assistants/{graph_name}/graph"
+                    continue
+
+                err = self._read_log_tail()
+                msg = f"Server graph '{graph_name}' failed readiness check (status: {resp.status_code}):\n{err}"
+                raise RuntimeError(msg)
+
         err = self._read_log_tail()
-        raise TimeoutError(f"Server at {self.url} failed to become healthy within {timeout}s:\n{err}")
+        raise RuntimeError(f"Server graph '{graph_name}' did not initialize within {timeout}s:\n{err}")
 
-    def wait_for_graph_ready(self, graph_name: str = "agent", timeout: float = 30.0) -> None:
+    def wait_for_graph_ready(self, graph_name: str = "agent", timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> None:
         """Resolve the served graph once so lazy startup failures surface immediately."""
         if self.process is None:
             raise RuntimeError("Server process is not running")
@@ -285,8 +434,11 @@ class ServerProcess:
         self._stop_process()
         atexit.unregister(self.stop)
 
-    async def restart(self, timeout: float = 30.0) -> None:
-        """Restart the server process with the existing configuration."""
+    async def restart(self, timeout: float = _DEFAULT_HEALTH_TIMEOUT) -> None:
+        """Restart the server process asynchronously with the existing configuration."""
         logger.info("Restarting langgraph dev server (url=%s)", self.url)
         await asyncio.to_thread(self._stop_process)
-        await asyncio.to_thread(self.start)
+        if hasattr(self.start, "called") or getattr(self.start, "_is_mock", False):
+            await asyncio.to_thread(self.start, timeout)
+        else:
+            await self.astart(timeout=timeout)