Robert Helewka
ca7d714a31
Add two new sections to the Pallas documentation: - Sampling parameters: explain that temperature/top_p/top_k are configured via the fast-agent decorator's `request_params`, with a provider support matrix and a note on Claude Opus 4.7 stripping these params in favor of `output_config.effort`. - Metrics: document the Prometheus `/metrics` endpoint exposed on the registry port, including scrape config, full metrics reference table, and notes on where each metric is captured.
722 lines
28 KiB
Python
722 lines
28 KiB
Python
"""
|
|
Health check module for Pallas.
|
|
|
|
Probes downstream MCP server connectivity and exposes a ``get_health`` MCP
|
|
tool. At startup, :func:`validate_llm_providers` runs a cheap, per-provider
|
|
preflight so Daedalus (or any headless consumer) can see *why* an agent
|
|
might be degraded when there is no fast-agent TUI to surface it — see
|
|
``docs/pallas_integration.md`` § Runtime get_health tool.
|
|
|
|
Preflight dispatch matrix:
|
|
|
|
========================== ======================================== =======================================
|
|
Provider Probe Success criterion
|
|
========================== ======================================== =======================================
|
|
``anthropic`` (direct) ``GET {base_url}/models/{model}`` HTTP 200
|
|
``anthropic`` (Mantle) ``GET {mantle_root}/v1/models/{wire}`` HTTP 200 (wire-name via mantle_shims)
|
|
``openai`` ``GET {base_url}/models`` HTTP 200 and active model in list
|
|
``generic`` ``GET {base_url}/models`` HTTP 200 (body not inspected)
|
|
``bedrock`` none ok if AWS creds resolvable
|
|
unknown provider none error — surfaces honestly to Daedalus
|
|
========================== ======================================== =======================================
|
|
|
|
API keys are resolved via :class:`fast_agent.llm.provider_key_manager.ProviderKeyManager`
|
|
so this module sees identical secret-loading behaviour to the real LLM
|
|
client path. We never duplicate key-resolution logic here.
|
|
|
|
Endpoints that cost inference tokens are deliberately avoided. Mantle's
|
|
anthropic probe uses the (token-free) model catalogue at the region root
|
|
(``/v1/models/{wire}``), *not* ``POST /anthropic/v1/messages``.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import httpx
|
|
import yaml
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _config_root() -> Path:
|
|
"""Return the working directory where agents.yaml and fastagent configs live."""
|
|
return Path.cwd()
|
|
|
|
|
|
def _load_deployment_name() -> str:
|
|
"""Read the deployment name from agents.yaml (or PALLAS_AGENTS_CONFIG override)."""
|
|
config_path = _config_root() / os.environ.get("PALLAS_AGENTS_CONFIG", "agents.yaml")
|
|
if config_path.exists():
|
|
data = yaml.safe_load(config_path.read_text()) or {}
|
|
return data.get("name", "pallas")
|
|
return "pallas"
|
|
|
|
|
|
_DEPLOY_NAME = _load_deployment_name()
|
|
|
|
# ── Default endpoints (only used when the provider section is missing) ───────
|
|
|
|
_ANTHROPIC_DEFAULT_API = "https://api.anthropic.com/v1"
|
|
_OPENAI_DEFAULT_API = "https://api.openai.com/v1"
|
|
_GENERIC_DEFAULT_API = "http://localhost:11434/v1"
|
|
|
|
# Populated by validate_llm_providers() at startup, refreshed on a TTL by
|
|
# _refresh_llm_status_if_stale() from inside the get_health MCP tool. The
|
|
# TTL re-probe is what makes the agent self-heal when an upstream LLM was
|
|
# briefly unreachable at boot — without it the cached error sticks until
|
|
# the process restarts. See docs/pallas.md § Health System.
|
|
_llm_status: dict[str, dict] = {}
|
|
_active_provider: str = ""
|
|
_llm_status_ts: float = 0.0 # time.monotonic() of last successful probe; 0 = never
|
|
_llm_refresh_lock = asyncio.Lock()
|
|
_LLM_HEALTH_TTL_S = float(os.environ.get("PALLAS_LLM_HEALTH_TTL", "60"))
|
|
|
|
|
|
# ── Config loading ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def _load_dotenv() -> None:
|
|
"""Load .env file into os.environ (without overwriting existing vars)."""
|
|
env_path = _config_root() / ".env"
|
|
if not env_path.exists():
|
|
return
|
|
for line in env_path.read_text().splitlines():
|
|
line = line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
key, _, value = line.partition("=")
|
|
key = key.strip()
|
|
value = value.strip()
|
|
if key and key not in os.environ:
|
|
os.environ[key] = value
|
|
|
|
|
|
def _expand_env(value: str) -> str:
|
|
"""Replace ${VAR} placeholders with environment variable values."""
|
|
return re.sub(
|
|
r"\$\{([^}]+)\}",
|
|
lambda m: os.environ.get(m.group(1), ""),
|
|
value,
|
|
)
|
|
|
|
|
|
def _expand_env_in_tree(obj: Any) -> Any:
|
|
"""Recursively expand ${ENV_VAR} placeholders inside a parsed YAML tree."""
|
|
if isinstance(obj, str):
|
|
return _expand_env(obj)
|
|
if isinstance(obj, dict):
|
|
return {k: _expand_env_in_tree(v) for k, v in obj.items()}
|
|
if isinstance(obj, list):
|
|
return [_expand_env_in_tree(v) for v in obj]
|
|
return obj
|
|
|
|
|
|
def _load_config() -> tuple[dict, dict]:
|
|
"""Load fastagent config and secrets YAML from the working directory."""
|
|
root = _config_root()
|
|
config = yaml.safe_load((root / "fastagent.config.yaml").read_text()) or {}
|
|
secrets_path = root / "fastagent.secrets.yaml"
|
|
secrets = yaml.safe_load(secrets_path.read_text()) if secrets_path.exists() else {}
|
|
return config, secrets
|
|
|
|
|
|
def _merge_for_key_manager(config: dict, secrets: dict) -> dict:
|
|
"""Produce the merged, env-expanded dict that ProviderKeyManager expects.
|
|
|
|
``ProviderKeyManager.get_config_file_key`` looks for ``<provider>.api_key``
|
|
in a single flat dict. It does not apply ``${ENV_VAR}`` expansion itself,
|
|
so we pre-expand the whole tree.
|
|
"""
|
|
merged: dict = {}
|
|
for source in (config or {}, secrets or {}):
|
|
for provider_name, settings in (source or {}).items():
|
|
if isinstance(settings, dict):
|
|
target = merged.setdefault(provider_name, {})
|
|
if isinstance(target, dict):
|
|
target.update(settings)
|
|
return _expand_env_in_tree(merged)
|
|
|
|
|
|
# ── Per-provider probes ──────────────────────────────────────────────────────
|
|
|
|
|
|
async def _check_anthropic(
|
|
client: httpx.AsyncClient, api_key: str, model_id: str, base_url: str
|
|
) -> str | None:
|
|
"""Validate an Anthropic model via ``GET {base_url}/models/{model_id}``.
|
|
|
|
Works for both the public ``api.anthropic.com`` endpoint and for AWS
|
|
Bedrock Mantle's region-root model catalogue
|
|
(``https://bedrock-mantle.{region}.api.aws/v1/models/{wire_id}``).
|
|
The caller is responsible for passing the correct ``base_url`` and the
|
|
wire-name form of ``model_id`` for Mantle — see ``pallas.mantle_shims``.
|
|
"""
|
|
try:
|
|
resp = await client.get(
|
|
f"{base_url.rstrip('/')}/models/{model_id}",
|
|
headers={
|
|
"x-api-key": api_key,
|
|
"anthropic-version": "2023-06-01",
|
|
},
|
|
)
|
|
except Exception as exc:
|
|
return f"API unreachable ({type(exc).__name__})"
|
|
if resp.status_code == 200:
|
|
return None
|
|
if resp.status_code == 404:
|
|
return f"model '{model_id}' not found"
|
|
return f"API request failed ({resp.status_code})"
|
|
|
|
|
|
async def _list_openai_models(
|
|
client: httpx.AsyncClient, api_key: str, base_url: str
|
|
) -> tuple[str | None, list[str]]:
|
|
"""List models from an OpenAI-compatible API. Returns (error, model_ids)."""
|
|
try:
|
|
resp = await client.get(
|
|
f"{base_url.rstrip('/')}/models",
|
|
headers={"Authorization": f"Bearer {api_key}"},
|
|
)
|
|
except Exception as exc:
|
|
return f"API unreachable ({type(exc).__name__})", []
|
|
if resp.status_code != 200:
|
|
return f"API request failed ({resp.status_code})", []
|
|
try:
|
|
data = resp.json()
|
|
except Exception:
|
|
return "response was not valid JSON", []
|
|
models = [m["id"] for m in data.get("data", []) if isinstance(m, dict) and "id" in m]
|
|
return None, models
|
|
|
|
|
|
async def _check_generic(
|
|
client: httpx.AsyncClient, base_url: str
|
|
) -> str | None:
|
|
"""Status-code-only probe against ``{base_url}/models``.
|
|
|
|
The generic provider targets local/on-prem OpenAI-compatible servers
|
|
(llama.cpp, Ollama, vLLM, …) whose ``/v1/models`` payloads are not all
|
|
identical — llama.cpp mixes an Ollama-style ``models`` list with the
|
|
OpenAI ``data`` list, for example. We deliberately don't require the
|
|
configured model name to appear in the response because users hot-swap
|
|
models by name all the time; as long as the server is up and returns
|
|
200 for its catalogue we call it ok.
|
|
"""
|
|
try:
|
|
resp = await client.get(f"{base_url.rstrip('/')}/models")
|
|
except Exception as exc:
|
|
return f"API unreachable ({type(exc).__name__})"
|
|
if resp.status_code == 200:
|
|
return None
|
|
return f"API request failed ({resp.status_code})"
|
|
|
|
|
|
# ── Mantle helpers ───────────────────────────────────────────────────────────
|
|
|
|
|
|
def _mantle_root_from_anthropic_base(base: str) -> str:
|
|
"""Return the region root for a Mantle anthropic base_url.
|
|
|
|
Mantle publishes its inference path under ``/anthropic`` but the
|
|
catalogue (``GET /v1/models/{wire_id}``) lives at the region root.
|
|
Example:
|
|
``https://bedrock-mantle.us-east-1.api.aws/anthropic``
|
|
→ ``https://bedrock-mantle.us-east-1.api.aws``
|
|
Any other trailing paths are returned untouched.
|
|
"""
|
|
stripped = base.rstrip("/")
|
|
if stripped.endswith("/anthropic"):
|
|
return stripped[: -len("/anthropic")]
|
|
return stripped
|
|
|
|
|
|
# ── Preflight orchestration ──────────────────────────────────────────────────
|
|
|
|
|
|
async def _preflight_anthropic(
|
|
client: httpx.AsyncClient, config: dict, secrets: dict, active_model: str
|
|
) -> dict:
|
|
from fast_agent.core.exceptions import ProviderKeyError
|
|
from fast_agent.llm.provider_key_manager import ProviderKeyManager
|
|
from pallas.mantle_shims import MANTLE_WIRE_NAMES, is_mantle_base_url
|
|
|
|
merged = _merge_for_key_manager(config, secrets)
|
|
try:
|
|
api_key = ProviderKeyManager.get_api_key("anthropic", merged)
|
|
except ProviderKeyError as exc:
|
|
return {"status": "error", "message": str(exc)}
|
|
|
|
anthropic_base = _expand_env(
|
|
(config.get("anthropic", {}) or {}).get("base_url", "")
|
|
) or os.environ.get("ANTHROPIC_BASE_URL", "") or _ANTHROPIC_DEFAULT_API
|
|
|
|
if is_mantle_base_url(anthropic_base):
|
|
# Mantle hosts the model catalogue at the region root, not under
|
|
# /anthropic. Wire-name translation (claude-opus-4-7 →
|
|
# anthropic.claude-opus-4-7) keeps us consistent with mantle_shims.
|
|
probe_base = f"{_mantle_root_from_anthropic_base(anthropic_base)}/v1"
|
|
wire_id = MANTLE_WIRE_NAMES.get(active_model, active_model)
|
|
err = await _check_anthropic(client, api_key, wire_id, probe_base)
|
|
if err:
|
|
logger.warning("anthropic (mantle, %s): %s", anthropic_base, err)
|
|
return {"status": "error", "model": wire_id, "message": err}
|
|
logger.info("anthropic (mantle, %s): %s ready", anthropic_base, wire_id)
|
|
return {"status": "ok", "model": wire_id}
|
|
|
|
err = await _check_anthropic(client, api_key, active_model, anthropic_base)
|
|
if err:
|
|
logger.warning("anthropic (%s): %s", anthropic_base, err)
|
|
return {"status": "error", "model": active_model, "message": err}
|
|
logger.info("anthropic (%s): %s ready", anthropic_base, active_model)
|
|
return {"status": "ok", "model": active_model}
|
|
|
|
|
|
async def _preflight_openai(
|
|
client: httpx.AsyncClient, config: dict, secrets: dict, active_model: str
|
|
) -> dict:
|
|
from fast_agent.core.exceptions import ProviderKeyError
|
|
from fast_agent.llm.provider_key_manager import ProviderKeyManager
|
|
|
|
merged = _merge_for_key_manager(config, secrets)
|
|
try:
|
|
api_key = ProviderKeyManager.get_api_key("openai", merged)
|
|
except ProviderKeyError as exc:
|
|
return {"status": "error", "message": str(exc)}
|
|
|
|
openai_base = _expand_env(
|
|
(config.get("openai", {}) or {}).get("base_url", "")
|
|
) or os.environ.get("OPENAI_BASE_URL", "") or _OPENAI_DEFAULT_API
|
|
|
|
err, models = await _list_openai_models(client, api_key, openai_base)
|
|
if err:
|
|
logger.warning("openai (%s): %s", openai_base, err)
|
|
return {"status": "error", "message": err}
|
|
if active_model and active_model not in models:
|
|
label = ", ".join(models) if models else "none"
|
|
msg = f"model '{active_model}' not found (available: {label})"
|
|
logger.warning("openai (%s): %s", openai_base, msg)
|
|
return {"status": "error", "model": active_model, "message": msg}
|
|
logger.info("openai (%s): %s ready", openai_base, active_model or "(any)")
|
|
return {"status": "ok", "model": active_model, "models": models}
|
|
|
|
|
|
async def _preflight_generic(
|
|
client: httpx.AsyncClient, config: dict, secrets: dict, active_model: str
|
|
) -> dict:
|
|
# generic has a synthetic "ollama" key via ProviderKeyManager, so there's
|
|
# nothing to authenticate against — we just need the endpoint.
|
|
generic_base = _expand_env(
|
|
(config.get("generic", {}) or {}).get("base_url", "")
|
|
) or os.environ.get("GENERIC_BASE_URL", "") or _GENERIC_DEFAULT_API
|
|
|
|
err = await _check_generic(client, generic_base)
|
|
if err:
|
|
logger.warning("generic (%s): %s", generic_base, err)
|
|
return {"status": "error", "model": active_model, "message": err}
|
|
logger.info("generic (%s): %s ready", generic_base, active_model or "(any)")
|
|
return {"status": "ok", "model": active_model}
|
|
|
|
|
|
def _preflight_bedrock(config: dict, secrets: dict, active_model: str) -> dict:
|
|
"""Bedrock uses the AWS credential chain — no outbound HTTP here.
|
|
|
|
We report ``ok`` whenever any of the usual credential sources is present
|
|
(long-term bedrock key, explicit access key pair, or a nonempty AWS
|
|
profile). If nothing is set we mark it degraded so Daedalus shows the
|
|
operator *why* the first real request will fail; we don't actually call
|
|
STS or Bedrock ourselves.
|
|
"""
|
|
have_bearer = bool(os.environ.get("AWS_BEARER_TOKEN_BEDROCK"))
|
|
have_access_key = bool(os.environ.get("AWS_ACCESS_KEY_ID")) and bool(
|
|
os.environ.get("AWS_SECRET_ACCESS_KEY")
|
|
)
|
|
have_profile = bool(os.environ.get("AWS_PROFILE"))
|
|
creds_path = Path.home() / ".aws" / "credentials"
|
|
have_file = creds_path.exists()
|
|
|
|
if have_bearer or have_access_key or have_profile or have_file:
|
|
logger.info("bedrock: credentials resolvable (no preflight request issued)")
|
|
return {"status": "ok", "model": active_model}
|
|
|
|
msg = "no AWS credentials found (set AWS_BEARER_TOKEN_BEDROCK or configure AWS CLI)"
|
|
logger.warning("bedrock: %s", msg)
|
|
return {"status": "error", "model": active_model, "message": msg}
|
|
|
|
|
|
async def _probe_active_provider(
|
|
client: httpx.AsyncClient,
|
|
config: dict,
|
|
secrets: dict,
|
|
active_provider: str,
|
|
active_model: str,
|
|
) -> dict:
|
|
"""Run the preflight probe for a single provider and return its result dict.
|
|
|
|
Shared by the startup path (:func:`validate_llm_providers`) and the
|
|
runtime TTL refresh path (:func:`_refresh_llm_status_if_stale`). Keeps
|
|
the dispatch matrix from being duplicated.
|
|
"""
|
|
if active_provider == "anthropic":
|
|
return await _preflight_anthropic(client, config, secrets, active_model)
|
|
if active_provider == "openai":
|
|
return await _preflight_openai(client, config, secrets, active_model)
|
|
if active_provider == "generic":
|
|
return await _preflight_generic(client, config, secrets, active_model)
|
|
if active_provider == "bedrock":
|
|
return _preflight_bedrock(config, secrets, active_model)
|
|
|
|
# Known to fast-agent? Surface that gap explicitly rather than silently
|
|
# reporting "error" from an empty dict lookup later.
|
|
try:
|
|
from fast_agent.llm.provider_types import Provider
|
|
|
|
Provider(active_provider) # raises ValueError if unknown
|
|
msg = (
|
|
f"preflight for provider '{active_provider}' is not "
|
|
"implemented in pallas.health; LLM health will be "
|
|
"validated on first inference call"
|
|
)
|
|
logger.info("%s: %s", active_provider, msg)
|
|
return {"status": "ok", "model": active_model, "message": msg}
|
|
except ValueError:
|
|
msg = f"unknown provider '{active_provider}' in default_model"
|
|
logger.warning(msg)
|
|
return {"status": "error", "message": msg}
|
|
|
|
|
|
async def _refresh_llm_status_if_stale(timeout: float = 5.0) -> None:
|
|
"""Re-probe the active LLM provider when the cached status is older than
|
|
``PALLAS_LLM_HEALTH_TTL`` seconds. No-op on a fresh cache or when no
|
|
active provider has been resolved yet.
|
|
|
|
On any unexpected exception, the previous cache entry is left intact —
|
|
a real probe result is better signal than our own internal failure.
|
|
"""
|
|
global _llm_status_ts
|
|
|
|
if not _active_provider:
|
|
return
|
|
if time.monotonic() - _llm_status_ts < _LLM_HEALTH_TTL_S:
|
|
return
|
|
|
|
async with _llm_refresh_lock:
|
|
# Re-check under the lock — another coroutine may have refreshed
|
|
# while we were waiting.
|
|
if time.monotonic() - _llm_status_ts < _LLM_HEALTH_TTL_S:
|
|
return
|
|
try:
|
|
_load_dotenv()
|
|
config, secrets = _load_config()
|
|
default_model = config.get("default_model", "") or ""
|
|
if "." not in default_model:
|
|
return
|
|
active_provider, active_model = default_model.split(".", 1)
|
|
if active_provider != _active_provider:
|
|
# default_model was edited under us — defer to the next
|
|
# full validate_llm_providers() rather than guessing.
|
|
return
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
result = await _probe_active_provider(
|
|
client, config, secrets, active_provider, active_model
|
|
)
|
|
_llm_status[_active_provider] = result
|
|
_llm_status_ts = time.monotonic()
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"llm health re-probe failed (%s); keeping previous cache: %s",
|
|
_active_provider,
|
|
exc,
|
|
)
|
|
|
|
|
|
async def validate_llm_providers(timeout: float = 5.0) -> dict[str, dict]:
|
|
"""
|
|
Validate the configured LLM provider and populate the module-level cache
|
|
read by :func:`get_health` on every MCP ``get_health`` tool call.
|
|
|
|
Only the *active* provider (the one named by ``default_model``) is
|
|
preflighted — that's the one whose failure would actually break the
|
|
agent, and it keeps the startup surface small. Other provider sections
|
|
are ignored here even if they're configured.
|
|
|
|
Returns a dict keyed by provider name with validation results. Shape:
|
|
|
|
.. code-block:: python
|
|
|
|
{"anthropic": {"status": "ok", "model": "anthropic.claude-opus-4-7"}}
|
|
{"generic": {"status": "ok", "model": "Qwen3.5-..."}}
|
|
{"openai": {"status": "error", "message": "API request failed (401)"}}
|
|
{"unknown": {"status": "error", "message": "unknown provider 'foo'"}}
|
|
"""
|
|
global _active_provider, _llm_status_ts
|
|
|
|
_load_dotenv()
|
|
config, secrets = _load_config()
|
|
default_model = config.get("default_model", "") or ""
|
|
|
|
if "." not in default_model:
|
|
msg = (
|
|
f"default_model '{default_model}' is missing a provider prefix "
|
|
"(expected '<provider>.<model>')"
|
|
)
|
|
logger.warning(msg)
|
|
results = {"unknown": {"status": "error", "message": msg}}
|
|
_llm_status.clear()
|
|
_llm_status.update(results)
|
|
_active_provider = "unknown"
|
|
_llm_status_ts = time.monotonic()
|
|
return results
|
|
|
|
active_provider, active_model = default_model.split(".", 1)
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
result = await _probe_active_provider(
|
|
client, config, secrets, active_provider, active_model
|
|
)
|
|
results: dict[str, dict] = {active_provider: result}
|
|
|
|
_llm_status.clear()
|
|
_llm_status.update(results)
|
|
_active_provider = active_provider
|
|
_llm_status_ts = time.monotonic()
|
|
|
|
try:
|
|
from pallas import metrics as _pallas_metrics
|
|
|
|
_pallas_metrics.llm_provider_up.labels(provider=active_provider).set(
|
|
1.0 if result.get("status") == "ok" else 0.0
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
return results
|
|
|
|
|
|
# ── Downstream MCP server probing ────────────────────────────────────────────
|
|
|
|
|
|
async def check_downstream_health(
|
|
servers: dict[str, dict], timeout: float = 3.0
|
|
) -> dict:
|
|
"""
|
|
Probe downstream MCP servers and return aggregate health status.
|
|
|
|
Args:
|
|
servers: Mapping of server name to {"url": str, "headers": dict}.
|
|
Headers may contain ${ENV_VAR} placeholders which are expanded
|
|
before the request is sent.
|
|
timeout: Per-request timeout in seconds.
|
|
|
|
Returns:
|
|
{"status": "ok"|"degraded", "timestamp": "...", "message": "..."}
|
|
"""
|
|
_load_dotenv()
|
|
|
|
async def _probe(
|
|
client: httpx.AsyncClient, name: str, cfg: dict
|
|
) -> tuple[str, bool, str]:
|
|
url = cfg.get("url", "")
|
|
raw_headers = cfg.get("headers", {})
|
|
headers = {k: _expand_env(str(v)) for k, v in raw_headers.items()}
|
|
try:
|
|
common_headers = {
|
|
"Accept": "application/json, text/event-stream",
|
|
"Content-Type": "application/json",
|
|
**headers,
|
|
}
|
|
resp = await client.post(
|
|
url,
|
|
headers=common_headers,
|
|
json={
|
|
"jsonrpc": "2.0",
|
|
"method": "initialize",
|
|
"id": 1,
|
|
"params": {
|
|
"protocolVersion": "2025-03-26",
|
|
"capabilities": {},
|
|
"clientInfo": {
|
|
"name": f"{_DEPLOY_NAME}-health",
|
|
"version": "1.0.0",
|
|
},
|
|
},
|
|
},
|
|
)
|
|
if resp.status_code >= 400:
|
|
return name, False, f"HTTP {resp.status_code}"
|
|
# Tear down the session so we don't leak server-side state.
|
|
session_id = resp.headers.get("mcp-session-id")
|
|
if session_id:
|
|
try:
|
|
await client.delete(
|
|
url,
|
|
headers={**headers, "mcp-session-id": session_id},
|
|
)
|
|
except Exception:
|
|
pass # best-effort cleanup
|
|
return name, True, ""
|
|
except Exception as exc:
|
|
return name, False, type(exc).__name__
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
results = await asyncio.gather(
|
|
*(_probe(client, name, cfg) for name, cfg in servers.items())
|
|
)
|
|
|
|
now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
failures = sorted(
|
|
(f"{name} ({reason})" if reason else name)
|
|
for name, ok, reason in results
|
|
if not ok
|
|
)
|
|
|
|
if not failures:
|
|
return {"status": "ok", "timestamp": now}
|
|
|
|
return {
|
|
"status": "degraded",
|
|
"timestamp": now,
|
|
"message": f"Unreachable: {', '.join(failures)}",
|
|
}
|
|
|
|
|
|
def register_health_tool(
|
|
mcp_server, servers: dict[str, dict], agent_name: str = "unknown"
|
|
) -> None:
|
|
"""Register a get_health MCP tool on the given FastMCP server instance.
|
|
|
|
``agent_name`` labels the Prometheus gauges populated as a side effect
|
|
of every probe (downstream reachability + overall status).
|
|
"""
|
|
from pallas import metrics as _pallas_metrics
|
|
|
|
@mcp_server.tool(
|
|
name="get_health",
|
|
description="Returns the health status of this agent and its downstream dependencies.",
|
|
)
|
|
async def get_health() -> str:
|
|
await _refresh_llm_status_if_stale()
|
|
result, per_server_ok = await _check_downstream_with_breakdown(servers)
|
|
# Include LLM provider status from startup preflight (active provider only)
|
|
llm_for_metrics: dict[str, str] = {}
|
|
if _active_provider:
|
|
active = _llm_status.get(_active_provider)
|
|
llm_for_metrics[_active_provider] = (
|
|
active.get("status", "error") if active else "error"
|
|
)
|
|
if active is None:
|
|
# Should be unreachable after the rewrite (validate_llm_providers
|
|
# always populates _llm_status for _active_provider). Keep a
|
|
# belt-and-braces path so a future refactor can't regress into
|
|
# silently reporting "error".
|
|
err_msg = (
|
|
f"LLM: {_active_provider}: provider not preflighted"
|
|
)
|
|
result["status"] = "degraded"
|
|
existing = result.get("message", "")
|
|
result["message"] = f"{existing}; {err_msg}" if existing else err_msg
|
|
elif active.get("status") != "ok":
|
|
err_msg = (
|
|
f"LLM: {_active_provider}: "
|
|
f"{active.get('message', 'unknown error')}"
|
|
)
|
|
result["status"] = "degraded"
|
|
existing = result.get("message", "")
|
|
result["message"] = f"{existing}; {err_msg}" if existing else err_msg
|
|
try:
|
|
_pallas_metrics.record_health_probe(
|
|
agent_name,
|
|
overall_status=result.get("status", "error"),
|
|
downstream=per_server_ok,
|
|
llm=llm_for_metrics,
|
|
)
|
|
except Exception:
|
|
pass
|
|
return json.dumps(result)
|
|
|
|
|
|
async def _check_downstream_with_breakdown(
|
|
servers: dict[str, dict], timeout: float = 3.0
|
|
) -> tuple[dict, dict[str, bool]]:
|
|
"""Like :func:`check_downstream_health` but also returns per-server ok flags.
|
|
|
|
Kept as a thin wrapper so external callers of ``check_downstream_health``
|
|
(if any) stay on the original dict-only contract.
|
|
"""
|
|
_load_dotenv()
|
|
|
|
async def _probe(
|
|
client: httpx.AsyncClient, name: str, cfg: dict
|
|
) -> tuple[str, bool, str]:
|
|
url = cfg.get("url", "")
|
|
raw_headers = cfg.get("headers", {})
|
|
headers = {k: _expand_env(str(v)) for k, v in raw_headers.items()}
|
|
try:
|
|
common_headers = {
|
|
"Accept": "application/json, text/event-stream",
|
|
"Content-Type": "application/json",
|
|
**headers,
|
|
}
|
|
resp = await client.post(
|
|
url,
|
|
headers=common_headers,
|
|
json={
|
|
"jsonrpc": "2.0",
|
|
"method": "initialize",
|
|
"id": 1,
|
|
"params": {
|
|
"protocolVersion": "2025-03-26",
|
|
"capabilities": {},
|
|
"clientInfo": {
|
|
"name": f"{_DEPLOY_NAME}-health",
|
|
"version": "1.0.0",
|
|
},
|
|
},
|
|
},
|
|
)
|
|
if resp.status_code >= 400:
|
|
return name, False, f"HTTP {resp.status_code}"
|
|
session_id = resp.headers.get("mcp-session-id")
|
|
if session_id:
|
|
try:
|
|
await client.delete(
|
|
url,
|
|
headers={**headers, "mcp-session-id": session_id},
|
|
)
|
|
except Exception:
|
|
pass
|
|
return name, True, ""
|
|
except Exception as exc:
|
|
return name, False, type(exc).__name__
|
|
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
results = await asyncio.gather(
|
|
*(_probe(client, name, cfg) for name, cfg in servers.items())
|
|
)
|
|
|
|
now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
per_server_ok = {name: ok for name, ok, _ in results}
|
|
failures = sorted(
|
|
(f"{name} ({reason})" if reason else name)
|
|
for name, ok, reason in results
|
|
if not ok
|
|
)
|
|
|
|
if not failures:
|
|
return {"status": "ok", "timestamp": now}, per_server_ok
|
|
|
|
return (
|
|
{
|
|
"status": "degraded",
|
|
"timestamp": now,
|
|
"message": f"Unreachable: {', '.join(failures)}",
|
|
},
|
|
per_server_ok,
|
|
)
|