diff --git a/.chainlit/config.toml b/.chainlit/config.toml index a06c5ee83..62d47ca1a 100644 --- a/.chainlit/config.toml +++ b/.chainlit/config.toml @@ -128,7 +128,8 @@ cot = "full" # Specify a CSS file that can be used to customize the user interface. # The CSS file can be served from the public directory or via an external link. -# custom_css = "/public/test.css" +# biomni.css lets the settings dialog's help text keep its paragraph breaks. +custom_css = "/public/biomni.css" # Specify additional attributes for a custom CSS file # custom_css_attributes = "media=\"print\"" @@ -157,6 +158,9 @@ alert_style = "classic" # custom_meta_image_url = "https://chainlit-cloud.s3.eu-west-3.amazonaws.com/logo/chainlit_banner.png" # Load assistant logo directly from URL. +# Only a signed-in browser is sent this setting. The sign-in page shows +# public/logo_dark.png or logo_light.png instead - copies of the same image, +# there so that page does not fall back to the Chainlit logo. logo_file_url = "/public/avatars/ad1_256.png" # Load assistant avatar image directly from URL. diff --git a/.env.example b/.env.example index 7bbe112b6..c85161628 100644 --- a/.env.example +++ b/.env.example @@ -28,8 +28,18 @@ AWS_BEARER_TOKEN_BEDROCK=your_bedrock_api_key_here AWS_REGION=us-east-1 # Optional: Custom model serving configuration -# CUSTOM_MODEL_BASE_URL=http://localhost:8000/v1 -# CUSTOM_MODEL_API_KEY=your_custom_api_key_here +# BIOMNI_CUSTOM_BASE_URL=http://localhost:8000/v1 +# BIOMNI_CUSTOM_API_KEY=your_custom_api_key_here + +# Optional: the platform LLM proxy (AD Workbench / GRIP; see docs/grip_deployment.md). +# When set, every model call goes through it and the provider keys above are +# unused. Each request carries the proxy's token as a bearer token and the +# session's user and workspace ids as X-User-Id / X-Workspace-Id - ids that come +# from the authentication gateway, so this needs BIOMNI_TRUST_AUTH_HEADERS below. +# BIOMNI_LLM_PROXY_URL=https:///v1 +# BIOMNI_LLM_PROXY_API_KEY= +# BIOMNI_LLM_PROXY_SCHEMA=openai # or anthropic, for /v1/messages +# BIOMNI_LLM= # Optional: Biomni data path (defaults to ~/.biomni/data) # BIOMNI_DATA_PATH=/path/to/your/data @@ -52,7 +62,9 @@ AWS_REGION=us-east-1 # Optional: put a sign-in in front of a deployment that has no authentication # gateway (a demo box on a public IP). One shared password, any username - and # the username becomes the identity, so each person gets their own conversation -# history. Leave unset wherever a gateway already authenticates. +# history. Separate, not private: anyone with the password can sign in under a +# username someone else used. Leave unset wherever a gateway already +# authenticates. # BIOMNI_DEMO_PASSWORD= # Optional: address the container publishes on. Set to 127.0.0.1 when a TLS @@ -97,7 +109,10 @@ AWS_REGION=us-east-1 # BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER=ai-app-workspace-context # # Single-value headers, for other gateways (oauth2-proxy, Envoy/ext_authz, ...). -# Used for any field the context headers above did not supply. +# Read only for a request that carries neither context header above, so a +# client cannot fill in a field the gateway left out. GRIP needs none of these: +# pointing one at an Ai-App-* header has no effect (those are parsed field by +# field), and the app logs a warning at startup saying so. # BIOMNI_AUTH_USER_ID_HEADER=x-auth-request-user-id # BIOMNI_AUTH_EMAIL_HEADER=x-auth-request-email # BIOMNI_AUTH_WORKSPACE_HEADER=x-workspace-id diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4f1e08935..44ba3789d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -90,4 +90,13 @@ jobs: biomni/config.py \ biomni/credentials.py \ biomni/env_probe.py \ - biomni/llm.py + biomni/health.py \ + biomni/identity.py \ + biomni/llm.py \ + biomni/llm_proxy.py \ + biomni/observability.py \ + biomni/status.py \ + biomni/tool/availability.py \ + biomni/workspace_prefs.py \ + chainlit_ui/graph_stream.py \ + chainlit_ui/llm_failures.py diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml index 83de5413f..493acae69 100644 --- a/.github/workflows/docker.yml +++ b/.github/workflows/docker.yml @@ -188,12 +188,13 @@ jobs: # Tag strategy: # - branch name (refs/heads/main → :main; / in feat branches # sanitised to - per Docker tag rules) - # - git SHA short (always - primary content-addressed tag, - # what production should pin to) + # - git SHA short (always, as :sha-<7 hex> - primary + # content-addressed tag, what production should + # pin to) # - :latest (only when pushing the default branch - # conventional "newest published" alias) # - semver aliases (v1.2.3 git tag → :1.2.3 and :1.2) - # Branch and :latest tags are convenience aliases; the : tag + # Branch and :latest tags are convenience aliases; the :sha- tag # is the canonical immutable pointer. - name: Extract image metadata id: meta diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7b5c4676b..16e542526 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -29,6 +29,8 @@ repos: - id: check-added-large-files - id: check-toml - id: check-yaml + # Kubernetes manifests hold several documents separated by ---. + args: [--allow-multiple-documents] - id: check-merge-conflict - id: no-commit-to-branch args: ["--branch=main"] diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index efd88e2ae..c3e57a8a8 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -36,32 +36,32 @@ graph TB Chainlit[Chainlit UI
Plan-Approve] API[Python API] end - + subgraph Agent Layer A1[A1 Agent
General Purpose] AD1[AD1 Agent
AD Specialized] ReAct[ReAct Engine] end - + subgraph Intelligence Layer LLM[LLM Providers
Claude/GPT/Gemini/Bedrock] Retriever[Tool Retriever] KnowHow[Know-How Library] end - + subgraph Tool Layer ToolRegistry[Tool Registry] Tools[180+ Domain Tools] Database[Database Queries] MCP[MCP Servers] end - + subgraph Data Layer - DataLake[Data Lake
77+ Files ~11GB] + DataLake[Data Lake
76 Files ~15GB] ADData[AD-Specific Data
30+ Datasets] External[External APIs
20+ Databases] end - + CLI --> A1 Gradio --> A1 Chainlit --> AD1 @@ -236,13 +236,13 @@ Each tool has a corresponding schema in `biomni/tool/tool_description/`: ### Overview -The data lake contains **77 curated datasets** (~11GB), catalogued in `biomni/env_desc.py` regardless of what's on disk. +The data lake contains **76 curated datasets** (about 15 GB), catalogued in `biomni/env_desc.py` regardless of what's on disk. Individual files are fetched lazily from S3, the first time a query actually selects them (`A1._ensure_data_lake_files`) - not downloaded in bulk on agent construction. Layout once files are present: ``` ./data/ -├── data_lake/ # ~11GB of curated datasets +├── data_lake/ # ~15GB of curated datasets │ ├── *.parquet # Tabular data files │ ├── *.pkl # Serialized Python objects │ ├── *.csv # CSV files @@ -273,12 +273,12 @@ Layout once files are present: | Dataset | Description | Size | |---------|-------------|------| -| `BindingDB_All_202409.tsv` | Drug-target binding affinities | Large | -| `DepMap_CRISPRGeneEffect.csv` | Genome-wide CRISPR effects | ~2GB | -| `gtex_tissue_gene_tpm.parquet` | GTEx expression across tissues | ~500MB | -| `gwas_catalog.pkl` | GWAS association results | ~100MB | -| `kg.csv` | Precision medicine knowledge graph (17,080 diseases, 4M+ relationships) | ~200MB | -| `DisGeNET.parquet` | Gene-disease associations | ~100MB | +| `BindingDB_All_202409.tsv` | Drug-target binding affinities | 6.3 GB | +| `DepMap_CRISPRGeneEffect.csv` | Genome-wide CRISPR effects | 430 MB | +| `gtex_tissue_gene_tpm.parquet` | GTEx expression across tissues | 11 MB | +| `gwas_catalog.pkl` | GWAS association results | 180 MB | +| `kg.csv` | Precision medicine knowledge graph (17,080 diseases, 4M+ relationships) | 980 MB | +| `DisGeNET.parquet` | Gene-disease associations | 3 MB | --- @@ -556,7 +556,7 @@ graph LR ACA --> AAnth subgraph Data - DataLake[Azure Files
shared /app/data
Biomni data lake ~11GB
read-only] + DataLake[Azure Files
shared /app/data
Biomni data lake ~15GB
read-only] ADDI[ADDI Workbench Mount
controlled-access datasets
read-only] end ACA --> DataLake @@ -575,15 +575,15 @@ graph LR | Concern | Azure Service | Maps to in Biomni-AD | |---------|---------------|----------------------| | **Edge / TLS / WAF** | Azure Front Door + WAF policy | Public ingress, OWASP rule set, DDoS Standard | -| **Identity** | Authentication gateway (GRIP) | **[Implemented]** The gateway validates access and forwards the caller as HTTP headers; `biomni/identity.py` reads them and keys per-user preferences and run records off the asserted subject. GRIP sends two composite headers, each a comma-separated `key=value` list - `Ai-App-User-Context` (`sub`, `email`, `given_name`, `family_name`) and `Ai-App-Workspace-Context` (`uuid`) - and `sub` is the stable subject everything is keyed by. Single-value headers from other gateways (oauth2-proxy, Envoy/ext_authz) fill in anything the context headers did not carry, so a non-GRIP deployment keeps working. Header names are configurable via `BIOMNI_AUTH_*_HEADER`, and the headers are ignored unless `BIOMNI_TRUST_AUTH_HEADERS` is set - it fails closed so an unprotected deployment cannot be impersonated. The app never authenticates anyone itself, so the gateway **must** strip client-supplied copies of these headers. The same identity is registered as Chainlit's `header_auth_callback`, so preferences, run records and chat threads all key off one string. With no gateway in front, identity falls back to a per-session anonymous key and nothing persists across sessions - unless `BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE` declares the deployment single-user, which makes every session one shared local user. | +| **Identity** | Authentication gateway (GRIP) | **[Implemented]** The gateway validates access and forwards the caller as HTTP headers; `biomni/identity.py` reads them and keys per-user preferences and run records off the asserted subject. GRIP sends two composite headers, each a comma-separated `key=value` list - `Ai-App-User-Context` (`sub`, `email`, `given_name`, `family_name`) and `Ai-App-Workspace-Context` (`uuid`) - and `sub` is the stable subject everything is keyed by. Single-value headers from other gateways (oauth2-proxy, Envoy/ext_authz) are read for a request that carries neither context header, so a non-GRIP deployment keeps working; next to GRIP's they are ignored, since a gateway strips only its own headers from what clients send and a client could otherwise fill in a field GRIP left out. A context header that valid CSV cannot produce - an unterminated quote, or `sub`, `email`, `iss` or `uuid` given twice - is ignored whole, and on the websocket the raw ASGI header list is read rather than the `HTTP_*` environ keys, which fold a client's `Ai_App_User_Context` into the gateway's header. Header names are configurable via `BIOMNI_AUTH_*_HEADER`, and the headers are ignored unless `BIOMNI_TRUST_AUTH_HEADERS` is set - it fails closed so an unprotected deployment cannot be impersonated. The app never authenticates anyone itself, so the gateway **must** strip client-supplied copies of these headers. The same identity is registered as Chainlit's `header_auth_callback`, so preferences, run records and chat threads all key off one string, and each chat session takes its identity from that login (`resolve_session_identity`) rather than re-deriving it. A login cookie naming a different user than the gateway's headers - someone else signing in on the same browser - or one that cannot say whom it was issued for, such as a cookie from an earlier release, is refused on every HTTP endpoint (`get_current_user` is overridden), and the page signs in afresh as the gateway's user; a websocket that gets through regardless is refused a session, a reopened conversation's transcript and its live run. The context headers are only ever parsed: a single-value override pointed at one of them has no effect and is logged as such at startup, and names sent as UTF-8 are repaired from the latin-1 decoding HTTP servers apply. With no gateway in front, identity falls back to a per-session anonymous key and nothing persists across sessions - unless `BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE` declares the deployment single-user, which makes every session one shared local user. | | **App runtime** | **Azure Container Apps** (preferred) or AKS | Runs the existing `Dockerfile` (micromamba + `chainlit run`) unmodified; per-revision rollouts | -| **Image registry** | Azure Container Registry (mirror of GHCR) | **[Implemented]** `.github/workflows/docker.yml` builds the `Dockerfile` on every PR and publishes to GHCR (`ghcr.io/kaimen-inc/biomni-ad`) on push to `main` / `feat/adworkbench` / tags. Tags: `:`, `:`, plus semver aliases for git tags. For ACR-based deployments, mirror from GHCR rather than rebuilding. | -| **Secrets** | Azure Key Vault + Container Apps secret refs | `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `AZURE_*` keys, DB password - never baked into the image | -| **Shared data lake** | Azure Files (Premium, SMB), mounted **read-only** at `/app/data` | The 77-file ~11GB Biomni data lake; downloaded once into the file share, then mounted by every replica | +| **Image registry** | Azure Container Registry (mirror of GHCR) | **[Implemented]** `.github/workflows/docker.yml` builds the `Dockerfile` on every PR and publishes to GHCR (`ghcr.io/kaimen-inc/biomni-ad`) on push to `main` / `feat/adworkbench` / tags. Tags: `:sha-`, `:`, plus semver aliases for git tags. For ACR-based deployments, mirror from GHCR rather than rebuilding. | +| **Secrets** | Azure Key Vault + Container Apps secret refs | `BIOMNI_LLM_PROXY_API_KEY` on GRIP, otherwise `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `AZURE_*` keys; DB password - never baked into the image | +| **Shared data lake** | Azure Files (Premium, SMB) at `BIOMNI_DATA_LAKE_PATH` | The Biomni data lake plus the `biomniAD` subtree, about 16.5 GB in all. **[Implemented]** each dataset is downloaded the first time a question needs it, so the location must be writable; `deploy/k8s/biomni-ad.yaml` puts it on the state volume. **[Target]** a share filled once and mounted by every replica. | | **Per-user scratch** | Azure Files (per-user share) at `/app/user-data` | User uploads + downloaded AD catalog files (`biomniAD//`) | -| **Run artifacts** | Persistent volume, optionally tiered to Blob | **[Implemented]** The output root resolves to the user's setting, then `BIOMNI_OUTPUT_ROOT`, then `/biomni-outputs/`, then container-local `./runs/` - which is flagged in the UI as ephemeral because a restart destroys it. `deploy/k8s/biomni-ad.yaml` mounts a PVC at `/data` for this. **[Target]** lifecycle rules to cool/archive tiers. | +| **Run artifacts** | Persistent volume, optionally tiered to Blob | **[Implemented]** The output root resolves to the user's setting, then `BIOMNI_OUTPUT_ROOT`, then `/biomni-outputs/`, then container-local `./runs/` - which is flagged in the UI as ephemeral because a restart destroys it. `deploy/k8s/biomni-ad.yaml` mounts a PVC at `/app/runs` for this, as the GRIP deployment does. **[Target]** lifecycle rules to cool/archive tiers. | | **Chat history** | PostgreSQL Flexible Server | **[Implemented]** Chainlit's SQLAlchemy data layer, wired in `chainlit_ui/persistence.py`: conversations are listed in the left sidebar, reopened, and continued (`on_chat_resume`). Defaults to SQLite on the state volume - which assumes a single replica, like the file-backed preferences - and switches to Postgres by setting `BIOMNI_THREADS_DB_URL`. Small attachments are archived inline with the conversation so a reopened thread renders as it did live; the SQLite schema is created on boot (and additively upgraded, since a column the data layer writes but the table lacks fails the insert *silently*), any other database is an operator migration. A run is not bound to the browser: closing the tab leaves it executing and writing into its conversation, and reopening that conversation redirects the run's output into the new connection (`chainlit_ui/live_runs.py`) so it streams on live and Stop cancels the run rather than the tab. It is still bound to the process - a restart ends it, recorded as `interrupted`. | -| **LLM** | Azure OpenAI **and/or** Azure AI Foundry Claude | Set `LLM_SOURCE=AzureOpenAI` / `AzureAnthropic` + endpoint/deployment env vars; no code change | +| **LLM** | Azure AI Foundry through the platform LLM proxy (GRIP); Azure OpenAI or Foundry Claude directly elsewhere | **[Implemented]** With `BIOMNI_LLM_PROXY_URL` set, `biomni/llm.py` routes every model call through the platform's proxy, ahead of any provider setting, in the OpenAI (`/v1/chat/completions`) or Anthropic (`/v1/messages`) schema per `BIOMNI_LLM_PROXY_SCHEMA`. Each request authenticates with the proxy's bearer token and names the session's user and workspace (`X-User-Id` = `sub`, `X-Workspace-Id` = `uuid`), read per request from the identity bound to the running session (`biomni/identity.py::bound_identity`, carried into worker threads by `capture_context`), so one model instance can never speak for another user. A call without a gateway identity is refused before it is sent (`biomni/llm_proxy.py`), and the proxy's refusals reach the user in plain words (`chainlit_ui/llm_failures.py`). The Claude web-search tool, which would call Anthropic directly, is withheld while the proxy is configured (`biomni/tool/availability.py`). Elsewhere, set `LLM_SOURCE=AzureOpenAI` / `Anthropic` + endpoint/deployment env vars; no code change. | | **Observability** | Application Insights + Log Analytics | **[Implemented]** Chainlit + stdlib `logging` write to stdout; Container Apps ships container logs to Log Analytics out of the box. **[Target]** OpenTelemetry instrumentation around LangGraph node transitions and tool calls - not wired up today; recommended before production rollout so per-turn latency and tool error rates are queryable. | | **CI/CD** | GitHub Actions → GHCR → Container Apps revision | **[Implemented]** GHCR publish on push (see Image registry row). **[Target]** Container Apps revision rollout from GHCR (`az containerapp update --image ghcr.io/...:`) with blue/green via traffic splits - operator-side wiring. | | **Container liveness** | Container Apps HTTP / TCP probe | **[Implemented]** `HEALTHCHECK` baked into the Dockerfile (TCP probe on `:8000` via `python -c`). `docker compose` inherits this directly; no duplicate block in `docker-compose.yml`. **[Target]** Container Apps probes are configured separately via `ingress.targetPort` + `probes.{startupProbe,livenessProbe}` in the app spec - Container Apps does **not** read Dockerfile `HEALTHCHECK`/`EXPOSE` directives. Point both probes at `:8000` to match. | @@ -626,8 +626,8 @@ The Alzheimer's Disease Data Initiative workbench provides a hosted Azure-based - **Image source.** Pull from public ACR (or GitHub Container Registry mirror) - no rebuild inside ADDI. - **Data lake.** Use the ADDI-provided read-only mount for the shared Biomni data lake instead of provisioning Azure Files separately. - **Controlled-access AD datasets.** Reference catalog URIs only; the agent reads bytes from the ADDI mount path (e.g., `/workbench/niagads//…`) when present, otherwise falls back to the public download path. The `BIOMNI_DATA_PATH` env var pins this. -- **Identity.** ADDI's existing OIDC flow gates Chainlit; no separate Entra ID tenant. -- **LLM.** Workbench-provided Azure OpenAI / Foundry Claude deployment by default; user-provided keys via the Chainlit settings panel for those who prefer their own quota. +- **Identity.** The platform's authentication gateway signs users in and forwards them in the `Ai-App-*` headers (see the Identity row above); no separate Entra ID tenant. +- **LLM.** Workbench-provided Azure AI Foundry models, reached through the platform's LLM proxy, which meters usage per user and per workspace (see the LLM row above and [docs/grip_deployment.md](docs/grip_deployment.md)). - **Egress.** Constrained to ADDI's allowed endpoints (LLM, NIAGADS, AD Workbench dataset APIs). The agent's database query tools that hit external public APIs (UniProt, Ensembl, etc.) are routed via the ADDI egress proxy. ### 7. Scaling & Cost Model @@ -686,7 +686,7 @@ The app is built to be operable as many pods behind a managed Kubernetes/Contain **Perceived liveness.** While a code step blocks, the Chainlit UI ticks an elapsed-time line on the running step (so it doesn't look frozen), and the backend `run_heartbeat` provides the same signal in the log pipeline. -**Health probes.** `GET /healthz` (liveness: process up, dependency-free) and `GET /readyz` (readiness: data dir mounted + an LLM credential present → `503` otherwise) are registered ahead of Chainlit's SPA catch-all. See [`deploy/k8s/biomni-ad.yaml`](deploy/k8s/biomni-ad.yaml) for probe wiring. +**Health probes.** `GET /healthz` (liveness: process up, dependency-free) and `GET /readyz` (readiness: the workspace mounted when one is configured, and a model reachable - an LLM credential, or the platform proxy's token with gateway trust on - → `503` otherwise) are registered ahead of Chainlit's SPA catch-all. See [`deploy/k8s/biomni-ad.yaml`](deploy/k8s/biomni-ad.yaml) for probe wiring. **Platform status endpoint.** `GET /status` answers the platform monitoring framework's question - is this application doing anything? - in its agreed schema: diff --git a/Dockerfile b/Dockerfile index eca07864d..2905d0488 100644 --- a/Dockerfile +++ b/Dockerfile @@ -47,11 +47,20 @@ COPY ${BIOMNI_ENV_FILE} /tmp/biomni_env.yml # BuildKit cache mounts: pkg downloads and pip's HTTP cache persist across # builds for speed, but never enter any image layer. ``sharing=locked`` so # parallel builds (e.g. arm64 + amd64) don't corrupt the cache. +# +# Chainlit is pinned (here and in adworkbench_env.yml) to the version the app +# is tested against: chainlit_app.py hooks its login dependency and styles its +# settings dialog, and an unpinned upgrade would change either without notice. +# Its data layer needs SQLAlchemy (in the env) and a database driver: aiosqlite +# for the default SQLite store, asyncpg for BIOMNI_THREADS_DB_URL on Postgres. +# Not chainlit[custom-data]: that also pulls the S3, GCS and Azure storage SDKs +# for attachments, which this app keeps inline (chainlit_ui/persistence.py), and +# the GCS one upgrades requests past what the env's arxiv client accepts. RUN --mount=type=cache,target=/opt/conda/pkgs,sharing=locked \ --mount=type=cache,target=/root/.cache/pip,sharing=locked \ micromamba create -y -n biomni_e1 -f /tmp/biomni_env.yml && \ micromamba run -n biomni_e1 pip install --upgrade pip && \ - micromamba run -n biomni_e1 pip install "chainlit[custom-data]>=2.8" "aiosqlite>=0.20" --index-url https://pypi.org/simple/ + micromamba run -n biomni_e1 pip install "chainlit==2.11.1" "aiosqlite>=0.20,<1" "asyncpg>=0.30,<1" --index-url https://pypi.org/simple/ # Editable install of the biomni package. Stage 2 copies the same source # tree to the same /app path so the .pth pointer resolves at runtime. diff --git a/README.md b/README.md index 5e16ee398..ed8df4b18 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ The underlying **Biomni** platform is a general-purpose biomedical AI agent that | [docs/configuration.md](docs/configuration.md) | Configuration management guide | | [docs/known_conflicts.md](docs/known_conflicts.md) | Known package conflicts and workarounds | | [docs/docker_vm_deployment.md](docs/docker_vm_deployment.md) | Docker and VM deployment guide | +| [docs/grip_deployment.md](docs/grip_deployment.md) | Deploying on the AD Workbench (GRIP): identity headers, the LLM proxy, storage, `/status`, resources | | [docs/mcp_integration.md](docs/mcp_integration.md) | Model Context Protocol (MCP) server integration | | [docs/building_documentation.md](docs/building_documentation.md) | Building Sphinx API documentation | @@ -145,10 +146,10 @@ AWS_BEARER_TOKEN_BEDROCK=your_bedrock_api_key_here AWS_REGION=us-east-1 # Optional: Custom model serving configuration -# CUSTOM_MODEL_BASE_URL=http://localhost:8000/v1 -# CUSTOM_MODEL_API_KEY=your_custom_api_key_here +# BIOMNI_CUSTOM_BASE_URL=http://localhost:8000/v1 +# BIOMNI_CUSTOM_API_KEY=your_custom_api_key_here -# Optional: Biomni data path (defaults to ./data) +# Optional: Biomni data path (defaults to ~/.biomni/data) # BIOMNI_DATA_PATH=/path/to/your/data # Optional: Timeout settings (defaults to 600 seconds) @@ -257,7 +258,7 @@ The recommended way to run Biomni-AD is the **Chainlit UI** with its **plan-then **Setup (one-time):** ```bash conda activate biomni_e1 -pip install "chainlit>=1.0" +pip install -e ".[chainlit]" # the Chainlit version the app is tested against, and its database drivers ``` **Single instance:** @@ -328,10 +329,11 @@ docker compose up -d Then open `http://localhost:8000` (or your VM public IP). For full VM deployment instructions (firewall/security group, operations, and hardening), see [docs/docker_vm_deployment.md](docs/docker_vm_deployment.md). +For the AD Workbench (GRIP) platform, see [docs/grip_deployment.md](docs/grip_deployment.md). #### Controlling Datalake Loading -Biomni does not download the full datalake (~11GB) when you create an agent. +Biomni does not download the full datalake (about 15 GB) when you create an agent. Individual files are fetched lazily, the first time a query actually needs them - a session that never touches DepMap never pays to fetch it. Set `BIOMNI_DATA_LAKE_PATH` if the datalake should live somewhere other than the repo-local `data/` folder, e.g. a dedicated volume or a different disk on the server. diff --git a/biomni/agent/a1.py b/biomni/agent/a1.py index 843726512..ab891d5d9 100644 --- a/biomni/agent/a1.py +++ b/biomni/agent/a1.py @@ -40,9 +40,10 @@ class _RunDirReady(Exception): from biomni.fs_scan import scan_directory from biomni.know_how import KnowHowLoader from biomni.llm import SourceType, get_llm, resolve_source -from biomni.llm_resilience import LLMUsageTracker, prepare_messages_for_cache +from biomni.llm_resilience import LLMUsageTracker, prepare_messages_for_cache, supports_prompt_caching from biomni.model.retriever import ToolRetriever from biomni.observability import emit_event +from biomni.tool.availability import CLAUDE_WEB_SEARCH from biomni.tool.support_tools import run_python_repl from biomni.tool.tool_registry import ToolRegistry from biomni.utils import ( @@ -260,7 +261,7 @@ def __init__( # get_llm() above would already have raised on a truly bad config. self._llm_source = None self._llm_prompt_caching = bool( - getattr(default_config, "enable_prompt_caching", False) and self._llm_source == "Anthropic" + getattr(default_config, "enable_prompt_caching", False) and supports_prompt_caching(self._llm_source) ) # Resolved model name for per-call telemetry (provider objects expose it # under different attributes; fall back to the configured/auto name). @@ -1588,6 +1589,35 @@ def _filter_know_how_for_commercial_mode(self): self.know_how_loader.remove_document(doc_id) print(f" ⚠️ Excluded know-how '{doc_name}' (non-commercial license)") + def _advertises_tool(self, name: str) -> bool: + """Whether the tool ``name`` is in this agent's catalogue.""" + return any(api.get("name") == name for apis in getattr(self, "module2api", {}).values() for api in apis) + + def _tool_priority_instructions(self) -> str: + """The search-first tool order, naming only functions this agent can call. + + A function the model is told to prefer but cannot call is a failed step + on nearly every question, so the Claude web search is named only where + this deployment advertises it (see biomni/tool/availability.py). + """ + claude = self._advertises_tool(CLAUDE_WEB_SEARCH) + web_search = f"`{CLAUDE_WEB_SEARCH}()` or `search_google()`" if claude else "`search_google()`" + protocol_search = f"{CLAUDE_WEB_SEARCH}()" if claude else "search_google()" + return f""" +TOOL PRIORITY - ALWAYS FOLLOW THIS ORDER: +1. **Web & literature search first**: Before writing any analysis code, use web/literature tools to gather information: + - {web_search} for general web searches + - `query_pubmed()`, `query_arxiv()`, `query_scholar()` for scientific literature + - the database query tools (`query_uniprot()`, `query_ensembl()`, `query_gwas_catalog()`, ...) for structured data +2. **Local data & built-in tools second**: Use data lake files, database tools, and domain-specific functions that are already available. +3. **Code generation last**: Only write custom analysis code (Python/R/Bash) when the above tools cannot provide the answer directly. Keep code minimal and focused. + +This priority order is especially important - retrieving information is faster and more reliable than generating it from scratch. When in doubt, search first. + +PROTOCOL GENERATION: +If the user requests an experimental protocol, use search_protocols(), {protocol_search}, list_local_protocols(), and read_local_protocol() to generate an accurate protocol. Include details such as reagents (with catalog numbers if available), equipment specifications, replicate requirements, error handling, and troubleshooting - but ONLY include information found in these resources. Do not make up specifications, catalog numbers, or equipment details. Prioritize accuracy over completeness. +""" + def _generate_system_prompt( self, tool_desc, @@ -1845,21 +1875,8 @@ def format_item_with_description(name, description): You may or may not receive feedbacks from human. If so, address the feedbacks by following the same procedure of multiple rounds of thinking, execution, and then coming up with a new solution. """ - # Add protocol generation instructions - prompt_modifier += """ -TOOL PRIORITY - ALWAYS FOLLOW THIS ORDER: -1. **Web & literature search first**: Before writing any analysis code, use web/literature tools to gather information: - - `advanced_web_search()` or `advanced_web_search_claude()` for general web searches - - `search_pubmed()`, `search_biorxiv()` for scientific literature - - `query_knowledge_base()` or database query tools for structured data -2. **Local data & built-in tools second**: Use data lake files, database tools, and domain-specific functions that are already available. -3. **Code generation last**: Only write custom analysis code (Python/R/Bash) when the above tools cannot provide the answer directly. Keep code minimal and focused. - -This priority order is especially important - retrieving information is faster and more reliable than generating it from scratch. When in doubt, search first. - -PROTOCOL GENERATION: -If the user requests an experimental protocol, use search_protocols(), advanced_web_search_claude(), list_local_protocols(), and read_local_protocol() to generate an accurate protocol. Include details such as reagents (with catalog numbers if available), equipment specifications, replicate requirements, error handling, and troubleshooting - but ONLY include information found in these resources. Do not make up specifications, catalog numbers, or equipment details. Prioritize accuracy over completeness. -""" + # Search-first tool order and protocol generation. + prompt_modifier += self._tool_priority_instructions() # Add custom resources section first (highlighted) has_custom_resources = any( diff --git a/biomni/agent/ad1.py b/biomni/agent/ad1.py index a85b381a7..cb70e53a7 100644 --- a/biomni/agent/ad1.py +++ b/biomni/agent/ad1.py @@ -12,6 +12,7 @@ from biomni.agent.a1 import A1 from biomni.agent.ad_data_downloader import download_ad_catalog_data from biomni.fs_scan import scan_directory +from biomni.tool.availability import CLAUDE_WEB_SEARCH try: from rich.console import Console @@ -662,6 +663,9 @@ def _inject_ad_context(self): ad_sourcing_content = f.read() data_root_dir = getattr(self, "data_root_dir", "not set") + web_search = "search_google()" + if self._advertises_tool(CLAUDE_WEB_SEARCH): + web_search = f"{CLAUDE_WEB_SEARCH}(), search_google()" ad_instruction = f""" @@ -669,7 +673,7 @@ def _inject_ad_context(self): 1. **Local data first**: Scan the built-in data lake ({getattr(self, "data_lake_dir", "not set")}) and user data directory ({data_root_dir}) for any locally available AD datasets. Use os.listdir() on both locations - the data lake has curated datasets; the user directory may contain additional data. 2. **BiomniAD catalogs second**: Load JSON catalogs from biomni/know_how/resource/ to find datasets with download URIs. -3. **Web & literature search third**: Use advanced_web_search(), search_pubmed(), search_biorxiv() to supplement. +3. **Web & literature search third**: Use {web_search}, query_pubmed(), query_arxiv() to supplement. 4. **Code generation last**: Write custom Python/R code only when the above cannot answer directly. Do NOT simulate or fabricate data at any step. diff --git a/biomni/health.py b/biomni/health.py index cbeff33b5..ba49629d8 100644 --- a/biomni/health.py +++ b/biomni/health.py @@ -7,9 +7,10 @@ pod restart. So ``/healthz`` returns 200 as long as the event loop can serve a request. * **readiness** (``/readyz``) - "should this pod receive traffic *right now*?". - Here we verify the agent could actually function: its data directory is - mounted and an LLM credential is configured. A misconfigured pod reports 503 - and is pulled from the Service endpoints instead of failing user requests. + Here we verify the agent could actually function: the user's workspace is + mounted, when one is configured, and a model can be reached. A misconfigured + pod reports 503 and is pulled from the Service endpoints instead of failing + user requests. Both endpoints are registered at the *front* of the router. Chainlit mounts a catch-all ``GET /{full_path:path}`` (for the SPA) as its last route; a route @@ -22,6 +23,8 @@ from typing import TYPE_CHECKING, Any from biomni import credentials +from biomni.identity import trust_auth_headers +from biomni.llm_proxy import llm_proxy_settings if TYPE_CHECKING: from collections.abc import Callable @@ -55,29 +58,63 @@ def build_info() -> dict[str, str]: return info -def _data_path() -> str: - """The data directory the app expects to be mounted (honors env overrides).""" - # Read env directly rather than importing biomni.config: readiness must - # reflect the *current* environment and stay import-light for the probe path. - return ( - os.getenv("BIOMNI_USER_DATA_PATH") - or os.getenv("BIOMNI_DATA_PATH") - or os.getenv("BIOMNI_PATH") - or os.path.join(os.path.expanduser("~"), ".biomni", "data") - ) +def _workspace_check() -> dict[str, Any]: + """Whether the user's workspace is mounted, when one is configured. + Worth holding traffic back for: without it every question about the user's + own data is answered as if there were none. With no workspace configured + there is nothing to wait for. ``BIOMNI_PATH`` is deliberately not consulted: + it is the app's own data directory, which the image does not contain and the + agent creates the first time it is built - and that takes a chat, which an + unready pod never receives. + """ + # Env read directly rather than through biomni.config: readiness must reflect + # the *current* environment and stay import-light for the probe path. + for env in ("BIOMNI_USER_DATA_PATH", "BIOMNI_DATA_PATH"): + path = os.getenv(env, "").strip() + if path: + return {"ok": os.path.isdir(path), "path": path} + return {"ok": True, "path": None} + + +def _llm_credential_check() -> dict[str, Any]: + """Whether some usable LLM credential / endpoint is configured. + + With an LLM proxy configured every model call goes through it, so its token + is the only credential that counts: a provider key left over in the + environment must not make a pod that cannot reach any model look ready. Nor + must a proxy that will refuse every call, because nothing tells it whom the + call is for: its user and workspace ids come only from a trusted gateway. + """ + try: + proxy = llm_proxy_settings() + except ValueError as exc: + return {"ok": False, "source": "BIOMNI_LLM_PROXY_URL", "detail": str(exc)} + if proxy is not None: + if not proxy.api_key: + return { + "ok": False, + "source": "BIOMNI_LLM_PROXY_URL", + "detail": "BIOMNI_LLM_PROXY_API_KEY is not set; the proxy needs its token", + } + if not trust_auth_headers(): + return { + "ok": False, + "source": "BIOMNI_LLM_PROXY_URL", + "detail": "BIOMNI_TRUST_AUTH_HEADERS is not enabled; the proxy needs the user and workspace ids " + "only the authentication gateway supplies", + } + return {"ok": True, "source": "BIOMNI_LLM_PROXY_URL"} -def _has_llm_credential() -> tuple[bool, str]: - """Whether some usable LLM credential / endpoint is configured.""" for env in _PROVIDER_KEY_ENVS: # Via the vault, not os.getenv: credentials are hidden from os.environ # while generated code runs, and a probe landing in that window must # not report the pod unready. if credentials.getenv(env): - return True, env + return {"ok": True, "source": env} if os.getenv("BIOMNI_CUSTOM_BASE_URL"): - return True, "BIOMNI_CUSTOM_BASE_URL" - return False, "" + return {"ok": True, "source": "BIOMNI_CUSTOM_BASE_URL"} + return {"ok": False, "source": None} def readiness_report() -> tuple[bool, dict[str, Any]]: @@ -86,14 +123,10 @@ def readiness_report() -> tuple[bool, dict[str, Any]]: Returns ``(ok, report)`` where ``report`` carries per-check detail suitable for the probe body and for debugging a pod stuck in ``NotReady``. """ - checks: dict[str, dict[str, Any]] = {} - - data_path = _data_path() - data_ok = os.path.isdir(data_path) - checks["data_path"] = {"ok": data_ok, "path": data_path} - - cred_ok, cred_source = _has_llm_credential() - checks["llm_credential"] = {"ok": cred_ok, "source": cred_source or None} + checks: dict[str, dict[str, Any]] = { + "data_path": _workspace_check(), + "llm_credential": _llm_credential_check(), + } ok = all(c["ok"] for c in checks.values()) report: dict[str, Any] = {"status": "ready" if ok else "not_ready", "checks": checks} diff --git a/biomni/identity.py b/biomni/identity.py index 12439e500..80b4ba9d3 100644 --- a/biomni/identity.py +++ b/biomni/identity.py @@ -24,12 +24,14 @@ at all - introducing it later re-keys every existing user (see :meth:`UserIdentity.storage_key`). - Values are read as CSV: a value containing a comma may be quoted, and a - literal quote inside a quoted value is doubled per RFC 4180. + Values are read as CSV: a field containing a comma may be quoted, either the + value alone or the whole ``key=value`` pair, and a literal quote inside a + quoted field is doubled per RFC 4180. A header that CSV cannot have produced + (an unterminated quote, or an identity key given twice) is rejected whole. * **Other gateways still work, and header names stay configurable.** Every field also has a list of single-value header names (oauth2-proxy, - Envoy/ext_authz, generic reverse proxies), consulted for whatever the context - headers did not supply and overridable per-deployment via + Envoy/ext_authz, generic reverse proxies), consulted only when neither + context header is sent, and overridable per-deployment via ``BIOMNI_AUTH_*_HEADER`` (comma-separated). A renamed header is therefore a values change in the manifest, never a code change. * **Header shape is normalised.** ``extract_headers`` accepts a plain mapping, @@ -44,19 +46,28 @@ defaults rather than by anything here. Storage keys are deliberately not raw identifiers - see :meth:`UserIdentity.storage_key`. + +One identity per chat session, decided once and then carried: +:func:`resolve_session_identity` settles who a websocket session belongs to, and +:func:`bound_identity` makes that answer visible to code acting on the user's +behalf further down - the LLM proxy needs the raw subject and workspace ids on +every model call (see :mod:`biomni.llm`). """ from __future__ import annotations +import contextlib +import contextvars import hashlib import logging import os import re +from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from collections.abc import Iterable, Mapping + from collections.abc import Iterable, Iterator logger = logging.getLogger(__name__) @@ -72,6 +83,11 @@ _CONTEXT_GIVEN_NAME_KEY = "given_name" _CONTEXT_FAMILY_NAME_KEY = "family_name" _CONTEXT_WORKSPACE_KEY = "uuid" +# The keys that decide whose data a request reaches. Given twice, the header is +# refused rather than guessed at (see parse_context_header). +_CONTEXT_IDENTITY_KEYS = frozenset( + {_CONTEXT_SUBJECT_KEY, _CONTEXT_ISSUER_KEY, _CONTEXT_EMAIL_KEY, _CONTEXT_WORKSPACE_KEY} +) # Candidate header names per field, tried in order, for gateways that send one # value per header instead. The defaults cover the conventions of the common @@ -119,15 +135,53 @@ def _header_candidates(env_name: str, defaults: tuple[str, ...]) -> tuple[str, . return overrides + tuple(d for d in defaults if d not in overrides) +def _context_header_names() -> frozenset[str]: + """Every header name currently treated as a composite context header.""" + return frozenset( + _header_candidates("BIOMNI_AUTH_USER_CONTEXT_HEADER", _DEFAULT_USER_CONTEXT_HEADERS) + + _header_candidates("BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER", _DEFAULT_WORKSPACE_CONTEXT_HEADERS) + ) + + +def _single_value_candidates(env_name: str, defaults: tuple[str, ...]) -> tuple[str, ...]: + """Header names to read one field from verbatim - never a composite header. + + Pointing ``BIOMNI_AUTH_USER_ID_HEADER`` at ``Ai-App-User-Context`` is the + natural reading of "the user id is in that header", and it is what the GRIP + team did. Read verbatim, that header's whole ``email=...,sub=...`` value + would become the user id whenever ``sub`` was missing, so context headers are + only ever parsed, and :func:`auth_config_warnings` says the override is moot. + """ + context = _context_header_names() + return tuple(name for name in _header_candidates(env_name, defaults) if name not in context) + + +def _repair_utf8(value: str) -> str: + """Undo a latin-1 decode of what was really UTF-8. + + Header bytes are decoded as latin-1 by Starlette and ASGI servers, which is + right only for ASCII. A gateway sending ``given_name=João`` as UTF-8 - as + any modern service does - arrives as ``João``. The value is re-decoded when + it round-trips cleanly; genuine latin-1 text is not valid UTF-8 and is left + exactly as it came. + """ + if value.isascii(): + return value + try: + return value.encode("latin-1").decode("utf-8") + except (UnicodeEncodeError, UnicodeDecodeError): + return value + + def _decode(value: Any) -> str: """Coerce a header value (``str`` or ``bytes``) to a stripped ``str``.""" if isinstance(value, bytes): try: - return value.decode("latin-1").strip() + value = value.decode("latin-1") except (UnicodeDecodeError, AttributeError): return "" if isinstance(value, str): - return value.strip() + return _repair_utf8(value).strip() return "" @@ -149,33 +203,38 @@ def extract_headers(source: Any) -> dict[str, str]: headers: dict[str, str] = {} try: - # ASGI scope (or anything else carrying a raw header pair list). - raw_pairs = source.get("headers") if isinstance(source, dict) else None - if isinstance(raw_pairs, list | tuple): - for pair in raw_pairs: - if not isinstance(pair, list | tuple) or len(pair) != 2: - continue - name = _decode(pair[0]).lower() - if name: - headers[name] = _decode(pair[1]) - if isinstance(source, dict): - for key, value in source.items(): - if not isinstance(key, str): - continue - if key.startswith("HTTP_"): - # WSGI/socket.io environ: HTTP_X_AUTH_REQUEST_EMAIL -> x-auth-request-email - headers[key[5:].replace("_", "-").lower()] = _decode(value) - elif key == "headers": - continue # already handled above - elif "-" in key or key.islower(): - # Already header-shaped (plain mapping of real header names). - headers.setdefault(key.lower(), _decode(value)) - # An ASGI scope is often nested under a socket.io environ. nested = source.get("asgi.scope") - if isinstance(nested, dict): - for name, value in extract_headers(nested).items(): - headers.setdefault(name, value) + if isinstance(nested, dict) and isinstance(nested.get("headers"), list | tuple): + # A socket.io environ: read the raw header list of the ASGI scope it + # carries, not the HTTP_* keys derived from it. Those fold a header a + # client sends as "Ai_App_User_Context" - which a gateway stripping + # "Ai-App-User-Context" by name lets through - into the same key as + # the gateway's own, and join repeated headers with commas. The raw + # list keeps them apart, exactly as the HTTP path sees them. + return extract_headers(nested) + + raw_pairs = source.get("headers") + if isinstance(raw_pairs, list | tuple): + # An ASGI scope. Its header list is all of its headers; the other + # keys (type, path, client, ...) are not. A repeated header keeps + # its last value, as it does on the Starlette path below. + for pair in raw_pairs: + if not isinstance(pair, list | tuple) or len(pair) != 2: + continue + name = _decode(pair[0]).lower() + if name: + headers[name] = _decode(pair[1]) + else: + for key, value in source.items(): + if not isinstance(key, str) or key == "headers": + continue + if key.startswith("HTTP_"): + # WSGI/socket.io environ: HTTP_X_AUTH_REQUEST_EMAIL -> x-auth-request-email + headers[key[5:].replace("_", "-").lower()] = _decode(value) + elif "-" in key or key.islower(): + # Already header-shaped (plain mapping of real header names). + headers.setdefault(key.lower(), _decode(value)) else: items = getattr(source, "items", None) if callable(items): @@ -196,7 +255,7 @@ def _first_present(headers: Mapping[str, str], names: Iterable[str]) -> str | No return None -def _split_context_pairs(value: str) -> list[str]: +def _split_context_pairs(value: str) -> list[str] | None: """Split a context header on commas, ignoring commas inside double quotes. A family name legitimately contains a comma ("Smith, Jr."), and a gateway @@ -210,6 +269,10 @@ def _split_context_pairs(value: str) -> list[str]: collapsed in :func:`parse_context_header` once the surrounding quotes come off. If the gateway settles on backslash escaping instead, this is the one function that changes. + + ``None`` when a quote is left open. Valid CSV never does that, and reading + on regardless would let an unescaped value swallow the fields after it: a + family name of ``M,sub=victim,x="`` would hide the gateway's own ``sub``. """ parts: list[str] = [] current: list[str] = [] @@ -230,6 +293,8 @@ def _split_context_pairs(value: str) -> list[str]: else: current.append(char) index += 1 + if in_quotes: + return None parts.append("".join(current)) return parts @@ -237,22 +302,36 @@ def _split_context_pairs(value: str) -> list[str]: def parse_context_header(value: str | None) -> dict[str, str]: """Parse ``key=value,key=value`` into a lowercase-keyed mapping. - Tolerant by design - this is untrusted input from another team's service, - and one malformed field must not cost us the rest of the identity: + Tolerant of formatting - this is input from another team's service, and a + stray space must not cost us the identity: * whitespace around keys, values and separators is stripped (the published example has a space after one of the commas); * a value may itself contain ``=`` (base64-ish subject ids do), so only the first one separates key from value; - * optional surrounding double quotes are removed, and a doubled quote inside - them collapses to one (RFC 4180 - see :func:`_split_context_pairs`); + * double quotes around the value, or around the whole ``key=value`` field as + a standard CSV writer puts them, are removed, and a doubled quote inside + collapses to one (RFC 4180 - see :func:`_split_context_pairs`); * fragments with no ``=``, or with an empty key or value, are dropped. + + Strict about anything that looks like injection, since these values decide + whose data a request reaches: a header with an unterminated quote, or one + naming an identity key twice, yields nothing at all. Neither can come from a + gateway writing valid CSV, and both are what an unescaped value smuggling in + a ``sub`` of its own would look like. """ if not value: return {} + parts = _split_context_pairs(value) + if parts is None: + logger.warning("ignoring a context header with an unterminated quote") + return {} fields: dict[str, str] = {} - for fragment in _split_context_pairs(value): + for fragment in parts: + fragment = fragment.strip() + if len(fragment) >= 2 and fragment[0] == '"' and fragment[-1] == '"': + fragment = fragment[1:-1].replace('""', '"') key, separator, raw = fragment.partition("=") if not separator: continue @@ -260,14 +339,24 @@ def parse_context_header(value: str | None) -> dict[str, str]: item = raw.strip() if len(item) >= 2 and item[0] == '"' and item[-1] == '"': item = item[1:-1].replace('""', '"').strip() - if name and item: - fields[name] = item + if not (name and item): + continue + if name in fields and name in _CONTEXT_IDENTITY_KEYS: + logger.warning("ignoring a context header that gives %r twice", name) + return {} + fields[name] = item return fields -def _context_fields(headers: Mapping[str, str], env_name: str, defaults: tuple[str, ...]) -> dict[str, str]: - """Parsed contents of the first context header present, or an empty dict.""" - return parse_context_header(_first_present(headers, _header_candidates(env_name, defaults))) +def _context_fields(headers: Mapping[str, str], env_name: str, defaults: tuple[str, ...]) -> dict[str, str] | None: + """Parsed contents of the first context header present, or ``None`` if none is. + + Present but unusable - malformed, or missing a field - is an empty or + partial dict, not ``None``: the caller must still treat the header as the + gateway's word on those fields. + """ + value = _first_present(headers, _header_candidates(env_name, defaults)) + return None if value is None else parse_context_header(value) @dataclass(frozen=True) @@ -291,7 +380,7 @@ class UserIdentity: given_name: str | None = None family_name: str | None = None issuer: str | None = None - source: str = "anonymous" # "headers" | "local" | "anonymous" + source: str = "anonymous" # "headers" | "password" | "local" | "anonymous" @property def is_authenticated(self) -> bool: @@ -363,6 +452,39 @@ def scoped_key(self) -> str: ws = _SLUG_UNSAFE_RE.sub("-", self.workspace_id.lower()).strip("-._")[:32] return f"{base}@{ws}" if ws else base + def to_metadata(self) -> dict[str, str]: + """Every field, for carrying the identity through a login session. + + The session a login creates is the only thing the websocket connection + is guaranteed to receive, so it has to hold the whole identity - the raw + subject id above all, which the LLM proxy needs and which a storage key + cannot be turned back into. + """ + return {"source": self.source, **{name: getattr(self, name) or "" for name in _METADATA_FIELDS}} + + @classmethod + def from_metadata(cls, metadata: Any) -> UserIdentity | None: + """Rebuild an identity from :meth:`to_metadata`, or ``None``. + + ``None`` for anything that is not such a record, including one written + before ``user_id`` was part of it: reading that as an identity would key + the user by e-mail instead of by subject and show them an empty history. + """ + if not isinstance(metadata, Mapping) or "user_id" not in metadata: + return None + source = metadata.get("source") + if source not in _IDENTITY_SOURCES: + return None + fields = {name: str(metadata.get(name) or "").strip() or None for name in _METADATA_FIELDS} + return cls(source=source, **fields) + + +# The fields to_metadata carries, and the sources an identity can have. "headers" +# is the gateway; "password" the demo login; "local" the single-user mode; +# "anonymous" nobody in particular. +_METADATA_FIELDS = ("user_id", "email", "workspace_id", "given_name", "family_name", "issuer") +_IDENTITY_SOURCES = frozenset({"headers", "password", "local", "anonymous"}) + def identity_from_headers(source: Any) -> UserIdentity | None: """Build an identity from gateway headers, or ``None`` when absent. @@ -375,28 +497,38 @@ def identity_from_headers(source: Any) -> UserIdentity | None: if not headers: return None - # GRIP's composite headers first; a gateway that sends them is authoritative - # about every field it carries. + # GRIP's composite headers first. A gateway that sends them is the only word + # on every identity field - including a field it leaves out. Filling such a + # gap from a single-value header would take the value from the client: a + # gateway strips its own headers from what clients send, not every name some + # other gateway might use, so an X-Auth-Issuer, X-User-Id or X-Workspace-Id + # the browser added would arrive untouched - and re-key somebody's storage, + # hand it over, or bill another workspace for the model calls. user_context = _context_fields(headers, "BIOMNI_AUTH_USER_CONTEXT_HEADER", _DEFAULT_USER_CONTEXT_HEADERS) workspace_context = _context_fields( headers, "BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER", _DEFAULT_WORKSPACE_CONTEXT_HEADERS ) - # Single-value headers fill in whatever the context headers did not carry, - # which is how non-GRIP gateways (and a partially populated context) keep - # working unchanged. - user_id = user_context.get(_CONTEXT_SUBJECT_KEY) or _first_present( - headers, _header_candidates("BIOMNI_AUTH_USER_ID_HEADER", _DEFAULT_USER_ID_HEADERS) - ) - email = user_context.get(_CONTEXT_EMAIL_KEY) or _first_present( - headers, _header_candidates("BIOMNI_AUTH_EMAIL_HEADER", _DEFAULT_EMAIL_HEADERS) - ) - workspace_id = workspace_context.get(_CONTEXT_WORKSPACE_KEY) or _first_present( - headers, _header_candidates("BIOMNI_AUTH_WORKSPACE_HEADER", _DEFAULT_WORKSPACE_HEADERS) - ) - issuer = user_context.get(_CONTEXT_ISSUER_KEY) or _first_present( - headers, _header_candidates("BIOMNI_AUTH_ISSUER_HEADER", _DEFAULT_ISSUER_HEADERS) - ) + if user_context is not None or workspace_context is not None: + user_fields = user_context or {} + user_id = user_fields.get(_CONTEXT_SUBJECT_KEY) + email = user_fields.get(_CONTEXT_EMAIL_KEY) + issuer = user_fields.get(_CONTEXT_ISSUER_KEY) + given_name = user_fields.get(_CONTEXT_GIVEN_NAME_KEY) + family_name = user_fields.get(_CONTEXT_FAMILY_NAME_KEY) + workspace_id = (workspace_context or {}).get(_CONTEXT_WORKSPACE_KEY) + else: + # No context header at all: a gateway of another kind, which sends one + # value per header. Those are read as before. + user_id = _first_present( + headers, _single_value_candidates("BIOMNI_AUTH_USER_ID_HEADER", _DEFAULT_USER_ID_HEADERS) + ) + email = _first_present(headers, _single_value_candidates("BIOMNI_AUTH_EMAIL_HEADER", _DEFAULT_EMAIL_HEADERS)) + issuer = _first_present(headers, _single_value_candidates("BIOMNI_AUTH_ISSUER_HEADER", _DEFAULT_ISSUER_HEADERS)) + workspace_id = _first_present( + headers, _single_value_candidates("BIOMNI_AUTH_WORKSPACE_HEADER", _DEFAULT_WORKSPACE_HEADERS) + ) + given_name = family_name = None if not (user_id or email): return None @@ -405,8 +537,8 @@ def identity_from_headers(source: Any) -> UserIdentity | None: user_id=user_id, email=email, workspace_id=workspace_id, - given_name=user_context.get(_CONTEXT_GIVEN_NAME_KEY), - family_name=user_context.get(_CONTEXT_FAMILY_NAME_KEY), + given_name=given_name, + family_name=family_name, issuer=issuer, source="headers", ) @@ -507,3 +639,203 @@ def resolve_identity(header_source: Any = None, *, session_id: str | None = None if identity is not None: return identity return anonymous_identity(session_id) + + +# --------------------------------------------------------------------------- # +# One identity per session +# --------------------------------------------------------------------------- # + +# Shown wherever the app explains why nothing is being remembered. They name +# the setting, because the people reading them during an integration are the +# operators who can change it. +HEADERS_IGNORED_NOTICE = ( + "sign-in details from the authentication gateway were received but ignored, " + "because BIOMNI_TRUST_AUTH_HEADERS is not enabled" +) +NO_GATEWAY_NOTICE = ( + "no authentication gateway is configured; set BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE=true for a single-user deployment" +) +GATEWAY_SILENT_NOTICE = "the authentication gateway did not identify this session" + + +@dataclass(frozen=True) +class SessionIdentity: + """Who a chat session belongs to, and anything wrong with that answer. + + ``notice`` explains an identity that is not a signed-in user. ``conflict`` + means the browser's login belongs to someone other than whom the gateway now + says is calling: the session must not be served under either name. + """ + + identity: UserIdentity + notice: str | None = None + conflict: bool = False + + +@dataclass(frozen=True) +class Login: + """What a Chainlit login records: the user identifier and its metadata. + + The identifier is what Chainlit keys a user's conversations by; the metadata + carries the identity it was issued for (see :meth:`UserIdentity.to_metadata`), + including the raw ids the identifier cannot be turned back into. + """ + + identifier: str | None + metadata: Any = None + + def identity(self) -> UserIdentity | None: + """The identity this login stands for, or ``None`` if it cannot say. + + Both halves have to agree. A login issued by an earlier release carries + no ``user_id`` in its metadata, and one issued before the gateway was + trusted names the shared ``local`` user; neither may stand for anybody, + because conversations are keyed by the identifier, and one browser's + leftover login would open every user's session onto the same history. + """ + who = UserIdentity.from_metadata(self.metadata) + if who is None or not self.identifier or who.scoped_key() != self.identifier: + return None + return who + + +def _login_conflicts(login: Login, asserted: UserIdentity) -> bool: + """Whether a login contradicts whom the gateway says is calling right now. + + A login is a cache of what the gateway said when it was issued, and it lasts + for days. When somebody else then signs in to the platform in the same + browser, the gateway names them while the old login still names the + previous user - who would otherwise be handed the new user's requests, + conversations and model usage. A login that is not a gateway login at all, + or that cannot say whom it was issued for (see :meth:`Login.identity`), is + equally stale. + """ + who = login.identity() + return who is None or who.source != "headers" or who.scoped_key() != asserted.scoped_key() + + +def login_is_current(login: Login, header_source: Any) -> bool: + """False when a trusted gateway names someone other than the login does. + + What an HTTP endpoint checks before honouring a login cookie, so a stale one + is refused and the browser signs in again as whoever the gateway names. + """ + if not trust_auth_headers(): + return True + asserted = identity_from_headers(header_source) + if asserted is None: + return True + return not _login_conflicts(login, asserted) + + +def resolve_session_identity( + login: Login | None, header_source: Any, *, session_id: str | None = None +) -> SessionIdentity: + """Settle who a websocket session belongs to. + + ``login`` is what the login step recorded, or ``None`` for a connection + with no login at all; ``header_source`` is the connection's own request, + which carries the gateway's headers when the gateway forwards them on the + websocket too. + + Behind a trusted gateway its word on this very connection wins, provided it + agrees with the login. Failing that the login stands in for it, since it + records what the gateway asserted when it was issued - which is what keeps a + gateway that does not forward headers on websocket upgrades working. + Without a gateway, a password login is the identity; otherwise the session + is the shared local user or anonymous, with a notice saying which and why. + """ + asserted = identity_from_headers(header_source) + who = login.identity() if login is not None else None + + if trust_auth_headers(): + if asserted is not None: + if login is not None and _login_conflicts(login, asserted): + return SessionIdentity( + anonymous_identity(session_id), + notice=( + "this browser is still signed in to Biomni as a different user than the one the " + f"platform reports ({asserted.display_name}); reload the page to continue" + ), + conflict=True, + ) + return SessionIdentity(asserted) + if who is not None and who.source == "headers" and who.is_authenticated: + return SessionIdentity(who) + return SessionIdentity(anonymous_identity(session_id), notice=GATEWAY_SILENT_NOTICE) + + if who is not None and who.source == "password" and who.is_authenticated: + return SessionIdentity(who) + if single_user_mode(): + return SessionIdentity(local_identity()) + notice = HEADERS_IGNORED_NOTICE if asserted is not None else NO_GATEWAY_NOTICE + return SessionIdentity(anonymous_identity(session_id), notice=notice) + + +# The identity of the session that code is running for. Set at the chat +# handlers' boundaries and carried into worker threads with the rest of the +# context (observability.capture_context), so anything acting on the user's +# behalf - the LLM proxy above all - can ask who that is without it being +# threaded through every call. +_current_identity: contextvars.ContextVar[UserIdentity | None] = contextvars.ContextVar("biomni_identity", default=None) + + +def current_identity() -> UserIdentity | None: + """The identity bound to the running session, or ``None`` outside one.""" + return _current_identity.get() + + +def set_current_identity(identity: UserIdentity | None) -> contextvars.Token: + """Bind ``identity`` for the rest of this context. Returns a reset token.""" + return _current_identity.set(identity) + + +@contextlib.contextmanager +def bound_identity(identity: UserIdentity | None) -> Iterator[None]: + """Bind ``identity`` for a block, restoring whatever was bound before.""" + token = _current_identity.set(identity) + try: + yield + finally: + _current_identity.reset(token) + + +# --------------------------------------------------------------------------- # +# Configuration report +# --------------------------------------------------------------------------- # + +_SINGLE_VALUE_HEADER_ENVS = ( + "BIOMNI_AUTH_USER_ID_HEADER", + "BIOMNI_AUTH_EMAIL_HEADER", + "BIOMNI_AUTH_WORKSPACE_HEADER", + "BIOMNI_AUTH_ISSUER_HEADER", +) + + +def auth_config_warnings() -> list[str]: + """Identity settings that cannot do what whoever set them intended.""" + warnings: list[str] = [] + context = _context_header_names() + for env_name in _SINGLE_VALUE_HEADER_ENVS: + raw = os.getenv(env_name, "") + moot = sorted({part.strip().lower() for part in raw.split(",") if part.strip()} & context) + if moot: + warnings.append( + f"{env_name}={raw.strip()} has no effect: {', '.join(moot)} is a composite context header, " + "which Biomni parses field by field with no configuration. Remove the variable." + ) + return warnings + + +def describe_auth_config() -> dict[str, Any]: + """The identity settings in force, for one structured line at startup.""" + return { + "trust_auth_headers": trust_auth_headers(), + "user_context_headers": list( + _header_candidates("BIOMNI_AUTH_USER_CONTEXT_HEADER", _DEFAULT_USER_CONTEXT_HEADERS) + ), + "workspace_context_headers": list( + _header_candidates("BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER", _DEFAULT_WORKSPACE_CONTEXT_HEADERS) + ), + "anonymous_persistence": single_user_mode(), + } diff --git a/biomni/know_how/sgRNA_design_guide.md b/biomni/know_how/sgRNA_design_guide.md index 5bdd56b98..fe5d45f10 100644 --- a/biomni/know_how/sgRNA_design_guide.md +++ b/biomni/know_how/sgRNA_design_guide.md @@ -87,13 +87,13 @@ print(results_filtered[['Target_Gene', 'Target_Sequence', **CRITICAL**: Even if Method 1 found no results, you MUST perform this literature search before moving to Option 2. Many validated sgRNAs are published in literature but not in the Addgene database. -Use `advanced_web_search_claude` from `biomni.tool.literature` to find validated sgRNAs from literature and databases: +Search the literature and databases for validated sgRNAs with the web search tool your tool list offers: `advanced_web_search_claude` from `biomni.tool.literature` where it is listed, otherwise `search_google` from the same module: ```python -from biomni.tool.literature import advanced_web_search_claude +from biomni.tool.literature import search_google # Example usage -results = advanced_web_search_claude("sgRNA TP53 validated H. sapiens experimental") +print(search_google("sgRNA TP53 validated H. sapiens experimental", num_results=10)) ``` **Search queries to try (use multiple):** diff --git a/biomni/llm.py b/biomni/llm.py index 628b98871..f7379195b 100644 --- a/biomni/llm.py +++ b/biomni/llm.py @@ -1,11 +1,15 @@ +import functools import logging import os +import re from typing import TYPE_CHECKING, Any, Literal, Optional, cast, get_args from dotenv import load_dotenv from langchain_core.language_models.chat_models import BaseChatModel +from pydantic import SecretStr from biomni import credentials +from biomni.llm_proxy import LLMProxySettings, llm_proxy_settings, proxy_identity_headers load_dotenv(override=True) @@ -14,10 +18,140 @@ if TYPE_CHECKING: from biomni.config import BiomniConfig -SourceType = Literal["OpenAI", "AzureOpenAI", "Anthropic", "Ollama", "Gemini", "Bedrock", "Groq", "Custom"] +SourceType = Literal["OpenAI", "AzureOpenAI", "Anthropic", "Ollama", "Gemini", "Bedrock", "Groq", "Custom", "LLMProxy"] ALLOWED_SOURCES: set[str] = set(get_args(SourceType)) +# --------------------------------------------------------------------------- # +# The platform LLM proxy +# --------------------------------------------------------------------------- # +# +# Settings and the identity rule live in biomni.llm_proxy (import-light, shared +# with the readiness probe); what follows is how our chat models carry them. + +# Reasoning models reject a non-default temperature and stop sequences, on the +# proxy exactly as on OpenAI's own API (see the gpt-5 handling in get_llm). +_REASONING_MODEL_RE = re.compile(r"^(?:gpt-5|o\d)", re.IGNORECASE) + + +def _with_identity_headers(payload: dict) -> dict: + """Add the proxy's identity headers to one request's SDK arguments. + + Both SDKs take ``extra_headers`` on every create call, and both langchain + wrappers build every request - sync or async, streamed or not - through + ``_get_request_payload``, which is why that is the one place this happens. + """ + payload["extra_headers"] = {**(payload.get("extra_headers") or {}), **proxy_identity_headers()} + return payload + + +@functools.cache +def _openai_proxy_class() -> type: + from langchain_openai import ChatOpenAI + + class ChatOpenAIViaProxy(ChatOpenAI): + """ChatOpenAI that names the calling session on every request.""" + + def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[override] + payload = super()._get_request_payload(input_, stop=stop, **kwargs) + if _REASONING_MODEL_RE.match(self.model_name or ""): + payload.pop("stop", None) + payload.pop("temperature", None) + return _with_identity_headers(payload) + + return ChatOpenAIViaProxy + + +@functools.cache +def _anthropic_proxy_class() -> type: + from functools import cached_property + + import anthropic + from langchain_anthropic import ChatAnthropic + + class ChatAnthropicViaProxy(ChatAnthropic): + """ChatAnthropic that authenticates to the proxy and names the session.""" + + @cached_property + def _client_params(self) -> dict[str, Any]: + params: dict[str, Any] = { + # The proxy takes a bearer token, not an API key. Left as None + # the SDK would fill api_key from ANTHROPIC_API_KEY and forward a + # direct Anthropic credential to the proxy, so the header it + # would travel in is omitted outright. + "api_key": None, + "auth_token": self.anthropic_api_key.get_secret_value(), + "base_url": self.anthropic_api_url, + "max_retries": self.max_retries, + "default_headers": {**(self.default_headers or {}), "X-Api-Key": anthropic.Omit()}, + } + if self.default_request_timeout is None or self.default_request_timeout > 0: + params["timeout"] = self.default_request_timeout + return params + + def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[override] + return _with_identity_headers(super()._get_request_payload(input_, stop=stop, **kwargs)) + + return ChatAnthropicViaProxy + + +def _proxy_chat_model( + settings: LLMProxySettings, + model: str, + *, + temperature: float, + stop_sequences: list[str] | None, + max_retries: int, + request_timeout: float | None, +) -> BaseChatModel: + """A chat model that reaches ``model`` through the platform LLM proxy.""" + if not settings.api_key: + raise ValueError("BIOMNI_LLM_PROXY_URL is set but BIOMNI_LLM_PROXY_API_KEY is not; the proxy needs its token.") + if settings.schema == "anthropic": + return cast( + "BaseChatModel", + _anthropic_proxy_class()( + model=model, + temperature=temperature, + max_tokens=8192, + stop_sequences=stop_sequences, + max_retries=max_retries, + default_request_timeout=request_timeout, + api_key=_secret(settings.api_key), + base_url=settings.anthropic_base_url, + ), + ) + return cast( + "BaseChatModel", + _openai_proxy_class()( + model=model, + temperature=temperature, + stop_sequences=stop_sequences, + base_url=settings.openai_base_url, + api_key=_secret(settings.api_key), + max_retries=max_retries, + timeout=request_timeout, + ), + ) + + +def _secret(value: str | None) -> SecretStr | None: + """A credential as the chat-model classes type it, or ``None`` when unset.""" + return SecretStr(value) if value else None + + +def _own_key(env_name: str, provider: str) -> SecretStr: + """A provider's own key for an OpenAI-compatible endpoint, or a clear error. + + Never ``None``: given no key, ChatOpenAI falls back to ``OPENAI_API_KEY`` - + and would send the OpenAI key to Google or Groq. + """ + value = credentials.getenv(env_name) + if not value: + raise ValueError(f"{provider} models need {env_name}, which is not set") + return SecretStr(value) + + def _openai_key_kwargs() -> dict[str, Any]: """``api_key`` for the OpenAI client, sourced through the credential vault. @@ -44,17 +178,28 @@ def resolve_source( independently so callers (e.g. the A1 agent) can know the resolved source without re-implementing the heuristic. - Precedence: explicit ``source`` arg -> ``LLM_SOURCE`` / ``BIOMNI_SOURCE`` - env vars -> model-name prefix heuristics -> ``base_url`` presence -> - Azure-Anthropic env detection. Raises ``ValueError`` if the source can't - be determined. + Precedence: a configured LLM proxy (``BIOMNI_LLM_PROXY_URL``) -> explicit + ``source`` arg -> ``LLM_SOURCE`` / ``BIOMNI_SOURCE`` env vars -> model-name + prefix heuristics -> ``base_url`` presence -> Azure-Anthropic env detection. + Raises ``ValueError`` if the source can't be determined. + + The proxy outranks even an explicit ``source``: on a platform that routes + model access through one, it is the only way to a model, and a call site + that hardcodes a provider would otherwise bypass the metering the platform + requires - or simply fail for want of a key the deployment does not have. """ + if llm_proxy_settings() is not None: + return "LLMProxy" if source is not None: if source not in ALLOWED_SOURCES: raise ValueError(f"Unknown source: {source!r}. Valid: {sorted(ALLOWED_SOURCES)}") + if source == "LLMProxy": + raise ValueError("source 'LLMProxy' needs BIOMNI_LLM_PROXY_URL to be set") return source env_source = os.getenv("LLM_SOURCE") or os.getenv("BIOMNI_SOURCE") + if env_source == "LLMProxy": + raise ValueError("LLM_SOURCE=LLMProxy needs BIOMNI_LLM_PROXY_URL to be set") if env_source in ALLOWED_SOURCES: return cast("SourceType", env_source) @@ -104,11 +249,13 @@ def get_llm( """ Get a language model instance based on the specified model name and source. This function supports models from OpenAI, Azure OpenAI, Anthropic, Ollama, Gemini, Bedrock, and custom model serving. + When ``BIOMNI_LLM_PROXY_URL`` is set, every model is reached through that LLM proxy instead, whatever the + other arguments say (see :func:`llm_proxy_settings`). Args: model (str): The model name to use temperature (float): Temperature setting for generation stop_sequences (list): Sequences that will stop generation - source (str): Source provider: "OpenAI", "AzureOpenAI", "Anthropic", "Ollama", "Gemini", "Bedrock", or "Custom" + source (str): Source provider: "OpenAI", "AzureOpenAI", "Anthropic", "Ollama", "Gemini", "Bedrock", "Groq" or "Custom" If None, will attempt to auto-detect from model name base_url (str): The base URL for custom model serving (e.g., "http://localhost:8000/v1"), default is None api_key (str): The API key for the custom llm @@ -136,7 +283,12 @@ def get_llm( # Use defaults if still not specified if model is None: - model = "claude-3-5-sonnet-20241022" + # The deployment's configured model (BIOMNI_LLM), not a hardcoded id: + # a pinned model is retired sooner or later, and one the deployment + # never set up is one its proxy or key cannot reach anyway. + from biomni.config import default_config + + model = default_config.llm if temperature is None: temperature = 0.7 if api_key is None: @@ -147,6 +299,18 @@ def get_llm( # Only default-fill when no config was passed; explicit None from a # configured caller means "disable per-call timeout". request_timeout = 120.0 + + proxy = llm_proxy_settings() + if proxy is not None: + return _proxy_chat_model( + proxy, + model, + temperature=temperature, + stop_sequences=stop_sequences, + max_retries=max_retries, + request_timeout=request_timeout, + ) + # Resolve source via shared helper (auto-detection + env precedence). source = resolve_source(model, source, base_url) @@ -225,10 +389,10 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ return payload return _AzureChatOpenAINoTemp( - openai_api_key=credentials.getenv("AZURE_OPENAI_API_KEY") or credentials.getenv("OPENAI_API_KEY"), + api_key=_secret(credentials.getenv("AZURE_OPENAI_API_KEY") or credentials.getenv("OPENAI_API_KEY")), azure_endpoint=os.getenv("ENDPOINT_URL") or os.getenv("OPENAI_ENDPOINT"), azure_deployment=deployment, - openai_api_version=API_VERSION, + api_version=API_VERSION, temperature=1, # default; will be stripped from payload by subclass max_retries=max_retries, timeout=request_timeout, @@ -289,7 +453,10 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ if anthropic_base_url: anthropic_kwargs["base_url"] = anthropic_base_url - return ChatAnthropic( + # Field names, not ChatAnthropic's legacy aliases (model_name, + # max_tokens_to_sample, stop, timeout). Its model accepts both, but a + # type checker without the pydantic plugin only knows the aliases. + return ChatAnthropic( # type: ignore[call-arg] model=model, temperature=temperature, max_tokens=8192, @@ -315,7 +482,7 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ return ChatOpenAI( model=model, temperature=temperature, - api_key=credentials.getenv("GEMINI_API_KEY"), + api_key=_own_key("GEMINI_API_KEY", "Gemini"), base_url="https://generativelanguage.googleapis.com/v1beta/openai/", stop_sequences=stop_sequences, max_retries=max_retries, @@ -332,7 +499,7 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ return ChatOpenAI( model=model, temperature=temperature, - api_key=credentials.getenv("GROQ_API_KEY"), + api_key=_own_key("GROQ_API_KEY", "Groq"), base_url="https://api.groq.com/openai/v1", stop_sequences=stop_sequences, max_retries=max_retries, @@ -348,7 +515,7 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ ) # ChatOllama exposes a transport-level timeout, not max_retries; pass # what's supported and let local Ollama retries stay manual. - ollama_kwargs: dict[str, object] = {"model": model, "temperature": temperature} + ollama_kwargs: dict[str, Any] = {"model": model, "temperature": temperature} if request_timeout is not None: ollama_kwargs["timeout"] = request_timeout return ChatOllama(**ollama_kwargs) @@ -363,7 +530,7 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ # Bedrock retry config lives on the boto3 client; pass through via # ``config`` kwarg when available. Older langchain-aws versions don't # accept ``config`` directly, so build defensively. - bedrock_kwargs: dict[str, object] = { + bedrock_kwargs: dict[str, Any] = { "model": model, "temperature": temperature, "stop_sequences": stop_sequences, @@ -383,10 +550,10 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ llm = ChatOpenAI( model=model, temperature=temperature, - max_tokens=8192, + max_completion_tokens=8192, stop_sequences=stop_sequences, base_url=base_url, - api_key=api_key, + api_key=_secret(api_key), max_retries=max_retries, timeout=request_timeout, ) @@ -394,5 +561,5 @@ def _get_request_payload(self, input_, *, stop=None, **kwargs): # type: ignore[ else: raise ValueError( - f"Invalid source: {source}. Valid options are 'OpenAI', 'AzureOpenAI', 'Anthropic', 'Gemini', 'Groq', 'Bedrock', or 'Ollama'" + f"Invalid source: {source}. Valid options are 'OpenAI', 'AzureOpenAI', 'Anthropic', 'Gemini', 'Groq', 'Bedrock', 'Ollama' or 'Custom'" ) diff --git a/biomni/llm_proxy.py b/biomni/llm_proxy.py new file mode 100644 index 000000000..84c983297 --- /dev/null +++ b/biomni/llm_proxy.py @@ -0,0 +1,183 @@ +"""The platform LLM proxy: where it is, and who each request is for. + +On GRIP every model call goes through the platform's LLM proxy in front of +Azure AI Foundry ("AI Apps Onboarding", section Azure Foundry). It speaks the +provider's own API, authenticates the app with its own bearer token, and meters +usage per user and per workspace - so each request must name both, in two +headers it refuses to go without:: + + X-User-Id: + X-Workspace-Id: + +Those are the ids the authentication gateway asserted for the session a call is +made for (see :mod:`biomni.identity`). They are read when the request is built, +from the identity bound to the running session, rather than baked into a client +- which is what makes it impossible for a model instance to speak for anyone +but the session it is running in. + +This module is deliberately import-light (no LangChain), so the readiness probe +and the session-start check can use it. The chat-model classes that attach the +headers live in :mod:`biomni.llm`. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, cast, get_args + +from biomni import credentials +from biomni.identity import current_identity, trust_auth_headers + +if TYPE_CHECKING: + from biomni.identity import UserIdentity + +PROXY_USER_HEADER = "X-User-Id" +PROXY_WORKSPACE_HEADER = "X-Workspace-Id" +ProxySchema = Literal["openai", "anthropic"] + + +@dataclass(frozen=True) +class LLMProxySettings: + """Where the proxy is and how to talk to it, from the environment.""" + + root_url: str + api_key: str | None + schema: ProxySchema + + @property + def openai_base_url(self) -> str: + """Base URL for the OpenAI SDK, which appends ``/chat/completions``.""" + return f"{self.root_url}/v1" + + @property + def anthropic_base_url(self) -> str: + """Base URL for the Anthropic SDK, which appends ``/v1/messages``.""" + return self.root_url + + +def llm_proxy_settings() -> LLMProxySettings | None: + """The configured LLM proxy, or ``None`` when model calls go direct. + + ``BIOMNI_LLM_PROXY_URL`` switches it on, and accepts either the root or the + ``/v1`` base the platform documents - ``https://proxy.example`` and + ``https://proxy.example/v1`` mean the same thing, since which of the two an + operator is handed is a coin toss. ``BIOMNI_LLM_PROXY_SCHEMA`` picks the API + to speak: ``openai`` (the default, and what the platform documents) or + ``anthropic``, for a Claude deployment the proxy serves in its native schema. + + Raises ``ValueError`` for a schema it does not know, so a typo fails loudly + instead of quietly speaking the wrong API. + """ + raw = os.getenv("BIOMNI_LLM_PROXY_URL", "").strip().rstrip("/") + if not raw: + return None + root = raw[: -len("/v1")] if raw.endswith("/v1") else raw + schema = (os.getenv("BIOMNI_LLM_PROXY_SCHEMA") or "openai").strip().lower() + if schema not in get_args(ProxySchema): + raise ValueError(f"BIOMNI_LLM_PROXY_SCHEMA must be 'openai' or 'anthropic', not {schema!r}") + return LLMProxySettings( + root_url=root.rstrip("/"), + # Through the vault: os.environ loses credentials while generated code + # runs, and a chat can start during another chat's code step. + api_key=credentials.getenv("BIOMNI_LLM_PROXY_API_KEY") or None, + schema=cast("ProxySchema", schema), + ) + + +def describe_llm_proxy() -> dict[str, Any]: + """The proxy settings in force, for one structured line at startup.""" + try: + settings = llm_proxy_settings() + except ValueError as exc: + return {"enabled": True, "error": str(exc)} + if settings is None: + return {"enabled": False} + return { + "enabled": True, + "url": settings.root_url, + "schema": settings.schema, + "token_configured": bool(settings.api_key), + } + + +class LLMProxyIdentityError(RuntimeError): + """A model call through the LLM proxy has no user or workspace to name. + + Raised before anything is sent: the proxy would refuse the request anyway, + and its refusal would not say which of our sessions lacked what. + + Code calling a model from Python outside the app binds the identity itself, + with :func:`biomni.identity.bound_identity`. The message does not say so, + because generated code reads it too, and the fix it must reach for is to + stay inside the session it runs for, not to make up an identity. + """ + + +_NEEDS = "the LLM proxy needs the user and workspace every model call is made for" + + +def _header_safe(value: str) -> bool: + return value.isascii() and value.isprintable() + + +def proxy_identity_problem(identity: UserIdentity | None) -> str | None: + """Why ``identity`` cannot be named to the proxy, or ``None`` if it can. + + Only a gateway identity will do. The proxy meters and limits usage per + user, and every other kind of session - a demo password login, the shared + local user, an anonymous visitor - has an id the platform has never heard + of, so sending it would charge the usage to nobody, or to a stranger. + """ + if identity is None: + return ( + f"{_NEEDS}, and this call is not being made for any chat session: it runs in a thread started " + "without one. Make model calls from the code step itself or from a concurrent.futures thread " + "pool, which carry the session with them." + ) + if identity.source != "headers": + if trust_auth_headers(): + return ( + f"{_NEEDS}, and the platform's authentication gateway did not identify this session. " + "Reload the page to sign in again." + ) + return ( + f"{_NEEDS}, and this session was not signed in through the platform's authentication gateway " + f"(identity source: {identity.source}). Set BIOMNI_TRUST_AUTH_HEADERS=true on a deployment " + "behind the gateway." + ) + missing = [] + if not identity.user_id: + missing.append("the user id (sub in Ai-App-User-Context)") + if not identity.workspace_id: + missing.append("the workspace id (uuid in Ai-App-Workspace-Context)") + if missing: + return f"{_NEEDS}, and the authentication gateway did not send {' or '.join(missing)}." + if not (_header_safe(identity.user_id or "") and _header_safe(identity.workspace_id or "")): + return f"{_NEEDS}, and the ids the authentication gateway sent cannot be sent as HTTP headers." + return None + + +def proxy_identity_headers() -> dict[str, str]: + """``X-User-Id`` and ``X-Workspace-Id`` for the session this call is for. + + Raises :class:`LLMProxyIdentityError` when there is no such session, or it + is not one the proxy can bill (see :func:`proxy_identity_problem`). + """ + identity = current_identity() + problem = proxy_identity_problem(identity) + if problem is not None or identity is None: + raise LLMProxyIdentityError(problem) + return {PROXY_USER_HEADER: identity.user_id or "", PROXY_WORKSPACE_HEADER: identity.workspace_id or ""} + + +__all__ = [ + "PROXY_USER_HEADER", + "PROXY_WORKSPACE_HEADER", + "LLMProxyIdentityError", + "LLMProxySettings", + "describe_llm_proxy", + "llm_proxy_settings", + "proxy_identity_headers", + "proxy_identity_problem", +] diff --git a/biomni/llm_resilience.py b/biomni/llm_resilience.py index 44dd83bde..1a494bc3d 100644 --- a/biomni/llm_resilience.py +++ b/biomni/llm_resilience.py @@ -33,7 +33,22 @@ def supports_prompt_caching(source: str | None) -> bool: - """Return True if ``source`` is a provider known to honor cache_control.""" + """Return True if ``source`` is a provider known to honor cache_control. + + The platform LLM proxy speaks whichever API it is configured for, and with + the Anthropic schema that is the Messages API, cache markers included. + Caching matters more there, not less: the proxy meters every input token + against per-user and per-workspace limits, and the system prompt is resent + on every turn. + """ + if source == "LLMProxy": + from biomni.llm_proxy import llm_proxy_settings + + try: + proxy = llm_proxy_settings() + except ValueError: + return False + return proxy is not None and proxy.schema == "anthropic" return source in _CACHE_SUPPORTED_SOURCES diff --git a/biomni/model/retriever.py b/biomni/model/retriever.py index 7476f048e..58d1055ee 100644 --- a/biomni/model/retriever.py +++ b/biomni/model/retriever.py @@ -95,11 +95,13 @@ def prompt_based_retrieval(self, query: str, resources: dict, llm=None) -> dict: if llm is None: # Route the fallback through get_llm so it inherits LLM resilience # config (max_retries / request_timeout -> provider-SDK 429/5xx - # backoff). A fresh BiomniConfig() honors BIOMNI_LLM_* env overrides. + # backoff). A fresh BiomniConfig() honors BIOMNI_LLM_* env overrides, + # the model among them - not a pinned provider the deployment may + # have no key for. from biomni.config import BiomniConfig from biomni.llm import get_llm - llm = get_llm("gpt-4o", source="OpenAI", config=BiomniConfig()) + llm = get_llm(config=BiomniConfig()) # Invoke the LLM if hasattr(llm, "invoke"): diff --git a/biomni/observability.py b/biomni/observability.py index 48343642d..461748fa6 100644 --- a/biomni/observability.py +++ b/biomni/observability.py @@ -124,6 +124,40 @@ def capture_context() -> contextvars.Context: return contextvars.copy_context() +_pool_patch_lock = threading.Lock() + + +def carry_context_into_thread_pools() -> None: + """Run every ``ThreadPoolExecutor`` task in the context that submitted it. + + For the code the model writes, which runs in this process and routinely fans + work out to a thread pool. Each task would otherwise start with an empty + context and so belong to no chat session: the LLM proxy refuses its model + calls for want of a user to bill, and what it prints misses the step that + ran it. The model writes ``from concurrent.futures import ThreadPoolExecutor``, + not ours, so the stdlib class itself is patched - the same wrapping LangChain + gives its own pools (``ContextThreadPoolExecutor``). + + Per task, not per worker thread: a pool's threads outlive the task that + started them and go on to serve whichever session submits next, so a thread + that kept its creator's context would run the next session's work as the + previous user. Installed once, idempotently, and never removed. + """ + from concurrent.futures import ThreadPoolExecutor + + with _pool_patch_lock: + if getattr(ThreadPoolExecutor.submit, "_carries_context", False): + return + submit = ThreadPoolExecutor.submit + + def submit_in_context(self, fn, /, *args, **kwargs): + return submit(self, contextvars.copy_context().run, fn, *args, **kwargs) + + submit_in_context.__doc__ = submit.__doc__ + submit_in_context._carries_context = True # type: ignore[attr-defined] + ThreadPoolExecutor.submit = submit_in_context # type: ignore[method-assign] + + # --------------------------------------------------------------------------- # Redaction # --------------------------------------------------------------------------- @@ -135,6 +169,11 @@ def capture_context() -> contextvars.Context: # triggering collateral redaction of unrelated log text. _SECRET_ENV_HINTS = ("KEY", "TOKEN", "SECRET", "PASSWORD", "PASSWD", "CREDENTIAL", "AUTH") +# Settings that trip a hint but hold no secret: the names of the gateway's +# identity headers. Scrubbing them turned every mention of Ai-App-User-Context +# into "[REDACTED]" - the warnings explaining a misconfigured one included. +_NOT_SECRET_ENV = re.compile(r"^BIOMNI_(?:AUTH_\w+_HEADER|TRUST_AUTH_HEADERS)$") + REDACTED = "[REDACTED]" # High-precision provider-key shapes. Ordered longest/most-specific first. @@ -182,7 +221,7 @@ def from_environ(cls, environ: Mapping[str, str] | None = None, *, redact_emails if not value: continue upper = name.upper() - if any(hint in upper for hint in _SECRET_ENV_HINTS): + if any(hint in upper for hint in _SECRET_ENV_HINTS) and not _NOT_SECRET_ENV.match(upper): secrets.append(value) if redact_emails is None: redact_emails = _env_flag(env.get("BIOMNI_LOG_REDACT_EMAILS"), default=False) @@ -520,7 +559,7 @@ def diff_usage_summary(before: Mapping[str, Any], after: Mapping[str, Any]) -> d the delta (a straight subtraction of ratios would be meaningless). """ counters = ("calls", "input_tokens", "output_tokens", "cache_read_tokens", "cache_creation_tokens", "total_tokens") - delta = {key: int(after.get(key, 0)) - int(before.get(key, 0)) for key in counters} + delta: dict[str, Any] = {key: int(after.get(key, 0)) - int(before.get(key, 0)) for key in counters} cached = delta["cache_read_tokens"] fresh = delta["input_tokens"] delta["cache_hit_ratio"] = round(cached / (cached + fresh), 4) if (cached + fresh) > 0 else 0.0 @@ -586,12 +625,11 @@ def __enter__(self) -> RunHeartbeat: self._thread.start() return self - def __exit__(self, *_exc: object) -> bool: + def __exit__(self, *_exc: object) -> None: # None: never suppresses an exception self._stop.set() thread = self._thread if thread is not None: thread.join(timeout=2.0) - return False # never suppress exceptions def _loop(self) -> None: # Event.wait returns True once stopped; the loop ends without a final emit. diff --git a/biomni/tool/availability.py b/biomni/tool/availability.py new file mode 100644 index 000000000..0916a77a6 --- /dev/null +++ b/biomni/tool/availability.py @@ -0,0 +1,55 @@ +"""Tools that some deployments cannot run, and why. + +Most tools need nothing beyond a Python package, and a deployment missing one +simply fails the import. A few need a way out that not every deployment allows, +and advertising those to the agent is worse than useless: the system prompt +tells it to search first, so it reaches for the tool on nearly every question, +and every call is a failed step the user watches. + +Import-light (no tool implementations, no LangChain), because the tool +catalogue is filtered through it on every agent build. +""" + +from __future__ import annotations + +from biomni import credentials +from biomni.llm_proxy import llm_proxy_settings + +CLAUDE_WEB_SEARCH = "advanced_web_search_claude" + + +def claude_web_search_problem() -> str | None: + """Why ``advanced_web_search_claude`` cannot run here, or ``None`` if it can. + + It calls Anthropic's API directly, with Anthropic's server-side web search + tool. Behind the platform's LLM proxy that is off limits: every model call + there must go through the proxy, which meters usage per user and per + workspace, and an ``ANTHROPIC_API_KEY`` left in the environment would let + this one tool quietly go around it. + """ + try: + proxied = llm_proxy_settings() is not None + except ValueError: + proxied = True + if proxied: + return ( + f"{CLAUDE_WEB_SEARCH} is not available in this deployment: model calls go through the platform's " + "LLM proxy, and this tool would call Anthropic directly. Use search_google() instead." + ) + + from biomni.config import default_config + + if "claude" not in default_config.llm.lower(): + return f"{CLAUDE_WEB_SEARCH} needs a Claude model, and this deployment uses {default_config.llm}." + if not credentials.getenv("ANTHROPIC_API_KEY"): + return f"{CLAUDE_WEB_SEARCH} needs ANTHROPIC_API_KEY, which is not set." + return None + + +def unavailable_tools() -> dict[str, str]: + """Tool name -> why this deployment cannot run it.""" + problems = {CLAUDE_WEB_SEARCH: claude_web_search_problem()} + return {name: problem for name, problem in problems.items() if problem} + + +__all__ = ["CLAUDE_WEB_SEARCH", "claude_web_search_problem", "unavailable_tools"] diff --git a/biomni/tool/genomics.py b/biomni/tool/genomics.py index 5ef96f710..3983fe512 100644 --- a/biomni/tool/genomics.py +++ b/biomni/tool/genomics.py @@ -473,7 +473,7 @@ def annotate_celltype_scRNA( data_info, data_lake_path, cluster="leiden", - llm="claude-3-5-sonnet-20241022", + llm=None, composition=None, ): """Annotate cell types based on gene markers and transferred labels using LLM. @@ -486,7 +486,7 @@ def annotate_celltype_scRNA( - data_dir (str): Directory containing the data files - data_info (str): Information about the scRNA-seq data (e.g., "homo sapiens, brain tissue, normal") - data_lake_path (str): Path to the data lake - - llm (str): Language model instance for cell type prediction, such as 'claude-3-haiku-20240307' + - llm (str, optional): Model for cell type prediction. Defaults to the deployment's configured model (BIOMNI_LLM) - composition (pd.DataFrame, optional): Transferred cell type composition for each cluster Returns: - str: Steps performed and file paths where results were saved diff --git a/biomni/tool/literature.py b/biomni/tool/literature.py index be39a4e4d..227c2bfd7 100644 --- a/biomni/tool/literature.py +++ b/biomni/tool/literature.py @@ -185,17 +185,16 @@ def query_pubmed(query: str, max_papers: int = 10, max_retries: int = 3) -> str: return f"Error querying PubMed: {e}" -def search_google(query: str, num_results: int = 3, language: str = "en") -> list[dict]: +def search_google(query: str, num_results: int = 3, language: str = "en") -> str: """Search using Google search. Args: - query (str): The search query (e.g., "protocol text or seach question") - num_results (int): Number of results to return (default: 10) + query (str): The search query (e.g., "protocol text or search question") + num_results (int): Number of results to return (default: 3) language (str): Language code for search results (default: 'en') - pause (float): Pause between searches to avoid rate limiting (default: 2.0 seconds) Returns: - List[dict]: List of dictionaries containing search results with title and URL + str: One Title / URL / Description block per result, or an empty string """ try: @@ -221,7 +220,7 @@ def advanced_web_search_claude( query: str, max_searches: int = 1, max_retries: int = 3, -) -> tuple[str, list[dict[str, str]], list]: +) -> str: """ Initiate an advanced web search by launching a specialized agent to collect relevant information and citations through multiple rounds of web searches for a given query. Craft the query carefully for the search agent to find the most relevant information. @@ -244,24 +243,21 @@ def advanced_web_search_claude( import anthropic - try: - from biomni.config import default_config - - model = default_config.llm - api_key = default_config.api_key - if not api_key: - api_key = credentials.getenv("ANTHROPIC_API_KEY") - except ImportError: - model = "claude-4-sonnet-latest" - api_key = credentials.getenv("ANTHROPIC_API_KEY") - - if "claude" not in model: - raise ValueError("Model must be a Claude model.") - - if not api_key: - raise ValueError("Set your api_key explicitly.") - - client = anthropic.Anthropic(api_key=api_key) + from biomni.config import default_config + from biomni.tool.availability import claude_web_search_problem + + # Checked here as well as when the tool list is built, because generated + # code can import the function whether or not it was advertised - and + # behind the LLM proxy a direct call to Anthropic must never be made. + problem = claude_web_search_problem() + if problem: + raise RuntimeError(problem) + + # ANTHROPIC_API_KEY only. default_config.api_key belongs to a custom model + # endpoint, and sending it to api.anthropic.com would hand that endpoint's + # credential to a third party. + client = anthropic.Anthropic(api_key=credentials.getenv("ANTHROPIC_API_KEY")) + model = default_config.llm tool_def = { "type": "web_search_20250305", "name": "web_search", diff --git a/biomni/tool/support_tools.py b/biomni/tool/support_tools.py index 9f48f0580..081d20f7d 100644 --- a/biomni/tool/support_tools.py +++ b/biomni/tool/support_tools.py @@ -41,7 +41,7 @@ from biomni import credentials from biomni.credentials import scrubbed_environ -from biomni.observability import session_id_var +from biomni.observability import carry_context_into_thread_pools, session_id_var logger = logging.getLogger(__name__) @@ -336,6 +336,9 @@ def execute_in_repl(command: str) -> str: preexisting_figures = set(_open_figure_numbers()) buffer = StringIO() _ensure_stdout_router() + # Thread pools the snippet starts stay in its session: their model calls + # name its user to the LLM proxy, and their prints land in this buffer. + carry_context_into_thread_pools() token = _active_buffer.set(buffer) try: diff --git a/biomni/tool/tool_description/genomics.py b/biomni/tool/tool_description/genomics.py index 787ff80a1..4908a464c 100644 --- a/biomni/tool/tool_description/genomics.py +++ b/biomni/tool/tool_description/genomics.py @@ -13,8 +13,8 @@ "type": "str", }, { - "default": "claude-3-5-sonnet-20241022", - "description": "Language model instance for cell type prediction", + "default": None, + "description": "Model for cell type prediction. Leave unset to use the deployment's configured model", "name": "llm", "type": "str", }, diff --git a/biomni/utils.py b/biomni/utils.py index d30de36a0..759159e1c 100644 --- a/biomni/utils.py +++ b/biomni/utils.py @@ -1,12 +1,15 @@ import ast import concurrent.futures import enum +import glob import importlib import json import os import pickle import subprocess import tempfile +import threading +import time import traceback import zipfile from typing import Any, ClassVar @@ -342,8 +345,10 @@ def write_python_code(request: str): # Route through get_llm so this call inherits the project's LLM resilience # config (max_retries / request_timeout -> provider-SDK 429/5xx backoff). - # A fresh BiomniConfig() honors BIOMNI_LLM_* env overrides at call time. - model = get_llm("claude-3-5-sonnet-20240620", source="Anthropic", config=BiomniConfig()) + # A fresh BiomniConfig() honors BIOMNI_LLM_* env overrides at call time, + # including the model: a pinned model id is eventually retired, and one the + # deployment did not configure is one its key or LLM proxy cannot reach. + model = get_llm(config=BiomniConfig()) template = """Write some python code to solve the user's problem. Return only python code in Markdown format, e.g.: @@ -898,11 +903,17 @@ def read_module2api(): "protocols", ] + from biomni.tool.availability import unavailable_tools + + # Tools this deployment cannot run are left out, not listed and left to + # fail: the agent is told to search first, so an advertised tool that + # cannot work gets called on nearly every question. + unavailable = unavailable_tools() module2api = {} for field in fields: module_name = f"biomni.tool.tool_description.{field}" module = importlib.import_module(module_name) - module2api[f"biomni.tool.{field}"] = module.description + module2api[f"biomni.tool.{field}"] = [api for api in module.description if api["name"] not in unavailable] return module2api @@ -948,6 +959,79 @@ def download_and_unzip(url: str, dest_dir: str) -> str: return f"Error: {e}" +# A download goes to a hidden partial file beside its target, named like this, +# and is renamed into place once complete (see _download_file). +_PARTIAL_SUFFIX = ".part" +# (connect, read) seconds. Without a timeout, a stalled connection held the chat +# turn that needed the file for as long as the socket stayed open. +_DOWNLOAD_TIMEOUT = (10, 60) +# A live download writes continuously, and gives up after the read timeout, so +# a partial file untouched this long was left by a process that is gone: killed +# mid-download, its pod evicted. +_STALE_PARTIAL_S = 15 * 60 + +_download_locks: dict[str, threading.Lock] = {} +_download_locks_guard = threading.Lock() + + +def _download_lock(path: str) -> threading.Lock: + """The lock that lets one thread of this process download ``path`` at a time.""" + with _download_locks_guard: + return _download_locks.setdefault(os.path.abspath(path), threading.Lock()) + + +def _remove_stale_partials(file_path: str) -> None: + """Delete partial downloads of ``file_path`` that their process abandoned.""" + directory, name = os.path.split(os.path.abspath(file_path)) + cutoff = time.time() - _STALE_PARTIAL_S + for partial in glob.glob(os.path.join(glob.escape(directory), f".{glob.escape(name)}.*{_PARTIAL_SUFFIX}")): + try: + if os.path.getmtime(partial) < cutoff: + os.remove(partial) + except OSError: + pass + + +def _download_file(url: str, file_path: str, desc: str) -> bool: + """Download ``url`` to ``file_path``, which appears only once it is complete. + + The data lake counts a file that exists as downloaded. Written in place, a + download cut short - the process killed, the pod evicted - left a truncated + dataset that no later query fetched again, and a chat could read a file + another chat was still downloading. So the bytes go to a hidden partial file + beside the target, renamed into place once all of them have arrived. + """ + partial = None + try: + directory = os.path.dirname(os.path.abspath(file_path)) + os.makedirs(directory, exist_ok=True) + _remove_stale_partials(file_path) + fd, partial = tempfile.mkstemp(dir=directory, prefix=f".{os.path.basename(file_path)}.", suffix=_PARTIAL_SUFFIX) + with os.fdopen(fd, "wb") as f, requests.get(url, stream=True, timeout=_DOWNLOAD_TIMEOUT) as response: + response.raise_for_status() + expected = int(response.headers.get("content-length", 0)) + received = 0 + with tqdm.tqdm(total=expected or None, unit="B", unit_scale=True, desc=desc, ncols=80) as pbar: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + received += len(chunk) + pbar.update(len(chunk)) + if expected and received != expected: + raise OSError(f"the connection closed after {received} of {expected} bytes") + os.replace(partial, file_path) + return True + except Exception as e: + print(f"✗ Failed to download {desc}: {e}") + return False + finally: + # Renamed away on success; still here, it is an incomplete download. + if partial is not None and os.path.exists(partial): + try: + os.remove(partial) + except OSError: + pass + + def check_and_download_s3_files( s3_bucket_url: str, local_data_lake_path: str, expected_files: list[str], folder: str = "data_lake" ) -> dict[str, bool]: @@ -966,38 +1050,6 @@ def check_and_download_s3_files( os.makedirs(local_data_lake_path, exist_ok=True) download_results = {} - def download_with_progress(url: str, file_path: str, desc: str) -> bool: - """Download file with progress bar.""" - try: - # Ensure directory exists for nested files - os.makedirs(os.path.dirname(file_path), exist_ok=True) - - response = requests.get(url, stream=True) - response.raise_for_status() - - total_size = int(response.headers.get("content-length", 0)) - - with open(file_path, "wb") as f: - if total_size > 0: - with tqdm.tqdm(total=total_size, unit="B", unit_scale=True, desc=desc, ncols=80) as pbar: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - pbar.update(len(chunk)) - else: - for chunk in response.iter_content(chunk_size=8192): - if chunk: - f.write(chunk) - return True - except Exception as e: - print(f"✗ Failed to download {desc}: {e}") - if os.path.exists(file_path): - try: - os.remove(file_path) - except OSError: - pass - return False - def cleanup_file(file_path: str): """Clean up file if it exists.""" if os.path.exists(file_path): @@ -1014,7 +1066,7 @@ def cleanup_file(file_path: str): with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp_zip: tmp_zip_path = tmp_zip.name - if download_with_progress(s3_zip_url, tmp_zip_path, f"{folder}.zip"): + if _download_file(s3_zip_url, tmp_zip_path, f"{folder}.zip"): print(f"Extracting {folder}.zip...") try: with zipfile.ZipFile(tmp_zip_path, "r") as zip_ref: @@ -1051,7 +1103,13 @@ def process_file(filename): local_file_path = os.path.join(local_data_lake_path, filename) s3_file_url = urljoin(s3_bucket_url + "/" + folder + "/", filename) - success = download_with_progress(s3_file_url, local_file_path, filename) + # One download of a file at a time: two chats that need the same dataset + # at once would otherwise both fetch it, each into a partial file of its + # own, doubling the disk a multi-GB file takes while it arrives. + with _download_lock(local_file_path): + if os.path.exists(local_file_path): # fetched while this one waited + return filename, True + success = _download_file(s3_file_url, local_file_path, filename) if success: print(f"✓ Successfully downloaded: {filename}") return filename, success diff --git a/biomni/workspace_prefs.py b/biomni/workspace_prefs.py index 1c232047f..f8875c843 100644 --- a/biomni/workspace_prefs.py +++ b/biomni/workspace_prefs.py @@ -33,6 +33,8 @@ import json import logging import os +import re +import stat import tempfile from dataclasses import asdict, dataclass, field from datetime import UTC, datetime @@ -292,35 +294,86 @@ def _nearest_existing(path: str) -> str: # exactly what a restart discards. _RAM_FSTYPES = frozenset({"tmpfs", "ramfs", "devtmpfs"}) _CONTAINER_LAYER_FSTYPES = frozenset({"overlay", "overlayfs", "aufs"}) +# Network filesystems whose ownership the kubelet cannot change: fsGroup is +# silently ignored on them, so they need their own remedy. +_SMB_FSTYPES = frozenset({"cifs", "smb3", "smbfs"}) +_NFS_FSTYPES = frozenset({"nfs", "nfs4"}) +# SMB options under which the file server's ACLs, not the mount options, +# decide a file's owner and mode. +_SMB_ACL_OPTIONS = ("cifsacl", "modefromsid") +_MOUNTS_FILE = "/proc/self/mounts" +_OCTAL_ESCAPE_RE = re.compile(r"\\([0-7]{3})") + + +@dataclass(frozen=True) +class MountInfo: + """One line of the mount table: what is mounted where, and how.""" + + mountpoint: str + source: str + fstype: str + options: tuple[str, ...] + + def option(self, name: str) -> str | None: + """The value of ``name=value`` among the options, or ``None``.""" + prefix = f"{name}=" + for option in self.options: + if option.startswith(prefix): + return option[len(prefix) :] + return None + + @property + def read_only(self) -> bool: + return "ro" in self.options + + +def _unescape_mount_field(field: str) -> str: + # /proc escapes whitespace and backslashes octal-style (a space is \040). + return _OCTAL_ESCAPE_RE.sub(lambda m: chr(int(m.group(1), 8)), field) -def _mount_fstype(path: str) -> str | None: - """Filesystem type backing ``path``, per ``/proc/self/mounts``. + +def mount_info(path: str) -> MountInfo | None: + """The mount backing ``path``, per ``/proc/self/mounts``. ``None`` when the mount table is unavailable (macOS, a container with no ``/proc``), which callers treat as "cannot tell" rather than as an answer. The longest matching mount point wins, since mount points nest. + + Symlinks are resolved first: a configured path that links onto a share is + stored on the share, not on whatever holds the link. And the table is read + the way the OS reads file names, since the kernel escapes only whitespace + and backslashes and passes any other byte of a mount point through as is. """ try: - with open("/proc/self/mounts", encoding="utf-8") as handle: + with open(_MOUNTS_FILE, encoding="utf-8", errors="surrogateescape") as handle: lines = handle.readlines() except OSError: return None - target = os.path.abspath(path) - best_mountpoint = "" - fstype: str | None = None + target = os.path.realpath(path) + best: MountInfo | None = None for line in lines: fields = line.split() - if len(fields) < 3: + if len(fields) < 4: continue - # Mount points are escaped octal-style in /proc (a space is \040). - mountpoint = fields[1].replace("\\040", " ") + mountpoint = _unescape_mount_field(fields[1]) if target != mountpoint and not target.startswith(mountpoint.rstrip("/") + "/"): continue - if len(mountpoint) >= len(best_mountpoint): - best_mountpoint, fstype = mountpoint, fields[2] - return fstype + if best is None or len(mountpoint) >= len(best.mountpoint): + best = MountInfo( + mountpoint=mountpoint, + source=_unescape_mount_field(fields[0]), + fstype=fields[2], + options=tuple(fields[3].split(",")), + ) + return best + + +def _mount_fstype(path: str) -> str | None: + """Filesystem type backing ``path``, or ``None`` when it cannot be told.""" + info = mount_info(path) + return info.fstype if info else None def on_mounted_volume(path: str) -> bool: @@ -368,6 +421,202 @@ def is_writable_dir(path: str | None) -> bool: return os.access(target, os.W_OK | os.X_OK) +def _process_ids() -> tuple[int | None, int | None, frozenset[int]]: + """This process's uid, primary gid and every group it belongs to.""" + if not hasattr(os, "getuid"): + return None, None, frozenset() + return os.getuid(), os.getgid(), frozenset({os.getgid(), *os.getgroups()}) + + +def _mode_allows(st: os.stat_result, uid: int, groups: frozenset[int], *, write: bool) -> bool: + """Whether owner, group and mode bits alone let ``uid`` list (or add to) a directory.""" + if uid == 0: + return True + need = stat.S_IWUSR if write else stat.S_IRUSR + mode = st.st_mode + if st.st_uid == uid: + return bool(mode & need and mode & stat.S_IXUSR) + if st.st_gid in groups: + return bool(mode & (need >> 3) and mode & stat.S_IXGRP) + return bool(mode & (need >> 6) and mode & stat.S_IXOTH) + + +def _is_read_only(target: str, info: MountInfo | None) -> bool: + try: + return bool(os.statvfs(target).f_flag & os.ST_RDONLY) + except (OSError, AttributeError): + return bool(info and info.read_only) + + +def _explain_access(target: str, *, write: bool, wanted: str | None = None) -> str: + """Why this process cannot read or write the existing directory ``target``. + + Written for whoever has to fix it: who owns the directory, what storage it + is on, and the remedy for that kind of storage. "Not writable by uid 57439" + was true and still cost the GRIP team a round trip, because a share mounted + without ownership options, one mounted read-only, and one whose new options + had not reached the node yet all read exactly the same. + + ``wanted`` is the directory the caller actually needs, when ``target`` is + only its nearest existing parent. + """ + if not os.path.isdir(target): + return f"{target} is not a directory" if os.path.exists(target) else f"{target} does not exist" + info = mount_info(target) + share = _without_server(info.source) if info else "" + on = f" ({info.fstype} mount of {share})" if info else "" + in_image = bool(info and info.fstype in _CONTAINER_LAYER_FSTYPES) + if write and _is_read_only(target, info): + if in_image: + # readOnlyRootFilesystem: there is no volume here to remount. + return ( + f"{target} is inside the container image, which is mounted read-only; " + f"mount a writable volume at {wanted or target}" + ) + return f"{target} is on a read-only mount{on}; mount the volume read-write" + try: + st = os.stat(target) + except OSError as exc: + return f"{target} cannot be inspected ({exc.strerror or exc})" + + uid, gid, groups = _process_ids() + owner = f"owned by uid {st.st_uid}, gid {st.st_gid} with mode {stat.S_IMODE(st.st_mode):04o}" + if uid is None: + return f"{target} is {owner}" + access = "writing" if write else "reading" + if _mode_allows(st, uid, groups, write=write): + return ( + f"{target} is {owner}, which permits {access}, so something else refuses it: an ACL, " + "SELinux or AppArmor, or the file server itself" + ) + if info and info.fstype in _SMB_FSTYPES: + acl = next((name for name in _SMB_ACL_OPTIONS if name in info.options), None) + if acl: + return ( + f"{target} is {owner} on an SMB share ({share}) mounted with {acl}, so its owner and mode come " + f"from the file server's ACLs; grant uid {uid} {'write' if write else 'read'} access in the " + f"share's ACLs, or mount it without {acl} and with uid={uid},gid={gid},dir_mode=0770," + "file_mode=0770" + ) + # SMB has no Unix ownership of its own: the client presents whatever the + # mount options say, so the live options are the diagnosis. + live = ",".join(f"{name}={info.option(name)}" for name in ("uid", "gid", "dir_mode") if info.option(name)) + return ( + f"{target} is {owner} on an SMB share ({share}). Its owner and mode come only from the mount " + f"options, which are {live or 'unset (uid=0,gid=0)'} on the mount this container sees. Mount it with " + f"uid={uid},gid={gid},dir_mode=0770,file_mode=0770; new options only apply once the share is " + "mounted afresh on the node" + ) + if info and info.fstype in _NFS_FSTYPES: + return ( + f"{target} is {owner} on an NFS export ({share}), which ignores fsGroup; " + f"grant uid {uid} {'write' if write else 'read'} access on the NFS server" + ) + if info and info.fstype.startswith("fuse"): + # blobfuse2, s3fs, gcsfuse...: the driver presents ownership from its own + # options, and the kubelet cannot change it. The driver is named in the + # type ("fuse.s3fs") or, for plain "fuse", as the source ("blobfuse2"). + driver = info.fstype.partition(".")[2] or share + return ( + f"{target} is {owner} on a FUSE mount ({driver}), which ignores fsGroup; mount it " + f"with the driver's options that let uid {uid} {'write' if write else 'read'} (allow_other, and " + "its uid, gid or umask options)" + ) + if in_image: + # Usually a volume that is not mounted where the setting says. Opening + # up the image would not help: whatever is written there is lost when + # the pod restarts. + return ( + f"{target} is {owner}, inside the container image rather than on a mounted volume; " + f"mount {'a writable' if write else 'the'} volume at {wanted or target}" + ) + if st.st_uid == uid: + # Already ours: ownership advice would send someone chowning a + # directory to the uid that owns it. + return ( + f"{target} is {owner}{on}, which denies its own owner (uid {uid}); give the owner access with chmod u+rwx" + ) + if st.st_gid in groups: + return f"{target} is {owner}{on}, which denies its group; give the group access with chmod g+rwx" + return ( + f"{target} is {owner}{on}; set securityContext.fsGroup: {gid} on the pod (it does not apply to hostPath " + f"volumes), or chown it to uid {uid}" + ) + + +def _without_server(source: str) -> str: + """A mount source minus the server it names. + + ``//acct.file.core.windows.net/share`` becomes ``/share`` and + ``nfs.example:/export`` becomes ``/export``. The server of an Azure Files + share is the storage account, which stays out of what users are shown and + what is logged - the same reason :func:`storage_facts` drops the + ``username=`` option. Devices and pseudo filesystems pass through. + """ + if source.startswith("//"): + return "/" + source[2:].partition("/")[2] + server, separator, export = source.partition(":/") + if separator and server and "/" not in server: + return "/" + export + return source + + +def explain_unwritable(path: str) -> str: + """Why ``path`` cannot be written (or created), as specifically as the system can tell. + + About the directory actually in the way: ``path`` itself, or for a path + that does not exist yet, the nearest parent it would be created in. + """ + return _explain_access(_nearest_existing(path), write=True, wanted=os.path.abspath(path)) + + +def is_readable_dir(path: str | None) -> bool: + """Whether ``path`` is a directory this process can list and enter.""" + if not path: + return False + return os.path.isdir(path) and os.access(path, os.R_OK | os.X_OK) + + +def explain_unreadable(path: str) -> str: + """Why ``path`` cannot be listed, as specifically as the system can tell.""" + return _explain_access(os.path.abspath(path), write=False) + + +# The mount options worth a log line - the ones that decide who may write. The +# rest (an SMB mount's username, cache settings) are noise at best. +_LOGGED_FLAGS = frozenset({"ro", "rw", "noperm", "forceuid", "forcegid", "cifsacl", "noexec"}) +_LOGGED_KEYS = frozenset({"uid", "gid", "file_mode", "dir_mode", "vers", "sec"}) + + +def storage_facts(path: str) -> dict[str, Any]: + """What a log line needs to diagnose ``path`` without shell access to the pod.""" + target = _nearest_existing(path) + facts: dict[str, Any] = { + "path": os.path.abspath(path), + "checked": target, + "readable": is_readable_dir(target), + "writable": is_writable_dir(path), + } + try: + st = os.stat(target) + facts.update(owner_uid=st.st_uid, owner_gid=st.st_gid, mode=f"{stat.S_IMODE(st.st_mode):04o}") + except OSError: + pass + info = mount_info(target) + if info is not None: + facts.update( + mountpoint=info.mountpoint, + fstype=info.fstype, + source=_without_server(info.source), + options=",".join(o for o in info.options if o in _LOGGED_FLAGS or o.split("=", 1)[0] in _LOGGED_KEYS), + ) + if not facts["readable"]: + facts["read_reason"] = explain_unreadable(target) + if not facts["writable"]: + facts["write_reason"] = explain_unwritable(path) + return facts + + def build_prefs_store(workspace_root: str | None = None) -> PrefsStore: """Pick the best available preferences location for this deployment. @@ -473,19 +722,64 @@ def _output_candidates(prefs: WorkspacePrefs, workspace_root: str | None) -> lis # nor a writable workspace exists - but flagged as ephemeral so the UI can # warn instead of silently losing the user's results. candidates.append((os.path.abspath(os.path.join(os.getcwd(), "runs")), "cwd-fallback")) - return candidates + + # One entry per directory, under the name that put it first. Setting + # BIOMNI_OUTPUT_ROOT to /biomni-outputs - the GRIP setup - makes it + # the workspace default too, and a directory that cannot be used would + # otherwise be reported twice, under two names. + unique: dict[str, str] = {} + for path, source in candidates: + unique.setdefault(path, source) + return list(unique.items()) -def _describe_skipped(entries: list[tuple[str, str]]) -> str: - """Human explanation of which configured locations were passed over, and why. +def _describe_skipped(entries: list[tuple[str, str, str]]) -> str: + """Which configured locations were passed over, and exactly why. - Names the setting and the path, because "fell back from env" told a reader - that something was rejected without telling them what to go and fix. + ``entries`` are ``(path, source, reason)``. Names the setting and the path, + because "fell back from env" told a reader that something was rejected + without telling them what to go and fix. Grouped by the directory actually + in the way: ``BIOMNI_OUTPUT_ROOT`` and the workspace default usually fail + together, on the same mount, and explaining that twice buried the one fix + both needed. """ + groups: dict[str, list[tuple[str, str, str]]] = {} + for entry in entries: + groups.setdefault(_nearest_existing(entry[0]), []).append(entry) uid = getattr(os, "getuid", lambda: None)() by_whom = f" by uid {uid}" if uid is not None else "" - parts = [f"{_SOURCE_LABELS.get(source, source)} ({path}) is not writable{by_whom}" for path, source in entries] - return "; ".join(parts) + sentences = [] + for members in groups.values(): + names = _join_names([f"{_SOURCE_LABELS.get(source, source)} ({path})" for path, source, _ in members]) + verb = "is" if len(members) == 1 else "are" + sentence = f"{names} {verb} not writable{by_whom}: {members[0][2]}" + # Most labels start "the ...", and each one starts a sentence. + sentences.append(sentence[0].upper() + sentence[1:]) + return ". ".join(sentences) + + +def _join_names(names: list[str]) -> str: + """``a``, ``a and b``, ``a, b and c``.""" + return names[0] if len(names) == 1 else f"{', '.join(names[:-1])} and {names[-1]}" + + +# (source, path, reason) combinations already warned about. Resolution runs +# several times a session, and the same unwritable mount repeated on every one +# buried everything else in the log; a changed reason is new, and is logged. +_warned_unwritable: set[tuple[str, str, str]] = set() + + +def _warn_unwritable_once(source: str, path: str, reason: str) -> None: + key = (source, path, reason) + if key in _warned_unwritable: + return + _warned_unwritable.add(key) + logger.warning( + "output location %s (%s) is not writable, falling back to the next candidate: %s", + _SOURCE_LABELS.get(source, source), + path, + reason, + ) def resolve_output_dir(prefs: WorkspacePrefs, *, workspace_root: str | None = None) -> OutputTarget: @@ -502,36 +796,28 @@ def resolve_output_dir(prefs: WorkspacePrefs, *, workspace_root: str | None = No tell the two apart. """ candidates = _output_candidates(prefs, workspace_root) - skipped: list[tuple[str, str]] = [] + skipped: list[tuple[str, str, str]] = [] for path, source in candidates: if is_writable_dir(path): - reason = None - if skipped: - reason = f"fell back from {_describe_skipped(skipped)}" return OutputTarget( path=path, source=source, writable=True, - reason=reason, + reason=_describe_skipped(skipped) if skipped else None, mounted_volume=on_mounted_volume(path), ) + # Worked out once: it reads the mount table and asks the file system, + # which on an SMB share is a round trip to the server. + reason = explain_unwritable(path) if source != "cwd-fallback": - logger.warning( - "output location %s (%s) is not writable; falling back to the next candidate", - _SOURCE_LABELS.get(source, source), - path, - ) - skipped.append((path, source)) + _warn_unwritable_once(source, path, reason) + skipped.append((path, source, reason)) path, source = candidates[-1] - logger.error("no writable output location found; tried %s", _describe_skipped(skipped)) - return OutputTarget( - path=path, - source=source, - writable=False, - reason="no writable output location found; set BIOMNI_OUTPUT_ROOT to a mounted volume", - ) + tried = _describe_skipped(skipped) + logger.error("no writable output location found: %s", tried) + return OutputTarget(path=path, source=source, writable=False, reason=tried) def ensure_output_dir(target: OutputTarget) -> bool: diff --git a/biomni_env/adworkbench_env.yml b/biomni_env/adworkbench_env.yml index 1c9f92699..1dd92c513 100644 --- a/biomni_env/adworkbench_env.yml +++ b/biomni_env/adworkbench_env.yml @@ -651,7 +651,7 @@ dependencies: - zarr==2.18.7 - zict==3.0.0 - zstandard==0.23.0 - - chainlit>=1.0 + - chainlit==2.11.1 - mi-googlesearch-python==1.3.0.post1 - alabaster==1.0.0 - imagesize==1.4.1 diff --git a/chainlit_app.py b/chainlit_app.py index e5137c961..5830d7df5 100644 --- a/chainlit_app.py +++ b/chainlit_app.py @@ -20,11 +20,9 @@ import os import re import sys -import threading import time import uuid import zipfile -from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace from pathlib import Path @@ -78,12 +76,21 @@ from biomni.config import default_config, resolve_data_lake_root, resolve_default_llm from biomni.health import register_health_routes from biomni.identity import ( + Login, + SessionIdentity, UserIdentity, + auth_config_warnings, + bound_identity, + describe_auth_config, + login_is_current, password_identity, resolve_identity, + resolve_session_identity, + set_current_identity, single_user_mode, trust_auth_headers, ) +from biomni.llm_proxy import describe_llm_proxy, llm_proxy_settings, proxy_identity_problem from biomni.observability import ( RunHeartbeat, bind_run, @@ -104,11 +111,15 @@ ScopeResolution, WorkspacePrefs, build_prefs_store, + default_prefs, env_flag, + explain_unreadable, + is_readable_dir, load_prefs, normalize_scope_entries, resolve_output_dir, resolve_scope, + storage_facts, ) from chainlit.input_widget import InputWidget, MultiSelect, Switch, Tags, TextInput from langchain_core.messages import AIMessage, HumanMessage @@ -160,7 +171,9 @@ def _extract_final_answer(state: dict) -> str: # intentionally not re-exported - anything that needs them should # `from chainlit_ui.planning import PLANNING_SYSTEM_PROMPT, AD1_PLANNING_SYSTEM_PROMPT`. from chainlit_ui.datasets import build_suggested_prompts_markdown +from chainlit_ui.graph_stream import stream_langgraph from chainlit_ui.live_runs import LIVE_RUNS +from chainlit_ui.llm_failures import describe_llm_failure from chainlit_ui.persistence import build_data_layer, ensure_auth_secret from chainlit_ui.planning import PLAN_ACTIONS, STALE_PLAN_ACTION_NOTE from chainlit_ui.planning import ( @@ -472,6 +485,10 @@ class WorkspaceSession: top_level_dirs: list[str] = field(default_factory=list) inventory: str = "" runs: list[RunRecord] = field(default_factory=list) + # Why the workspace cannot be listed, when it cannot. It exists (or it + # would not be the workspace) but every folder in it is invisible, which + # without this reads as an empty workspace rather than a broken mount. + workspace_problem: str | None = None @property def persistence_label(self) -> str: @@ -499,22 +516,31 @@ def _refresh_workspace_view(ws: WorkspaceSession) -> WorkspaceSession: """ ws.scope = resolve_scope(ws.prefs, ws.workspace_root) ws.output = resolve_output_dir(ws.prefs, workspace_root=ws.workspace_root) + return _survey_workspace(ws) + + +def _survey_workspace(ws: WorkspaceSession) -> WorkspaceSession: + """What is in the workspace, given the session's scope and output directory.""" + root = ws.workspace_root + ws.workspace_problem = explain_unreadable(root) if root and not is_readable_dir(root) else None # The output directory usually sits inside the workspace and gains a run # folder per query; it is a destination, never an input, so keep it out of - # the picker and out of what the agent is told it can read. - ws.top_level_dirs = list_top_level_dirs(ws.workspace_root, exclude_paths=[ws.output.path]) - ws.inventory = build_scope_inventory(ws.scope, ws.workspace_root, top_level_dirs=ws.top_level_dirs) + # the picker and out of what the agent is told it can read. An unreadable + # workspace is not listed at all: why is already known, and the attempt + # would only log the same failure again on every start and save. + ws.top_level_dirs = [] if ws.workspace_problem else list_top_level_dirs(root, exclude_paths=[ws.output.path]) + ws.inventory = build_scope_inventory(ws.scope, root, top_level_dirs=ws.top_level_dirs) return ws -def _build_workspace_session(session_id: str, header_source: object) -> WorkspaceSession: - """Resolve identity, load stored preferences and derive the session view. +def _build_workspace_session(session: SessionIdentity) -> WorkspaceSession: + """Load stored preferences for the session's user and derive the session view. Blocking (filesystem + preference load), so callers must run it in an executor. Never raises: any failure degrades to defaults, because a user who cannot load their settings should still get a working session. """ - identity = resolve_identity(header_source, session_id=session_id) + identity = session.identity workspace_root = _resolve_workspace_root() prefs_key = identity.scoped_key() @@ -527,7 +553,10 @@ def _build_workspace_session(session_id: str, header_source: object) -> Workspac store = build_prefs_store(workspace_root) registry = build_run_registry(workspace_root) else: - store = NullPrefsStore("no authenticated user; set BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE=true to override") + # The notice says why there is no signed-in user - gateway headers being + # ignored is an operator's fix, while having no gateway at all may be + # exactly what the deployment intends. + store = NullPrefsStore(session.notice or "no signed-in user") registry = RunRegistry(None) prefs = load_prefs(store, prefs_key) @@ -545,36 +574,63 @@ def _build_workspace_session(session_id: str, header_source: object) -> Workspac # Reconciling here (rather than at write time) is what turns a pod restart # into an honest "interrupted" on the user's next visit. ws.runs = registry.reconcile(prefs_key) - return _refresh_workspace_view(ws) + return _survey_workspace(ws) + + +def _as_sentence(text: str) -> str: + text = text.strip() + return text if not text or text.endswith((".", "!", "?")) else f"{text}." + + +def _output_warnings(output: OutputTarget) -> list[str]: + """What a user must know about where results go, before a long run starts. + + Shared by the settings dialog and the message confirming a settings change, + so the two can never describe the same directory differently. + """ + if not output.writable: + return [_as_sentence(f"⚠️ Results cannot be saved: no output location is writable. {output.reason or ''}")] + warnings = [] + # Shown even when the resolved directory works: a working fallback is + # exactly the case where a mis-set BIOMNI_OUTPUT_ROOT goes unnoticed. + if output.reason: + warnings.append( + _as_sentence( + f"⚠️ The configured location could not be used, so results go to {output.path} instead. {output.reason}" + ) + ) + if output.is_ephemeral: + warnings.append( + "⚠️ This is container-local storage: results are lost when the application restarts. " + "Set a path on a mounted volume, or ask an operator to set BIOMNI_OUTPUT_ROOT." + ) + return warnings def _output_dir_description(ws: WorkspaceSession) -> str: """Help text for the output-directory field, including its warnings. - The two warnings used to live in the sidebar panel, and they are the only - part of it that was load-bearing: a user whose results are going somewhere + The warnings used to live in the sidebar panel, and they are the only part + of it that was load-bearing: a user whose results are going somewhere unwritable, or somewhere a restart will erase, has to be told before they - start a long run rather than after. + start a long run rather than after. Paragraphs are separated by blank lines, + which ``public/biomni.css`` makes the dialog render as breaks. """ - text = ( - f"Where run results are written. Currently resolved to {ws.output.path} " - f"(source: {ws.output.source}). Leave as-is to keep the deployment default." + # The field shows a saved directory even when it cannot be used, so the way + # back to the default is to clear it, not to leave it alone. + reset = ( + "Clear the field to use the deployment's default." + if ws.prefs.output_dir + else "Leave as-is to keep the deployment's default." ) - if not ws.output.writable: - text += f"\n\n⚠️ That directory is not writable. {ws.output.reason or ''}".rstrip() - else: - if ws.output.is_ephemeral: - text += ( - "\n\n⚠️ Container-local storage: results are lost when the application restarts. " - "Set a path on a mounted volume, or ask an operator to configure BIOMNI_OUTPUT_ROOT." - ) - # Shown even when the resolved directory is fine: a working fallback is - # exactly the case where a mis-set BIOMNI_OUTPUT_ROOT goes unnoticed. - if ws.output.reason: - text += f"\n\n⚠️ This is not the configured location - {ws.output.reason}." - if ws.persistence_label: - text += f"\n\nSettings are stored in {ws.persistence_label}." - return text + return "\n\n".join([f"Where run results are written. {reset}", *_output_warnings(ws.output)]) + + +def _remember_settings_description(ws: WorkspaceSession) -> str: + """Where settings are kept, or why they cannot be kept at all.""" + if isinstance(ws.store, NullPrefsStore): + return _as_sentence(f"Settings apply to this session only: {ws.store.reason}") + return f"Saved in {ws.persistence_label}." def _workspace_settings_widgets(ws: WorkspaceSession) -> list[InputWidget]: @@ -621,12 +677,17 @@ def _workspace_settings_widgets(ws: WorkspaceSession) -> list[InputWidget]: ) ) + extra_description = "Sub-folders deeper than the top level, e.g. studyA/processed" + if ws.workspace_problem: + extra_description += "\n\n" + _as_sentence( + f"⚠️ The workspace cannot be read, so there are no folders to choose from. {ws.workspace_problem}" + ) widgets.append( Tags( id="scope_extra_paths", label="Additional paths (workspace-relative)", initial=[p for p in ws.prefs.scope_paths if p not in known], - description="Sub-folders deeper than the top level, e.g. studyA/processed", + description=extra_description, ) ) @@ -644,6 +705,7 @@ def _workspace_settings_widgets(ws: WorkspaceSession) -> list[InputWidget]: id="remember_settings", label="Remember these settings for my next session", initial=ws.prefs.remember, + description=_remember_settings_description(ws), ) ) return widgets @@ -747,18 +809,15 @@ def _user_from_identity(identity: UserIdentity) -> cl.User: """One place that turns a resolved identity into Chainlit's user. Shared by both auth paths so preferences, run records and chat threads - always key off the same string, whichever way the caller signed in. + always key off the same string, whichever way the caller signed in. The + metadata carries the whole identity, because the login is the one record of + it every websocket connection is guaranteed to have (see + ``_resolve_session_identity``). """ return cl.User( identifier=identity.scoped_key(), display_name=identity.display_name, - metadata={ - "source": identity.source, - "email": identity.email or "", - "workspace_id": identity.workspace_id or "", - "given_name": identity.given_name or "", - "family_name": identity.family_name or "", - }, + metadata=identity.to_metadata(), ) @@ -820,17 +879,155 @@ def password_auth_callback(username: str, password: str) -> cl.User | None: logger.info("password sign-in enabled (BIOMNI_DEMO_PASSWORD is set)") +def _login_of(user) -> Login | None: + """What a Chainlit user object records about the login it came from.""" + if user is None: + return None + return Login(getattr(user, "identifier", None), getattr(user, "metadata", None)) + + +def _reject_stale_login() -> None: + """Refuse a login cookie that names someone other than the gateway does. + + A Chainlit login outlives the platform session that created it by days, and + the browser only signs in again when it has no valid cookie. So when a second + person signs in to the platform on a shared machine, the gateway names them + on every request while the leftover cookie still names the first - who would + otherwise be handed the new person's conversations, preferences and model + usage. + + Every Chainlit HTTP endpoint resolves its caller through ``get_current_user``, + so overriding that one dependency covers all of them: a contradicted cookie + gets a 401, and the browser responds exactly as it does to an expired one, + by signing in again as whoever the gateway names. The websocket + authenticates separately, and ``_resolve_session_identity`` holds the same + line there. + """ + from chainlit.auth import get_current_user, reuseable_oauth + from chainlit.server import app + from fastapi import Depends, HTTPException, Request + + # Chainlit's own token extractor (cookie or bearer header), so the override + # reads the login exactly the way the dependency it replaces does. + async def gateway_checked_user(request: Request, token: str | None = Depends(reuseable_oauth)): + user = await get_current_user(token) + login = _login_of(user) + if login is not None and not login_is_current(login, request.headers): + logger.warning("refusing a login that names a different user than the authentication gateway") + raise HTTPException(status_code=401, detail="login does not match the authentication gateway") + return user + + app.dependency_overrides[get_current_user] = gateway_checked_user + + +# Only behind a gateway: without one there is nothing to check a login against. +if trust_auth_headers(): + try: + _reject_stale_login() + except Exception: # pragma: no cover - defensive: never block startup + logger.warning("Could not install the stale-login check", exc_info=True) + + +def _log_access_config() -> None: + """One structured line saying how callers are identified and models reached. + + The two settings interact - the LLM proxy can only bill a gateway identity - + and the failure when they disagree is an error on every model call, long + after startup. Saying so here is what makes it a one-line diagnosis. + """ + proxy = describe_llm_proxy() + emit_event( + "access_config", + logger=logger, + **describe_auth_config(), + password_login=bool(demo_password()), + llm_proxy=proxy, + ) + for warning in auth_config_warnings(): + logger.warning(warning) + if proxy.get("error"): + logger.error("LLM proxy misconfigured: %s", proxy["error"]) + elif proxy.get("enabled"): + if not proxy.get("token_configured"): + logger.error("BIOMNI_LLM_PROXY_URL is set but BIOMNI_LLM_PROXY_API_KEY is not; every model call will fail") + if not trust_auth_headers(): + logger.error( + "the LLM proxy needs each caller's platform user and workspace id, which only a trusted " + "authentication gateway supplies; with BIOMNI_TRUST_AUTH_HEADERS off every model call will be " + "refused" + ) + + +_log_access_config() + + +def _log_storage_check() -> None: + """One structured line per storage location the app depends on, at startup. + + Storage problems otherwise surface one user at a time, in a settings + dialog, as a sentence an operator has to ask for a screenshot of. These are + logged before any session exists, from what only the pod can see: the live + mount options, the owner and mode this process meets, and whether it can + read and write - enough to tell a share mounted without ownership options + from one whose new options have not reached the node yet. + """ + workspace = _resolve_workspace_root() + if workspace is None: + configured = ", ".join(f"{env}={path}" for env, path in _resolve_user_data_roots()) + if configured: + logger.error("no configured workspace directory exists: %s", configured) + else: + facts = storage_facts(workspace) + level = logging.INFO if facts["readable"] else logging.ERROR + emit_event("storage_check", logger=logger, level=level, location="workspace", **facts) + + # Where results go for a user with no preference - which also logs, once, + # why each configured location before it was passed over. + output = resolve_output_dir(default_prefs(), workspace_root=workspace) + if not output.writable: + level = logging.ERROR + elif output.reason or output.is_ephemeral: + # Working, but not where the deployment meant results to go, or not + # somewhere they survive a restart. + level = logging.WARNING + else: + level = logging.INFO + emit_event( + "storage_check", + logger=logger, + level=level, + location="outputs", + output_source=output.source, + ephemeral=output.is_ephemeral, + passed_over=output.reason, + **storage_facts(output.path), + ) + + state_dir = os.getenv("BIOMNI_STATE_DIR", "").strip() + if state_dir: + facts = storage_facts(state_dir) + level = logging.INFO if facts["writable"] else logging.WARNING + emit_event("storage_check", logger=logger, level=level, location="state_dir", **facts) + + +try: + _log_storage_check() +except Exception: # pragma: no cover - defensive: a diagnostic must never block startup + logger.warning("Could not check storage locations", exc_info=True) + + @cl.data_layer def data_layer(): """Durable storage for conversations, or ``None`` to keep them in memory. Enabled only where the storage key is stable enough for a user to find their - own threads again: behind a gateway, or in the explicitly opted-in - single-user mode. Under a per-connection anonymous key every page load would - start a fresh identity, so the sidebar would fill with orphaned threads that - nobody could ever reopen. + own threads again: behind a gateway, behind the demo password (the username + is the key), or in the explicitly opted-in single-user mode. Under a + per-connection anonymous key every page load would start a fresh identity, + so the sidebar would fill with orphaned threads that nobody could ever + reopen. """ - if not (trust_auth_headers() or single_user_mode()): + if not (trust_auth_headers() or single_user_mode() or demo_password()): logger.info( "chat history is disabled: no authentication gateway is configured. " "Set BIOMNI_TRUST_AUTH_HEADERS (behind a gateway) or " @@ -862,50 +1059,6 @@ async def run_in_executor(fn, *args): return await loop.run_in_executor(None, lambda: ctx.run(fn, *args)) -async def stream_langgraph(agent_app, inputs, config): - """Yield LangGraph state dicts asynchronously from a sync stream. - - The graph runs in a worker thread, which cancellation cannot reach: an - asyncio task that is cancelled stops *reading*, and a thread that nobody - reads goes right on calling the model. Stop therefore used to end the - conversation while the run it stopped kept working - for as long as the - whole plan took. The flag below is the cooperative half of the stop: the - producer checks it between graph nodes, so the run ends after the step it - was already in. - """ - queue: asyncio.Queue = asyncio.Queue() - loop = asyncio.get_event_loop() - stop = threading.Event() - - def _producer(): - try: - for state in agent_app.stream(inputs, stream_mode="values", config=config): - asyncio.run_coroutine_threadsafe(queue.put(state), loop) - if stop.is_set(): - break - finally: - asyncio.run_coroutine_threadsafe(queue.put(None), loop) # sentinel - - executor = ThreadPoolExecutor(max_workers=1) - # Capture the active context here (caller thread) and replay it in the - # producer thread so the agent graph's logs (LLM calls, code-execution - # audit) carry the session_id/run_id. - ctx = capture_context() - executor.submit(ctx.run, _producer) - - try: - while True: - state = await queue.get() - if state is None: - break - yield state - finally: - stop.set() - # Not waited on: the worker may be mid-node, and holding the event loop - # for it would freeze the UI. It exits at the next check. - executor.shutdown(wait=False) - - # --------------------------------------------------------------------------- # Chat profiles (A1 vs AD1) # --------------------------------------------------------------------------- @@ -949,7 +1102,7 @@ async def set_chat_profiles(): "For each gene: (1) pull GWAS significance from any local summary stats " "(GCST90027158, NG00075, or NG00052), (2) check brain eQTL evidence in NG00105 or " "SingleBrain, (3) look up proteomic levels in NG00102 if available, and " - "(4) find CRISPR dependency scores from CRISPRbrain relevant screen from the biomni-AD datalake." + "(4) find CRISPR dependency scores from the relevant CRISPRbrain screen in the biomni-AD data lake. " "Compile everything into a single comparison table and a radar chart per gene." ), icon="/public/avatars/ad1.png", @@ -1046,15 +1199,21 @@ async def on_chat_resume(thread: dict): the tab and came back - it is redirected into this connection, so the rest of it streams in here instead of into a socket nobody is holding. """ - history = _history_from_thread(thread) - reattached = await _follow_live_run(thread.get("id")) - emit_event( - "chat_resume", - steps=len(thread.get("steps") or []), - history_turns=len(history), - reattached=reattached, - ) - await _start_session(resumed_history=history) + await _start_session(resumed_thread=thread) + + +def _withhold_transcript(thread: dict) -> None: + """Keep a refused session from receiving the conversation it asked to reopen. + + Chainlit sends the stored transcript to the browser after the resume handler + returns, whatever the handler decided - and it sends the very dict it handed + the handler, so emptying that is what makes refusing the session refuse the + conversation too. The first line of defence is ``_reject_stale_login``: a + browser holding a contradicted login is sent to sign in again before it ever + opens a websocket. This is the second. + """ + thread["steps"] = [] + thread["elements"] = [] async def _follow_live_run(thread_id: str | None) -> bool: @@ -1079,6 +1238,35 @@ async def _follow_live_run(thread_id: str | None) -> bool: _RESUME_ANNOUNCE_DELAY_S = 1.0 +async def _send_notice(content: str, *, resumed: bool = False) -> None: + """Tell the user something about this session, outside the conversation. + + Never stored with the conversation. A notice describes the session it was + sent in - a workspace that could not be read, a sign-in that changed - and a + stored copy would be shown again, and replayed to the model as something it + said, every time the conversation was reopened, long after the cause was + fixed. + + On a reopened conversation it is sent once the browser has restored the + transcript, which would otherwise replace it (see :func:`_announce_live_run`). + """ + if resumed: + asyncio.create_task(_send_notice_after_resume(content)) + return + message = cl.Message(content=content) + # Chainlit stores a message on its first send unless told it already has. + message.persisted = True + await message.send() + + +async def _send_notice_after_resume(content: str) -> None: + try: + await asyncio.sleep(_RESUME_ANNOUNCE_DELAY_S) + await _send_notice(content) + except Exception: + logger.warning("Could not deliver a session notice", exc_info=True) + + async def _announce_live_run() -> None: """Put a reopened conversation back into its running state. @@ -1123,7 +1311,7 @@ def _history_from_thread(thread: dict) -> list[dict]: return history -async def _start_session(resumed_history: list[dict] | None = None) -> None: +async def _start_session(resumed_thread: dict | None = None) -> None: """Build everything one chat session needs. Shared by start and resume.""" # One correlation id per chat session, bound for the lifetime of this # handler so agent-init logs carry it. on_message rebinds it per message. @@ -1136,6 +1324,47 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: # is new or a reopened conversation. ACTIVITY.session_opened("chat") + # Who this session belongs to is settled once, here, and everything after + # uses that answer: preferences, run records, and every model call made on + # the user's behalf (the LLM proxy bills them by these ids). Cheap - header + # parsing only - so it runs before the agent is built, which a refused + # session never needs. + resumed = resumed_thread is not None + session_identity = _resolve_session_identity(thread_id) + if session_identity.conflict: + emit_event("identity_conflict", logger=logger, level=logging.WARNING, resumed=resumed) + if resumed_thread is not None: + _withhold_transcript(resumed_thread) + await _send_notice(f"⚠️ Sign-in changed: {session_identity.notice}.", resumed=resumed) + return + + resumed_history: list[dict] | None = None + if resumed_thread is not None: + resumed_history = _history_from_thread(resumed_thread) + # Only once the session is known to be its owner's: a run is never + # handed to a connection that is about to be refused. + reattached = await _follow_live_run(resumed_thread.get("id")) + emit_event( + "chat_resume", + steps=len(resumed_thread.get("steps") or []), + history_turns=len(resumed_history), + reattached=reattached, + ) + + identity = session_identity.identity + set_current_identity(identity) + cl.user_session.set("identity", identity) + if llm_proxy_settings() is not None: + problem = proxy_identity_problem(identity) + if problem: + # Said now rather than as an error on the first question: every + # model call in this session would be refused for the same reason. + logger.error("session cannot use the LLM proxy: %s", problem) + await _send_notice( + f"⚠️ **The assistant cannot answer in this session.** {problem[:1].upper()}{problem[1:]}", + resumed=resumed, + ) + # Determine agent type: CLI env var > chat profile selection > default AD1 if FORCE_AGENT in ("a1", "ad1"): agent_type = FORCE_AGENT @@ -1144,7 +1373,7 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: agent_type = "a1" if str(profile).upper() == "A1" else "ad1" label = "AD1" if agent_type == "ad1" else "A1" - emit_event("chat_start", agent_type=agent_type, llm=DEFAULT_LLM, resumed=resumed_history is not None) + emit_event("chat_start", agent_type=agent_type, llm=DEFAULT_LLM, resumed=resumed) try: if agent_type == "ad1": from biomni.agent.ad1 import AD1 @@ -1160,7 +1389,7 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: cl.user_session.set("thread_id", thread_id) except Exception as exc: logger.exception("Failed to initialize %s agent", label) - await cl.Message(content=f"Failed to initialize {label}: {exc}").send() + await _send_notice(f"Failed to initialize {label}: {exc}", resumed=resumed) return # Resolve identity, preferences, data scope and output directory, then tell @@ -1171,7 +1400,7 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: # than failing an otherwise usable one. ws = None try: - ws = await run_in_executor(_build_workspace_session, thread_id, _session_header_source()) + ws = await run_in_executor(_build_workspace_session, session_identity) cl.user_session.set("workspace", ws) _apply_scope_to_agent(agent, ws) emit_event( @@ -1180,6 +1409,7 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: has_workspace=bool(ws.workspace_root), scope_folders=len(ws.scope.roots), scope_is_default=ws.scope.is_default, + workspace_readable=ws.workspace_problem is None, output_source=ws.output.source, output_ephemeral=ws.output.is_ephemeral, prefs_persisted=not isinstance(ws.store, NullPrefsStore), @@ -1194,6 +1424,15 @@ async def _start_session(resumed_history: list[dict] | None = None) -> None: except Exception: logger.exception("Failed to resolve workspace session (continuing with defaults)") + if ws is not None and ws.workspace_problem: + # Said up front: every question about the user's own data would + # otherwise be answered as if the workspace were empty. + await _send_notice( + f"⚠️ **This app cannot read your workspace**, so none of your data is visible to it. " + f"{_as_sentence(ws.workspace_problem[:1].upper() + ws.workspace_problem[1:])}", + resumed=resumed, + ) + # Expose the scope / output controls. if ws is not None: try: @@ -1236,9 +1475,8 @@ def _session_header_source() -> object: """The gateway headers for this session, or None outside a request context. Chainlit keeps the connection's WSGI/ASGI environ on the session; that is - where an authentication gateway's forwarded headers arrive. Best-effort by - design - with no gateway (local dev, today's deployment) this yields None - and identity falls back to an anonymous, session-scoped key. + where an authentication gateway's forwarded headers arrive, when it forwards + them on the websocket upgrade as well as on plain requests. """ try: session = cl.context.session @@ -1247,6 +1485,23 @@ def _session_header_source() -> object: return getattr(session, "environ", None) or getattr(session, "http_headers", None) +def _resolve_session_identity(session_id: str) -> SessionIdentity: + """Who this chat session belongs to (see ``identity.resolve_session_identity``). + + Two witnesses, reconciled there: the login Chainlit authenticated the + connection with - which records what the gateway, or the demo password, + established at sign-in - and the gateway's headers on the connection itself. + Deriving it from the headers alone, as this app used to, cost every + password login its own preferences and made the gateway's headers + mandatory on websocket upgrades, which not every gateway forwards. + """ + try: + user = getattr(cl.context.session, "user", None) + except Exception: + user = None + return resolve_session_identity(_login_of(user), _session_header_source(), session_id=session_id) + + def _apply_scope_to_agent(agent, ws: WorkspaceSession) -> None: """Point the agent at the selected scope and output directory. @@ -1320,13 +1575,8 @@ async def on_settings_update(settings: dict): summary = "**Data scope cleared.** The agent will only look inside workspace folders a task points it to." summary += f"\n\nOutputs: `{ws.output.path}`" - if not ws.output.writable: - summary += f"\n\n⚠️ That directory is not writable. {ws.output.reason or ''}".rstrip() - else: - if ws.output.is_ephemeral: - summary += "\n\n⚠️ Container-local storage: results are lost when the application restarts." - if ws.output.reason: - summary += f"\n\n⚠️ This is not the configured location - {ws.output.reason}." + for warning in _output_warnings(ws.output): + summary += f"\n\n{warning}" if ws.prefs.remember and not persisted: # NullPrefsStore knows exactly why it cannot persist - no signed-in user @@ -1338,6 +1588,16 @@ async def on_settings_update(settings: dict): await cl.Message(content=summary).send() + # The dialog's help text describes the state it was built from - where + # results go and why, what cannot be read - so it is rebuilt from the + # settings just applied. Otherwise it goes on describing the session as it + # began: a folder the user has just chosen, and that cannot be used, would go + # unmentioned the next time they open it. + try: + await cl.ChatSettings(_workspace_settings_widgets(ws)).send() + except Exception: + logger.exception("Failed to refresh workspace settings (session remains usable)") + # --------------------------------------------------------------------------- # Message handler @@ -1423,8 +1683,15 @@ async def on_message(message: cl.Message): started = time.monotonic() # An agent run is the work /status exists to report: it holds the process # for minutes, and the monitoring framework must never see the pod as idle - # while one is in flight. - with bind_run(session_id=thread_id, run_id=run_id), ACTIVITY.track_job("agent_run"): + # while one is in flight. The session's identity rides along with the ids: + # each message is its own task, and every model call made for it - triage, + # planning, the agent's turns, a tool's own lookups - must name the user it + # is made for (see biomni.llm_proxy). + with ( + bind_run(session_id=thread_id, run_id=run_id), + bound_identity(cl.user_session.get("identity")), + ACTIVITY.track_job("agent_run"), + ): try: await _process_message(message) finally: @@ -1445,7 +1712,7 @@ async def on_stop(): happened. Cancelling the run makes its handler unwind normally: the record closes out and the page is told the task ended. """ - if LIVE_RUNS.cancel(_thread_key()): + if LIVE_RUNS.cancel(_thread_key(), _websocket_session()): emit_event("run_stopped_by_user") @@ -1522,12 +1789,22 @@ async def _process_message(message: cl.Message): step.output = "Using the full tool set." except Exception as exc: logger.warning("Tool retrieval failed; falling back to full tool set", exc_info=True) - step.output = f"⚠️ Tool retrieval failed ({exc}); proceeding with all tools." + step.output = f"⚠️ Could not narrow the tool set: {describe_llm_failure(exc)} Proceeding with all tools." # ------------------------------------------------------------------ # Phase 3: Interactive planning # ------------------------------------------------------------------ - prompt = await _interactive_planning(agent, prompt, agent_type=agent_type, selected_data_lake=selected_data_lake) + try: + prompt = await _interactive_planning( + agent, prompt, agent_type=agent_type, selected_data_lake=selected_data_lake + ) + except Exception as exc: + # Only raised when the model is out of reach altogether; the run could + # not get past its first step, so it is not started. + logger.warning("Planning could not reach the language model; not starting the run", exc_info=True) + emit_event("run_not_started", logger=logger, level=logging.WARNING, error_type=type(exc).__name__) + await cl.Message(content=f"⚠️ **Could not start this request.** {describe_llm_failure(exc)}").send() + return if prompt is None: # User cancelled await cl.Message(content="Execution cancelled.").send() @@ -1617,10 +1894,14 @@ async def _process_message(message: cl.Message): ) run_status, run_error = "completed", None except Exception as exc: - # Close the record out before the exception propagates, otherwise the - # run stays "running" until a later session reconciles it as stale. + # Recorded as failed here - otherwise the run stays "running" until a + # later session reconciles it as stale - and told to the user in words + # rather than as Chainlit's raw exception text. run_status, run_error = "failed", str(exc) - raise + logger.exception("Agent run failed") + emit_event("run_failed", logger=logger, level=logging.WARNING, error_type=type(exc).__name__) + await cl.Message(content=f"⚠️ **The run stopped before it finished.** {describe_llm_failure(exc)}").send() + return finally: if record is not None: try: @@ -1715,89 +1996,98 @@ async def _stream_execution( # can tell a slow-but-alive run from one whose process died. on_tick=on_heartbeat, ): - async for state in stream_langgraph(agent.app, inputs, config): - final_state = state - message = state["messages"][-1] - content = message.content if isinstance(message.content, str) else "" - - if not content or content == prompt: - continue - - # ------------------------------------------------------------------ - # Parse XML tags from the agent's raw output - # ------------------------------------------------------------------ - - # Locate first structural tag to separate reasoning prefix - tag_positions = [] - for tag in ["", "", ""]: - pos = content.find(tag) - if pos != -1: - tag_positions.append(pos) - - # 1. Reasoning / thinking (text before the first tag) - if tag_positions: - first_tag = min(tag_positions) - thinking = content[:first_tag].strip() - if thinking: - async with cl.Step(name="🤔 Thinking", type="llm", show_input=False) as step: - step.output = thinking - - # 2. Solution (final answer) - solution_match = re.search(r"(.*?)", content, re.DOTALL) - if solution_match and not solution_found: - solution_found = True - solution_text = solution_match.group(1).strip() - await cl.Message(content=solution_text).send() - - # 3. Code execution block - execute_match = re.search(r"(.*?)", content, re.DOTALL) - if execute_match: - code = execute_match.group(1).strip() - language = "python" - if code.startswith("#!R"): - language = "r" - code = re.sub(r"^#!R\s*", "", code, count=1) - elif code.startswith("#!BASH") or code.startswith("#!CLI"): - language = "bash" - code = re.sub(r"^#!(BASH|CLI)\s*", "", code, count=1) - - code_step = cl.Step(name=f"⚡ Executing {language.upper()}", type="run", show_input=False) - await code_step.__aenter__() - code_step.output = f"```{language}\n{code}\n```" - code_steps.append(code_step) - # Tick an elapsed-time line while the (blocking) code runs so the - # step doesn't look frozen; cancelled when the observation lands. - code_timers.append(asyncio.create_task(_tick_step_timer(code_step, language, code, time.monotonic()))) - # Do NOT exit the step yet; we close it when the observation arrives - - # 4. Observation (result of code execution) - obs_match = re.search(r"(.*?)", content, re.DOTALL) - if obs_match: - observation = obs_match.group(1).strip() - - # Close the pending code step now that we have a result - if code_timers: - await _cancel_task(code_timers.pop()) - if code_steps: - finished_step = code_steps.pop() - await finished_step.__aexit__(None, None, None) - - async with cl.Step(name="👁 Observation", type="tool", show_input=False) as obs_step: - # Truncate very long output for display - display_obs = observation[:3000] + "\n...[truncated]" if len(observation) > 3000 else observation - obs_step.output = display_obs - - # Display any generated images mentioned in the observation - await _display_images(observation) - - # Close any code steps/timers that never received an observation (edge case) - for timer in code_timers: - await _cancel_task(timer) - for step in code_steps: - await step.__aexit__(None, None, None) - - # If no tag was found, surface the last message content - if not solution_found and final_state: + try: + async for state in stream_langgraph(agent.app, inputs, config): + final_state = state + message = state["messages"][-1] + content = message.content if isinstance(message.content, str) else "" + + if not content or content == prompt: + continue + + # ------------------------------------------------------------------ + # Parse XML tags from the agent's raw output + # ------------------------------------------------------------------ + + # Locate first structural tag to separate reasoning prefix + tag_positions = [] + for tag in ["", "", ""]: + pos = content.find(tag) + if pos != -1: + tag_positions.append(pos) + + # 1. Reasoning / thinking (text before the first tag) + if tag_positions: + first_tag = min(tag_positions) + thinking = content[:first_tag].strip() + if thinking: + async with cl.Step(name="🤔 Thinking", type="llm", show_input=False) as step: + step.output = thinking + + # 2. Solution (final answer) + solution_match = re.search(r"(.*?)", content, re.DOTALL) + if solution_match and not solution_found: + solution_found = True + solution_text = solution_match.group(1).strip() + await cl.Message(content=solution_text).send() + + # 3. Code execution block + execute_match = re.search(r"(.*?)", content, re.DOTALL) + if execute_match: + code = execute_match.group(1).strip() + language = "python" + if code.startswith("#!R"): + language = "r" + code = re.sub(r"^#!R\s*", "", code, count=1) + elif code.startswith("#!BASH") or code.startswith("#!CLI"): + language = "bash" + code = re.sub(r"^#!(BASH|CLI)\s*", "", code, count=1) + + code_step = cl.Step(name=f"⚡ Executing {language.upper()}", type="run", show_input=False) + await code_step.__aenter__() + code_step.output = f"```{language}\n{code}\n```" + code_steps.append(code_step) + # Tick an elapsed-time line while the (blocking) code runs so the + # step doesn't look frozen; cancelled when the observation lands. + code_timers.append( + asyncio.create_task(_tick_step_timer(code_step, language, code, time.monotonic())) + ) + # Do NOT exit the step yet; we close it when the observation arrives + + # 4. Observation (result of code execution) + obs_match = re.search(r"(.*?)", content, re.DOTALL) + if obs_match: + observation = obs_match.group(1).strip() + + # Close the pending code step now that we have a result + if code_timers: + await _cancel_task(code_timers.pop()) + if code_steps: + finished_step = code_steps.pop() + await finished_step.__aexit__(None, None, None) + + async with cl.Step(name="👁 Observation", type="tool", show_input=False) as obs_step: + # Truncate very long output for display + display_obs = ( + observation[:3000] + "\n...[truncated]" if len(observation) > 3000 else observation + ) + obs_step.output = display_obs + + # Display any generated images mentioned in the observation + await _display_images(observation) + finally: + # Close any code steps/timers that never received an observation - + # including those of a run that failed or was stopped mid-step, which + # would otherwise keep spinning in the transcript. + for timer in code_timers: + await _cancel_task(timer) + for step in code_steps: + await step.__aexit__(None, None, None) + + # If no tag was found, surface the last message content - the + # agent's, never the user's own question, which is what the last message is + # when the model produced nothing at all. + if not solution_found and final_state and not isinstance(final_state["messages"][-1], HumanMessage): last_content = final_state["messages"][-1].content if isinstance(last_content, str): cleaned = re.sub(r".*?", "", last_content, flags=re.DOTALL) diff --git a/chainlit_ui/graph_stream.py b/chainlit_ui/graph_stream.py new file mode 100644 index 000000000..6d94a82e0 --- /dev/null +++ b/chainlit_ui/graph_stream.py @@ -0,0 +1,102 @@ +"""Reading a LangGraph run from the event loop while it executes in a thread. + +The agent's graph is synchronous and runs for minutes, so it gets a worker +thread of its own and hands each state back to the event loop through a queue. +Two things have to cross that boundary besides the states, and both used to be +lost on the way: the order to stop, and the exception a failed run raised. The +states themselves have to cross as they were, not as they later became. + +No Chainlit import, so the stopping and failure behaviour is testable without a +running app (CI installs no chainlit extra). +""" + +from __future__ import annotations + +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from biomni.observability import capture_context + +if TYPE_CHECKING: + from collections.abc import AsyncIterator + + +def _snapshot(state: Any) -> Any: + """The state as it is at this moment. + + The agent's nodes append to one shared ``messages`` list and return it, so + every state the graph streams is the same list. Read when the event loop got + round to it, a state already held the next node's messages: a code step went + missing from the transcript and its observation was shown twice. A shallow + copy is enough, since messages are appended, never changed. + """ + if isinstance(state, dict) and isinstance(state.get("messages"), list): + return {**state, "messages": list(state["messages"])} + return state + + +@dataclass(frozen=True) +class _GraphFailed: + """What the graph raised, on its way from the worker thread to the reader.""" + + error: Exception + + +async def stream_langgraph(agent_app: Any, inputs: Any, config: Any) -> AsyncIterator[Any]: + """Yield LangGraph state dicts asynchronously from a sync stream. + + The graph runs in a worker thread, which cancellation cannot reach: an + asyncio task that is cancelled stops *reading*, and a thread that nobody + reads goes right on calling the model. Stop therefore used to end the + conversation while the run it stopped kept working - for as long as the + whole plan took. The flag below is the cooperative half of the stop: the + producer checks it between graph nodes, so the run ends after the step it + was already in. + + An exception the graph raises is re-raised here, in the reader. + """ + queue: asyncio.Queue = asyncio.Queue() + loop = asyncio.get_running_loop() + stop = threading.Event() + + def _producer(): + try: + for state in agent_app.stream(inputs, stream_mode="values", config=config): + asyncio.run_coroutine_threadsafe(queue.put(_snapshot(state)), loop) + if stop.is_set(): + break + except Exception as exc: + # Handed to the reader, not left on the executor's future where no + # one would ever look: a graph that died - a model call refused, a + # node that raised - used to read as a graph that finished, and the + # user was told "Run complete" over a run that never produced a word. + asyncio.run_coroutine_threadsafe(queue.put(_GraphFailed(exc)), loop) + finally: + asyncio.run_coroutine_threadsafe(queue.put(None), loop) # sentinel + + executor = ThreadPoolExecutor(max_workers=1) + # Capture the active context here (caller thread) and replay it in the + # producer thread, so the graph's logs carry the session_id/run_id and its + # model calls the session's identity (the LLM proxy bills each call to it). + ctx = capture_context() + executor.submit(ctx.run, _producer) + + try: + while True: + state = await queue.get() + if state is None: + break + if isinstance(state, _GraphFailed): + raise state.error + yield state + finally: + stop.set() + # Not waited on: the worker may be mid-node, and holding the event loop + # for it would freeze the UI. It exits at the next check. + executor.shutdown(wait=False) + + +__all__ = ["stream_langgraph"] diff --git a/chainlit_ui/live_runs.py b/chainlit_ui/live_runs.py index d7b83e454..6739fd65c 100644 --- a/chainlit_ui/live_runs.py +++ b/chainlit_ui/live_runs.py @@ -24,8 +24,13 @@ (``restore_existing_session``), which is why two plain attribute assignments are enough - ``ChainlitEmitter.emit`` reads ``session.emit`` on every call. -Sessions are held by duck type (anything with ``emit``/``emit_call``), so this -module stays importable without Chainlit and the bookkeeping stays testable. +A run is only ever handed to - or stopped from - a connection signed in as the +user who started it. Chainlit already lets nobody but a conversation's author +reopen it, so this is the second line of that defence, not the first. + +Sessions are held by duck type (anything with ``emit``/``emit_call``, and the +``user`` Chainlit signed it in as), so this module stays importable without +Chainlit and the bookkeeping stays testable. """ from __future__ import annotations @@ -37,6 +42,11 @@ logger = logging.getLogger(__name__) +def _owner(session: Any) -> str | None: + """Who a session is signed in as: the identifier of its Chainlit login.""" + return getattr(getattr(session, "user", None), "identifier", None) + + @dataclass class LiveRun: """An executing conversation: where its output goes, and what to cancel.""" @@ -97,6 +107,9 @@ def reattach(self, thread_id: str | None, session: Any) -> bool: live = self._runs.get(thread_id) if live is None or live.session is session: return False + if _owner(live.session) != _owner(session): + logger.warning("refused to hand the run for thread %s to a connection of another user", thread_id) + return False try: live.session.emit = session.emit live.session.emit_call = session.emit_call @@ -106,19 +119,23 @@ def reattach(self, thread_id: str | None, session: Any) -> bool: logger.info("reattached a live run to the reopened conversation %s", thread_id) return True - def cancel(self, thread_id: str | None) -> bool: + def cancel(self, thread_id: str | None, session: Any) -> bool: """Stop the run executing in ``thread_id``. True if one was cancelled. Chainlit's Stop cancels the *clicking* session's task, which for a reopened conversation is not the task doing the work. This cancels the run itself, and its own cleanup then closes the record out and tells the - page the task ended. + page the task ended. ``session`` is the one the click arrived on; only + the user who started the run can stop it. """ if not thread_id: return False live = self._runs.get(thread_id) task = live.task if live else None - if task is None or task.done(): + if live is None or task is None or task.done(): + return False + if _owner(live.session) != _owner(session): + logger.warning("refused a stop for thread %s from a connection of another user", thread_id) return False task.cancel() logger.info("cancelled the run for conversation %s", thread_id) diff --git a/chainlit_ui/llm_failures.py b/chainlit_ui/llm_failures.py new file mode 100644 index 000000000..bfb131711 --- /dev/null +++ b/chainlit_ui/llm_failures.py @@ -0,0 +1,130 @@ +"""Say why a model call failed, in words a user and an operator can act on. + +The provider SDKs raise exceptions whose text is a status code and a raw JSON +body - ``Error code: 400 - {'error': {'message': ...}}`` - and that is what the +UI used to show. Behind the platform's LLM proxy the causes that matter are few +and specific (a missing identity, a revoked token, a usage limit), so they get +sentences of their own; everything else still says what was raised. + +Pure Python: duck-typed on the SDKs' exception shape (``status_code``, ``body``) +rather than importing them, so it is testable without either installed. +""" + +from __future__ import annotations + +from biomni.llm_proxy import LLMProxyIdentityError, llm_proxy_settings + +# Failures after which trying again in the same run is pointless: the next call +# would meet the same refusal. The pipeline stops on these instead of carrying +# on "without a plan" into an execution that can only fail the same way. +_ACCESS_STATUSES = frozenset({401, 403, 404, 429}) +_CONNECTION_ERRORS = frozenset({"APIConnectionError", "APITimeoutError", "ConnectError", "ConnectTimeout"}) + +_MAX_DETAIL = 300 + + +def _status(exc: BaseException) -> int | None: + status = getattr(exc, "status_code", None) + return status if isinstance(status, int) else None + + +def _is_connection_error(exc: BaseException) -> bool: + return any(cls.__name__ in _CONNECTION_ERRORS for cls in type(exc).__mro__) + + +def _error_body(exc: BaseException) -> object: + """The error response's JSON body, however the SDK chose to keep it. + + Read from the response itself where there is one: the OpenAI SDK keeps only + the ``error`` member as ``body`` - a bare string when the proxy sent + ``{"error": "..."}`` - and a response that was not JSON at all as the raw + text, and the two cannot be told apart from ``body`` alone. + """ + read_json = getattr(getattr(exc, "response", None), "json", None) + if callable(read_json): + try: + return read_json() + except Exception: + return None # not JSON: a gateway's HTML error page, say + return getattr(exc, "body", None) + + +def _server_message(exc: BaseException) -> str: + """The message inside the error body, or an empty string. + + Reads the shapes providers and the proxies in front of them send: + ``{"error": {"message": ...}}`` (OpenAI, Anthropic), ``{"error": "..."}``, + ``{"message": ...}`` and FastAPI's ``{"detail": "..."}``. Anything else says + nothing a user can act on and is left out: a body that is not JSON, and the + SDK's own ``message``, which is only the status code followed by the raw + body. + """ + body = _error_body(exc) + if not isinstance(body, dict): + return "" + error = body.get("error") + candidates = (error.get("message") if isinstance(error, dict) else error, body.get("message"), body.get("detail")) + text = next((c.strip() for c in candidates if isinstance(c, str) and c.strip()), "") + return text if len(text) <= _MAX_DETAIL else text[: _MAX_DETAIL - 1] + "…" + + +def _sentence(text: str) -> str: + text = text.strip() + if not text: + return text + text = text[:1].upper() + text[1:] + return text if text.endswith((".", "!", "?")) else text + "." + + +def _behind_proxy() -> bool: + try: + return llm_proxy_settings() is not None + except ValueError: + return True + + +def is_access_failure(exc: BaseException) -> bool: + """Whether the model is unreachable for this session, not just this call.""" + if isinstance(exc, LLMProxyIdentityError): + return True + status = _status(exc) + if status is not None: + return status in _ACCESS_STATUSES or status >= 500 + return _is_connection_error(exc) + + +def describe_llm_failure(exc: BaseException) -> str: + """One or two sentences explaining ``exc``, naming the fix when there is one.""" + if isinstance(exc, LLMProxyIdentityError): + return _sentence(str(exc)) + + proxy = _behind_proxy() + service = "The platform's LLM proxy" if proxy else "The language model service" + status = _status(exc) + detail = _server_message(exc) + said = f": {detail}" if detail else "" + + if status == 401: + fix = "Check BIOMNI_LLM_PROXY_API_KEY." if proxy else "Check the deployment's API key." + return f"{service} rejected this deployment's credentials (HTTP 401). {fix}" + if status == 403: + return _sentence(f"{service} refused the request (HTTP 403){said}") + if status == 404: + return _sentence(f"{service} does not offer the configured model (HTTP 404{said}). Check BIOMNI_LLM") + if status == 429: + scope = " The platform limits model usage per user and per workspace." if proxy else "" + return f"The language model usage limit has been reached (HTTP 429).{scope} Try again later." + if status is not None and status >= 500: + return _sentence(f"{service} failed to answer (HTTP {status}){said}. Try again shortly") + if status is not None: + return _sentence(f"{service} rejected the request (HTTP {status}){said}") + if _is_connection_error(exc): + return _sentence(f"Could not reach {service[0].lower()}{service[1:]}. Try again shortly") + + text = str(exc).strip() or type(exc).__name__ + if len(text) > _MAX_DETAIL: + text = text[: _MAX_DETAIL - 1] + "…" + return _sentence(f"{type(exc).__name__}: {text}" if text != type(exc).__name__ else text) + + +__all__ = ["describe_llm_failure", "is_access_failure"] diff --git a/chainlit_ui/planning.py b/chainlit_ui/planning.py index b4699a64a..036183d1c 100644 --- a/chainlit_ui/planning.py +++ b/chainlit_ui/planning.py @@ -28,6 +28,8 @@ from biomni.observability import emit_event from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from chainlit_ui.llm_failures import describe_llm_failure, is_access_failure + if TYPE_CHECKING: from biomni.agent.a1 import A1 @@ -39,7 +41,7 @@ "Given the user's research question, write a concise numbered plan " "of 3 to 7 steps describing exactly how you will solve it. " "IMPORTANT: Always prefer this tool order in your plan: " - "(1) web/literature search tools first (advanced_web_search, search_pubmed, search_biorxiv), " + "(1) web/literature search tools first (search_google, query_pubmed, query_arxiv, query_scholar), " "(2) local data and database query tools second, " "(3) custom code generation only as a last resort. " "Mention specific tools, databases, or analyses you will use. " @@ -163,8 +165,10 @@ async def quick_answer(agent: A1, prompt: str, history: list[dict] | None = None messages.append(HumanMessage(content=prompt)) try: - loop = asyncio.get_event_loop() - response = await loop.run_in_executor(None, agent.llm.invoke, messages) + # to_thread, not a bare executor: it carries the caller's context into + # the worker, and with it the session's identity, which a model call + # through the platform's LLM proxy cannot be made without. + response = await asyncio.to_thread(agent.llm.invoke, messages) except Exception: logger.warning("Triage call failed; falling back to the planning path", exc_info=True) return None @@ -372,7 +376,9 @@ async def interactive_planning( Returns the (possibly modified) prompt on approval, or `None` if the user cancels. If plan generation itself raises, the approval gate is skipped - and the original prompt is returned so execution still happens. + and the original prompt is returned so execution still happens - unless the + model is out of reach altogether (`is_access_failure`: refused credentials, + a usage limit, no identity for the LLM proxy), which is raised to the caller. """ base_prompt = build_planning_system_prompt(agent, agent_type, selected_data_lake) modification_context = "" @@ -389,21 +395,28 @@ async def interactive_planning( # default_open: the plan is the thing the user is being asked to approve, # so it must be readable without first expanding a collapsed step. + unreachable: Exception | None = None async with cl.Step(name="📋 Generating Research Plan", type="llm", show_input=False, default_open=True) as step: try: - # NB: asyncio.to_thread would also work but copies the caller's - # contextvars into the worker - the rest of chainlit_app.py - # uses bare run_in_executor and we match that semantics so - # LangChain callback/tracing contextvars don't silently change - # which trace the LLM call attaches to. - loop = asyncio.get_event_loop() - response = await loop.run_in_executor(None, agent.llm.invoke, planning_messages) + # to_thread carries the caller's context into the worker: the + # session's identity (the LLM proxy bills the call to it) and + # the correlation ids its log lines need. + response = await asyncio.to_thread(agent.llm.invoke, planning_messages) plan_text = response.content if hasattr(response, "content") else str(response) step.output = plan_text except Exception as exc: - logger.warning("Plan generation failed; proceeding without approval gate", exc_info=True) - step.output = f"⚠️ Could not generate plan ({exc}). Proceeding without a plan." - return prompt + if not is_access_failure(exc): + logger.warning("Plan generation failed; proceeding without approval gate", exc_info=True) + step.output = f"⚠️ Could not generate a plan: {describe_llm_failure(exc)} Proceeding without one." + return prompt + # Out of reach for the whole run: "proceeding without a plan" + # would only fail again at the first step. Raised once the step + # has closed, since raising inside it replaces this text with + # the exception's raw repr. + step.output = "⚠️ Could not reach the language model." + unreachable = exc + if unreachable is not None: + raise unreachable # The plan already carries its own data-files section and is shown # expanded, so repeating the list here only duplicated it on screen. diff --git a/chainlit_ui/workspace_panel.py b/chainlit_ui/workspace_panel.py index 68178af15..422fe6651 100644 --- a/chainlit_ui/workspace_panel.py +++ b/chainlit_ui/workspace_panel.py @@ -238,6 +238,15 @@ def build_scope_inventory( """ if not workspace_root: return "" + if not os.access(workspace_root, os.R_OK | os.X_OK): + # Otherwise the agent is told the workspace has no folders and to look + # for itself, and a user whose data is all there hears that it is empty. + return ( + f"User workspace ({workspace_root})\n\n" + "This application is NOT permitted to read the workspace, so what it contains is unknown. " + "If a task needs the user's data, tell them it cannot be accessed from this app until the " + "deployment's storage permissions are fixed. Do not guess at what the workspace contains." + ) if scope.is_default or not scope.roots: names = list(top_level_dirs if top_level_dirs is not None else list_top_level_dirs(workspace_root)) diff --git a/deploy/k8s/biomni-ad.yaml b/deploy/k8s/biomni-ad.yaml index 22cdbfd24..73b1ec9fb 100644 --- a/deploy/k8s/biomni-ad.yaml +++ b/deploy/k8s/biomni-ad.yaml @@ -2,11 +2,16 @@ # # This is a starting point, not a turnkey production config. Review resource # limits, replica count, ingress, and the Secret/ConfigMap wiring for your -# cluster before applying. The pieces that matter for *observability* - the -# focus of this manifest - are the three probes and the log-shaping env vars. +# cluster before applying. # # kubectl apply -f deploy/k8s/biomni-ad.yaml # +# The volumes and paths are the ones the AD Workbench (GRIP) deployment uses. +# What it does not do is trust an authentication gateway or use the platform's +# LLM proxy, since it exposes the app through a plain Service with no gateway +# in front. docs/grip_deployment.md#kubernetes-manifest lists the changes that +# turn it into the GRIP deployment. +# # Logs: the container writes one JSON object per line to stdout. On AKS with # Container Insights enabled, those land in the Log Analytics workspace and are # queryable with KQL, e.g.: @@ -23,8 +28,11 @@ metadata: name: biomni-ad-secrets type: Opaque stringData: - # Provide whichever provider your BIOMNI_LLM resolves to. Readiness (/readyz) - # fails closed until at least one credential is present. + # Model access: the platform LLM proxy's token on GRIP (with + # BIOMNI_LLM_PROXY_URL below; see docs/grip_deployment.md), otherwise the key + # of whichever provider BIOMNI_LLM resolves to. Readiness (/readyz) fails + # closed until the one in use is present. + # BIOMNI_LLM_PROXY_API_KEY: "replace-me" ANTHROPIC_API_KEY: "replace-me" # OPENAI_API_KEY: "replace-me" # AZURE_OPENAI_API_KEY: "replace-me" @@ -37,11 +45,9 @@ stringData: # openssl rand -base64 48 # CHAINLIT_AUTH_SECRET: "replace-me" --- -# Persistent storage for run outputs, user preferences and run records. -# -# Without this the app falls back to a container-local ./runs directory, and a -# rollout or restart destroys every artifact the user produced. The UI warns -# when it detects that fallback. +# The app's own state: per-user preferences, run records, the chat-history +# database, and the reference datasets (BIOMNI_DATA_LAKE_PATH below), about 16.5 GB +# if every one is used, each downloaded the first time a question needs it. # # ReadWriteOnce is sufficient at replicas: 1 (as configured below). Scaling out # requires ReadWriteMany (e.g. an Azure Files share) because several pods would @@ -55,9 +61,52 @@ spec: - ReadWriteOnce resources: requests: - storage: 20Gi + storage: 30Gi # storageClassName: azurefile-csi # switch to a RWX class before scaling out --- +# Run outputs: a folder per run with its figures, tables and code. +# +# Without a volume here the app falls back to a container-local ./runs, and a +# rollout or restart destroys every result the user produced. The UI warns when +# it detects that fallback. +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: biomni-ad-runs +spec: + accessModes: + - ReadWriteOnce + resources: + requests: + storage: 50Gi +--- +# The user's workspace: their own input data, which the agent reads. +# +# Azure Files, so the same data is reachable from wherever else the user works +# with it. AKS's built-in azurefile-csi class mounts shares world-writable +# (dir_mode=0777), which the app can use as is. A share with tighter +# permissions - an existing one bound through a PersistentVolume, as on GRIP - +# needs ownership options on that PersistentVolume, because SMB ignores fsGroup +# (see the note on securityContext below): +# +# mountOptions: +# - uid=57439 +# - gid=57439 +# - dir_mode=0770 +# - file_mode=0770 +# - mfsymlinks +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: biomni-ad-workspace +spec: + accessModes: + - ReadWriteMany + storageClassName: azurefile-csi + resources: + requests: + storage: 100Gi +--- apiVersion: apps/v1 kind: Deployment metadata: @@ -76,27 +125,31 @@ spec: spec: # The image runs as the non-root mambauser (UID 57439). A freshly # provisioned block PVC mounts as root:root 0755, so without fsGroup the - # container cannot write /data - and because every write path here - # degrades gracefully, the failure is silent: outputs fall back to the - # ephemeral ./runs and preferences stop persisting, making the volume a - # no-op. fsGroup makes the kubelet chgrp the volume to that GID. + # container cannot write /app/state or /app/runs - and because every write + # path here degrades gracefully, the failure is silent: outputs fall back + # to the ephemeral ./runs and preferences stop persisting, making the + # volume a no-op. fsGroup makes the kubelet chgrp the volume to that GID. # # IMPORTANT: fsGroup only works on *block* volumes (Azure Disk, EBS, any # ReadWriteOnce PVC the kubelet formats). It is silently ignored by SMB - # and NFS, so a deployment putting BIOMNI_OUTPUT_ROOT or BIOMNI_STATE_DIR - # on Azure Files must instead set ownership in the StorageClass or PV: + # and NFS, so a share - the workspace, or BIOMNI_OUTPUT_ROOT or + # BIOMNI_STATE_DIR put on Azure Files - must instead get its ownership + # from the mountOptions on its StorageClass or PV (see the workspace claim + # above). Without them the share mounts as root:root 0755, os.access() + # reports it unwritable to UID 57439, and outputs fall back exactly as if + # the env var had never been set. resolve_output_dir() logs a WARNING + # naming the rejected path when this happens. # - # mountOptions: - # - uid=57439 - # - gid=57439 - # - dir_mode=0770 - # - file_mode=0770 - # - mfsymlinks + # Changed options only apply when the node mounts the share afresh. A + # StorageClass's options are copied into the PV when it is provisioned, so + # edit the PV's spec.mountOptions too; and the node keeps the share + # mounted while any pod there uses it, so scale this Deployment to 0 and + # back rather than restarting it. To see what the container actually got: # - # Without it the share mounts as root:root 0755, os.access() reports it - # unwritable to UID 57439, and outputs fall back exactly as if the env var - # had never been set. resolve_output_dir() logs a WARNING naming the - # rejected path when this happens. + # kubectl exec deploy/biomni-ad -- sh -c 'id; grep " /app/user-data " /proc/self/mounts; ls -ldn /app/user-data' + # + # The app logs the same facts at startup, one `storage_check` line per + # location, with the reason and the fix for any it cannot use. securityContext: fsGroup: 57439 fsGroupChangePolicy: OnRootMismatch @@ -149,17 +202,28 @@ spec: # BIOMNI_TRUST_AUTH_HEADERS below is the mechanism; a password would # add a second, weaker door in front of a real one. + # The user's workspace, mounted below. /readyz reports not ready + # until it is there. + - name: BIOMNI_USER_DATA_PATH + value: "/app/user-data" + # Run artifacts. Without this the app writes to a container-local # ./runs and every result is lost on restart. Points at the PVC # mounted below; users can override per session in the UI. - name: BIOMNI_OUTPUT_ROOT - value: "/data/outputs" + value: "/app/runs" # Per-user preferences (data scope, output directory), durable run # records and the chat-history database, so a user who closes the # tab still finds their jobs and their conversations. - name: BIOMNI_STATE_DIR - value: "/data/state" + value: "/app/state" + + # Reference datasets from the public data lake, fetched on first use. + # The default is inside the container, which re-downloads them after + # every restart and counts them against the node's ephemeral storage. + - name: BIOMNI_DATA_LAKE_PATH + value: "/app/state/data_lake" # Chat history is stored in SQLite under $BIOMNI_STATE_DIR by # default, which assumes one replica - the same assumption the @@ -183,7 +247,10 @@ spec: # (Ai-App-User-Context, carrying sub/email/given_name/family_name, # and Ai-App-Workspace-Context, carrying uuid), plus the usual # single-value headers for other gateways. Set these only when the - # gateway in front uses different names. + # gateway in front uses different names. GRIP needs none of them: + # the composite headers are parsed field by field, so pointing a + # single-value setting (USER_ID, EMAIL, WORKSPACE) at one of them + # has no effect, and the app logs a warning saying so. # - name: BIOMNI_AUTH_USER_CONTEXT_HEADER # value: "ai-app-user-context" # - name: BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER @@ -217,13 +284,30 @@ spec: # survive). Never applied as shorter than the window above. # - name: BIOMNI_STATUS_OPEN_SESSION_INACTIVITY_SECONDS # value: "43200" + # The platform LLM proxy (GRIP, in front of Azure AI Foundry). With + # it set every model call goes through the proxy, authenticated by + # BIOMNI_LLM_PROXY_API_KEY from the Secret above and naming the + # session's user and workspace in X-User-Id / X-Workspace-Id. Those + # ids come from the gateway's headers, so this needs + # BIOMNI_TRUST_AUTH_HEADERS on too: without it every call is refused. + # BIOMNI_LLM is the name of the model the proxy serves. + # - name: BIOMNI_LLM_PROXY_URL + # value: "https:///v1" + # - name: BIOMNI_LLM_PROXY_SCHEMA + # value: "openai" # or "anthropic", for /v1/messages + # - name: BIOMNI_LLM + # value: "" envFrom: - secretRef: name: biomni-ad-secrets volumeMounts: - name: state - mountPath: /data + mountPath: /app/state + - name: runs + mountPath: /app/runs + - name: workspace + mountPath: /app/user-data # Startup probe: gates the other two until the heavy agent stack has # finished importing (~30-60s cold start). failureThreshold*period = @@ -246,8 +330,9 @@ spec: timeoutSeconds: 3 failureThreshold: 3 - # Readiness: pull the pod from the Service if it can't serve - data - # dir unmounted or no LLM credential. Returns 503 when not ready. + # Readiness: pull the pod from the Service if it can't serve - the + # workspace not mounted, or no way to reach a model. Returns 503, with + # the failing check named in the body, when not ready. readinessProbe: httpGet: path: /readyz @@ -256,16 +341,21 @@ spec: timeoutSeconds: 3 failureThreshold: 3 + # The documented minimum as requests and the recommendation as limits + # (docs/grip_deployment.md#resources): a full AD1 run peaks around + # 1 GB, but large reference datasets take several GB once loaded + # (BindingDB about 13 GiB, if loaded whole), and each open chat + # keeps what it loaded in memory. resources: requests: - cpu: "1" - memory: 2Gi + cpu: "2" + memory: 8Gi limits: # Agent-generated code runs in-process; cap memory so a runaway # snippet is OOMKilled (visible in kube events) instead of taking # the node down. - cpu: "2" - memory: 6Gi + cpu: "4" + memory: 16Gi # SIGTERM is sent on rollout/scale-down; give in-flight agent runs a # chance to drain before SIGKILL. Runs that do not finish in time are @@ -277,6 +367,12 @@ spec: - name: state persistentVolumeClaim: claimName: biomni-ad-state + - name: runs + persistentVolumeClaim: + claimName: biomni-ad-runs + - name: workspace + persistentVolumeClaim: + claimName: biomni-ad-workspace --- apiVersion: v1 kind: Service diff --git a/docs/configuration.md b/docs/configuration.md index 1642830ed..0d18241b5 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -109,9 +109,22 @@ ANTHROPIC_BASE_URL=https://api.anthropic.com OPENAI_BASE_URL=https://api.openai.com/v1 # Biomni Settings -BIOMNI_PATH=/path/to/data # Default: ./data +BIOMNI_PATH=/path/to/data # Default: ~/.biomni/data (the app's own data directory; + # /app/data in the container image) +BIOMNI_DATA_LAKE_PATH=/data/data_lake # Default: data/biomni_data/data_lake in the checkout + # (/app/data/biomni_data/data_lake in the image). + # Reference datasets, ~16.5 GB if all are used, + # each downloaded the first time a question needs it. +BIOMNI_USER_DATA_PATH=/app/user-data # Default: unset. The user's workspace: their own input + # data, shown in Settings and read by the agent. + # /readyz waits for it when set. +BIOMNI_DATA_PATH=/app/user-data # Default: unset. Older name for the workspace, read + # when BIOMNI_USER_DATA_PATH is not set +BIOMNI_USER_DATA_HOST_PATH=/mnt/data # Default: unset. Compose only: the host folder mounted + # at /app/user-data (also read as a workspace + # when the path exists inside the container) BIOMNI_TIMEOUT_SECONDS=1200 # Default: 600 (per code/tool step) -BIOMNI_LLM=model_name # Default: claude-sonnet-4-20250514 +BIOMNI_LLM=model_name # Default: claude-sonnet-4-5 (BIOMNI_LLM_MODEL is an alias) BIOMNI_TEMPERATURE=0.7 # Default: 0.7 BIOMNI_USE_TOOL_RETRIEVER=true # Default: true BIOMNI_AUTO_NETWORK_LIMITED_MODE=true # Default: true @@ -119,11 +132,34 @@ LLM_SOURCE=Anthropic # Preferred source selector BIOMNI_SOURCE=Anthropic # Also supported (backward compatibility) BIOMNI_CUSTOM_BASE_URL=http://localhost:8000/v1 BIOMNI_CUSTOM_API_KEY=custom_key +BIOMNI_COMMERCIAL_MODE=true # Default: false (leave out datasets and know-how whose + # licence does not allow commercial use) + +# Chat interface +BIOMNI_AGENT=ad1 # Default: unset (the user picks AD1 or A1 as the chat + # profile; set a1 or ad1 to fix the agent) +BIOMNI_ALWAYS_PLAN=true # Default: false (short questions are answered directly; + # true sends every question through a plan the + # user approves) +CHAINLIT_HOST=0.0.0.0 # Default: 0.0.0.0 (container entrypoint only) +CHAINLIT_PORT=8000 # Default: 8000 (container entrypoint only) + +# Tokens for individual data services, used by the tools that query them +PROTOCOLS_IO_ACCESS_TOKEN=token # protocols.io (BIOMNI_PROTOCOLS_IO_ACCESS_TOKEN also read) +SYNAPSE_AUTH_TOKEN=token # Synapse + +# Platform LLM proxy (GRIP; see docs/grip_deployment.md). When set, every model call +# goes through it, whatever the provider settings above say, and each request names +# the session's user and workspace in X-User-Id and X-Workspace-Id. +BIOMNI_LLM_PROXY_URL=https://proxy/v1 # Default: unset (model calls go direct). With or without /v1 +BIOMNI_LLM_PROXY_API_KEY=proxy_token # The proxy's token, sent as Authorization: Bearer +BIOMNI_LLM_PROXY_SCHEMA=openai # Default: openai (/v1/chat/completions); or anthropic (/v1/messages) # LLM resilience BIOMNI_LLM_MAX_RETRIES=3 # Default: 3 (provider-SDK 429/5xx backoff) BIOMNI_LLM_REQUEST_TIMEOUT=120 # Default: 120 (seconds per LLM HTTP call; none/0 disables) -BIOMNI_ENABLE_PROMPT_CACHING=true # Default: true (Anthropic system-prompt cache) +BIOMNI_ENABLE_PROMPT_CACHING=true # Default: true (Anthropic system-prompt cache; also + # through the LLM proxy with the anthropic schema) BIOMNI_RUN_TIMEOUT_SECONDS=600 # Default: unset (total wall-clock budget per run; bounds # the number of ReAct turns. Recommended for # interactive/demo so slow queries fail fast.) @@ -135,6 +171,12 @@ BIOMNI_ENABLE_LLM_TELEMETRY=true # Default: false (library); true in # Emits a per-run llm_usage token/cost event. BIOMNI_RUN_HEARTBEAT_SECONDS=15 # Default: 15 (run-liveness heartbeat cadence; Chainlit) BIOMNI_LOG_REDACT_EMAILS=false # Default: false (opt-in scrub of e-mail-shaped PII in logs) +# Stamped on every log line, to tell services, builds and environments apart +OTEL_SERVICE_NAME=biomni-ad # Default: biomni (BIOMNI_SERVICE_NAME also read) +BIOMNI_VERSION=1.2.3 # Default: the installed package version +BIOMNI_ENV=prod # Default: unset (DEPLOY_ENV, then ENVIRONMENT, also read) +BIOMNI_GIT_SHA=abc1234 # Default: set by the image build (GIT_SHA also read); +BIOMNI_GIT_REF=feat/adworkbench # reported by /healthz and /readyz (GIT_REF also read) # Workspace scanning (bounds every directory walk; see biomni/fs_scan.py) BIOMNI_WORKSPACE_MAX_FILES=10000 # Default: 10000 (hard file cap per scan) @@ -173,11 +215,24 @@ CHAINLIT_AUTH_SECRET=… # Default: generated once and sto # Authentication gateway. Identity headers are IGNORED unless this is enabled: # without a gateway stripping client-supplied copies, anyone could send -# `x-user-id: ` and read or overwrite that person's settings and -# run history. Enable it only when a gateway is actually in front. +# `Ai-App-User-Context: sub=` and read or overwrite that person's +# settings and run history. Enable it only when a gateway is actually in front. BIOMNI_TRUST_AUTH_HEADERS=true # Default: false (fails closed) -# Header names (each accepts a comma-separated list, additive to the defaults) +# Header names. GRIP's composite headers, Ai-App-User-Context (sub, email, +# given_name, family_name) and Ai-App-Workspace-Context (uuid), are read by +# default and need no setting. Change these two only for a gateway that sends +# the same key=value lists under other names: +BIOMNI_AUTH_USER_CONTEXT_HEADER=ai-app-user-context +BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER=ai-app-workspace-context + +# Single-value headers, for gateways that send one field per header +# (oauth2-proxy, Envoy). Each takes a comma-separated list, tried before the +# defaults. They are read only for a request that carries neither context +# header: a gateway strips its own headers from what clients send, not every +# name another gateway might use, so reading them next to GRIP's would let a +# client fill in a field the gateway left out. Pointing one at a composite +# header has no effect, and the app logs a warning at startup saying so. BIOMNI_AUTH_USER_ID_HEADER=x-auth-request-user-id BIOMNI_AUTH_EMAIL_HEADER=x-auth-request-email BIOMNI_AUTH_WORKSPACE_HEADER=x-workspace-id @@ -223,6 +278,11 @@ Two mount-level causes account for nearly all of it: `mountOptions: [uid=57439, gid=57439, dir_mode=0770, file_mode=0770, mfsymlinks]` on the StorageClass or PV. +Changed mount options reach a running deployment only when the node mounts the share afresh. +A StorageClass's options are copied into a PersistentVolume when it is created, and the node keeps a share mounted while any pod there uses it. +So change the PersistentVolume's `spec.mountOptions` as well, then scale the deployment to zero and back up; [grip_deployment.md](grip_deployment.md#storage) has the details and a command to check what the container sees. +At startup the app logs a `storage_check` line for the workspace, the output location and `BIOMNI_STATE_DIR`, with the mount options in effect and, for any location it cannot use, the reason and the fix. + Files attached to a message with 📎 are copied into `/uploads/` under their original names, and it is that path the agent is given. Chainlit's own copy stays where it puts it - a scratch tree it deletes when the @@ -283,12 +343,14 @@ Check this before scaling the deployment out. ### Python Configuration ```python +import os + from biomni.config import default_config # All available settings -default_config.path = "./data" +default_config.path = os.path.expanduser("~/.biomni/data") # "~" is not expanded for you default_config.timeout_seconds = 600 -default_config.llm = "claude-sonnet-4-20250514" +default_config.llm = "claude-sonnet-4-5" default_config.temperature = 0.7 default_config.use_tool_retriever = True default_config.auto_network_limited_mode = True @@ -306,15 +368,25 @@ All are optional; the defaults are what a single-user local run wants. | Variable | Default | What it does | |---|---|---| -| `BIOMNI_DEMO_PASSWORD` | unset | Puts Chainlit's own login in front of the app: one shared password, any username, and the username becomes the identity - so each person gets a separate conversation history. For a deployment with **no** authentication gateway. Leave unset behind a gateway, where it would add a second and weaker door in front of a real one. | +| `BIOMNI_DEMO_PASSWORD` | unset | Puts Chainlit's own login in front of the app: one shared password, any username, and the username becomes the identity - so each person gets a separate conversation history. The histories are separate, not private: anyone with the password can sign in under a username someone else used and see that history. For a deployment with **no** authentication gateway. Leave unset behind a gateway, where it would add a second and weaker door in front of a real one. | | `BIOMNI_BIND_ADDRESS` | `0.0.0.0` | Address the compose deployment publishes on. Set to `127.0.0.1` when a TLS proxy sits in front, so the app is reachable only through it. Note a `ufw deny` will **not** close a published port - Docker writes iptables rules ahead of ufw's chain - so binding to loopback is the reliable way. | | `BIOMNI_TRUST_AUTH_HEADERS` | off | Believe the gateway's identity headers. Fails closed: until this is set, headers are ignored and every session is anonymous. Turn it on in the same change that puts a gateway in front. | | `BIOMNI_AUTH_ISSUER_HEADER` | see identity.py | Header carrying the OIDC issuer. When an issuer is present it is folded into the storage key, because a subject id is unique within a realm rather than globally. **Send it from the first deployment or not at all** - introducing it later re-keys every existing user, who then finds an empty history. | +### Model access + +| Variable | Default | What it does | +|---|---|---| +| `BIOMNI_LLM_PROXY_URL` | unset | Routes every model call through a platform LLM proxy, such as GRIP's in front of Azure AI Foundry, ahead of every provider setting. Each request carries the proxy's token as a bearer token and the session's user and workspace ids as `X-User-Id` and `X-Workspace-Id`, and a call with no gateway identity behind it is not sent. See [grip_deployment.md](grip_deployment.md#llm-proxy). | +| `BIOMNI_LLM_PROXY_API_KEY` | unset | The proxy's token. `/readyz` reports not ready until it is set, and while `BIOMNI_TRUST_AUTH_HEADERS` is off, since the proxy then refuses every call. | +| `BIOMNI_LLM_PROXY_SCHEMA` | `openai` | The API to speak to the proxy: `openai` calls `/v1/chat/completions`, `anthropic` calls `/v1/messages`. | + ### Storage | Variable | Default | What it does | |---|---|---| +| `BIOMNI_USER_DATA_PATH` | unset | The user's workspace. `/readyz` reports not ready until it exists. | +| `BIOMNI_DATA_LAKE_PATH` | inside the checkout or image | The reference datasets. Point it at a volume: inside the container they count against the node's ephemeral storage and are downloaded again after every restart. | | `BIOMNI_STATE_DIR` | unset | Preferences, run records and the chat-history database. **The conversation list on the left needs this to be writable**; without it the app falls back to the workspace root, and if that is mounted read-only there is no thread list at all. Set by both the compose file and the Kubernetes manifest. | | `BIOMNI_STATE_HOST_PATH` | `./state` | Host path compose mounts at `BIOMNI_STATE_DIR`. | | `BIOMNI_OUTPUT_ROOT` | unset | Where run artifacts go. Declare it whenever outputs land on a volume: without it resolution falls through to a candidate flagged *ephemeral*, and the UI warns that results are lost on restart even when they are not. | @@ -351,6 +423,9 @@ All are optional; the defaults are what a single-user local run wants. ## Troubleshooting +**Deployed on GRIP**: +- See the troubleshooting table in [grip_deployment.md](grip_deployment.md#troubleshooting) + **API Key Not Found**: - Check `.env` file exists in your working directory - Verify with: `echo $ANTHROPIC_API_KEY` diff --git a/docs/grip_deployment.md b/docs/grip_deployment.md new file mode 100644 index 000000000..b90f69621 --- /dev/null +++ b/docs/grip_deployment.md @@ -0,0 +1,317 @@ +# Deploying Biomni-AD on the AD Workbench (GRIP) + +This guide is for the team that runs Biomni-AD as an AI App on the AD Workbench GRIP platform. +It follows the "AI Apps Onboarding V3" checklist: what Biomni-AD provides for each item, and what the deployment has to set. +[configuration.md](configuration.md) describes every setting; this page is the part that matters on GRIP. + +## Checklist + +| Onboarding item | Biomni-AD | +|---|---| +| Docker image | `ghcr.io/kaimen-inc/biomni-ad`, see [Image](#image) | +| Kubernetes manifest | [`deploy/k8s/biomni-ad.yaml`](../deploy/k8s/biomni-ad.yaml), with the changes in [Kubernetes manifest](#kubernetes-manifest), or [`docker-compose.yml`](../docker-compose.yml) for a single machine | +| Environment variables | [Settings](#settings) below, and all of them in [configuration.md](configuration.md) | +| Sample input data | Not required. A small synthetic table for a smoke test is in [`samples/`](../samples/README.md) | +| Minimum and recommended resources | [Resources](#resources) | +| User and workspace headers | `Ai-App-User-Context` and `Ai-App-Workspace-Context` are read exactly as the gateway sends them, see [Identity](#identity) | +| `GET /status` | Implemented, `schema_version: 1`, see [Status and health](#status-and-health) | +| Azure Foundry through the LLM proxy | OpenAI and Anthropic schemas, see [LLM proxy](#llm-proxy) | +| Clear mount points for input and output | [Storage](#storage) | +| Configuration bundled, overridable by environment | Defaults are built into the image, and every setting is an environment variable | +| Automated initialisation | Nothing to run: the app creates what it needs when it starts, and fetches reference data the first time a question needs it | + +## Image + +`ghcr.io/kaimen-inc/biomni-ad:`, where `` is one of: + +- `sha-`, the commit's first seven hex digits: one per build, never moved. + Deploy this one. +- `feat-adworkbench`: the latest build of the AD Workbench branch. +- `latest`: the latest build of the default branch. + +The container runs as the non-root user `mambauser`, uid and gid 57439, and serves HTTP on port 8000. +It needs no init container and no setup step. +Logs are one JSON object per line on stdout. + +## Settings + +### Set these + +| Variable | Value on GRIP | Why | +|---|---|---| +| `BIOMNI_TRUST_AUTH_HEADERS` | `true` | Declares that the platform's authentication gateway is in front, so its `Ai-App-*` headers can be believed. It is off by default, so that an app reachable without a gateway cannot be impersonated with a hand-written header. While it is off, every session is anonymous: nothing is saved, and every model call is refused, because the LLM proxy needs a user and a workspace for each call. `/readyz` reports not ready for that reason while the proxy is configured. | +| `BIOMNI_LLM_PROXY_URL` | the proxy's address, for example `https:///v1` | Sends every model call through the platform's LLM proxy. The address works with or without the trailing `/v1`. | +| `BIOMNI_LLM_PROXY_API_KEY` | the proxy's API token, from a Secret | Sent as `Authorization: Bearer ` on every call. | +| `BIOMNI_LLM` | the name of the model the proxy serves | Sent as the `model` of every request. | +| `BIOMNI_LLM_PROXY_SCHEMA` | `openai` (the default) or `anthropic` | The API the proxy speaks for that model. | +| `BIOMNI_USER_DATA_PATH` | `/app/user-data` | The user's workspace. | +| `BIOMNI_OUTPUT_ROOT` | a directory on a writable volume, see [Where results go](#where-results-go) | Where each run's results are written. | +| `BIOMNI_STATE_DIR` | `/app/state` | Each user's settings, run records and chat history. | +| `BIOMNI_DATA_LAKE_PATH` | a directory on a volume, for example `/app/state/data_lake` | The reference datasets, downloaded the first time a question needs each one. Left unset they go inside the container, see [Resources](#resources). | +| `CHAINLIT_AUTH_SECRET` | a long random string, from a Secret | Signs the session cookie. With one replica it can be left unset: the app generates one and keeps it in `BIOMNI_STATE_DIR`. Set it before running more than one replica, or users are signed out whenever a request reaches a different pod. | + +### Do not set these + +| Variable | Why not | +|---|---| +| `BIOMNI_AUTH_USER_ID_HEADER`, `BIOMNI_AUTH_EMAIL_HEADER`, `BIOMNI_AUTH_WORKSPACE_HEADER` | They name headers that carry one plain value each, for gateways other than GRIP's. Pointed at `Ai-App-User-Context` or `Ai-App-Workspace-Context` they have no effect, and the app logs a warning at startup telling you to remove them. | +| `BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE` | For a deployment with no gateway: it makes every session one shared user. Remove it if it was set for testing before the gateway was in place. | +| `BIOMNI_DEMO_PASSWORD` | Puts a password login in front of the app, for a deployment with no gateway. | +| Provider keys: `ANTHROPIC_API_KEY`, `OPENAI_API_KEY`, `AZURE_*` | Unused while `BIOMNI_LLM_PROXY_URL` is set, since every model call goes to the proxy. | + +## Kubernetes manifest + +[`deploy/k8s/biomni-ad.yaml`](../deploy/k8s/biomni-ad.yaml) deploys the app with the volumes and paths described under [Storage](#storage): Azure Disks at `/app/state` and `/app/runs`, and an Azure Files share at `/app/user-data`. +It is a generic manifest. +It exposes the app through a plain Service with no gateway in front, so it leaves the gateway and the LLM proxy off. +On GRIP, change three things, each of which the manifest shows where it applies: + +1. Set `BIOMNI_TRUST_AUTH_HEADERS` to `"true"`, and let traffic reach the Service only through the authentication gateway. +2. Set `BIOMNI_LLM_PROXY_URL`, `BIOMNI_LLM_PROXY_SCHEMA` and `BIOMNI_LLM`, and in the Secret replace `ANTHROPIC_API_KEY` with `BIOMNI_LLM_PROXY_API_KEY`. +3. Bind the `biomni-ad-workspace` claim to the workspace's own Azure Files share, through a PersistentVolume with the mount options in [Azure Files over SMB](#azure-files-over-smb-mount-options), instead of letting it provision a new share. + +## Identity + +The platform's authentication gateway sends two headers: + +```text +Ai-App-User-Context: email=,given_name=,family_name=,sub= +Ai-App-Workspace-Context: uuid= +``` + +Biomni-AD reads both as they are, with no configuration beyond `BIOMNI_TRUST_AUTH_HEADERS=true`. + +- `sub` identifies the user. + Their settings, run records and chat history are stored under it, separately for each workspace `uuid`. +- `email`, `given_name` and `family_name` are only displayed. +- Values are read as CSV. + A value containing a comma is quoted, either on its own (`family_name="Smith, Jr."`) or together with its key (`"family_name=Smith, Jr."`), and a quote inside a quoted value is doubled. + A header with an unterminated quote, or with `sub`, `email`, `iss` or `uuid` given twice, is ignored whole, since valid CSV cannot produce it. +- Names sent as UTF-8, such as `João`, display correctly. +- An optional `iss` key, the OIDC issuer, becomes part of the storage key when present. + Send it from the first deployment or never: adding it later gives every existing user a new, empty history. + +The gateway must remove `Ai-App-User-Context` and `Ai-App-Workspace-Context` from what clients send, and set its own on every request it forwards. +Biomni-AD cannot tell a header the gateway set from one a client forged, which is why trusting them is a setting. +While a request carries either context header, Biomni-AD takes the identity from those headers alone. +It then ignores the single-value headers that other gateways use (`X-User-Id`, `X-Workspace-Id`, `X-Auth-Request-*` and similar), so a client cannot fill in a field the gateway left out. + +A browser signs in once: Chainlit exchanges the gateway's headers for a signed session cookie, valid for 15 days. +If the gateway's headers later name a different user than the cookie, as when someone else signs in on the same browser, the cookie is refused and the page signs in again as the user the gateway names. +A cookie issued by an earlier version of Biomni-AD is refused the same way, once. + +To check that users are recognised: + +- At startup the `access_config` log line shows `trust_auth_headers=true` and the header names in use, followed by a warning for any setting that has no effect. +- Each chat logs a `workspace_session` line with `identity_source=headers`. + Any other value means the session is not signed in through the gateway. +- In the app, the menu under the avatar at the top right shows the user's name. + With `BIOMNI_TRUST_AUTH_HEADERS` off it shows `anonymous`, or `local` under `BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE`. + +If the gateway is ever reconfigured to send the same fields under other header names, `BIOMNI_AUTH_USER_CONTEXT_HEADER` and `BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER` are the settings to change. + +## LLM proxy + +With `BIOMNI_LLM_PROXY_URL` set, every model call Biomni-AD makes goes to the proxy: the agent's reasoning, its research plans, and the model calls some of its tools make. +Each request carries: + +| Header | Value | +|---|---| +| `Authorization` | `Bearer ` | +| `X-User-Id` | `sub` from the session's `Ai-App-User-Context` | +| `X-Workspace-Id` | `uuid` from the session's `Ai-App-Workspace-Context` | + +The ids are those of the chat the call is made for, attached as each request is built, so when several people use the app at once each is charged for their own calls. +That includes the calls made by the analysis code the agent writes, from the code step itself and from the `concurrent.futures` thread pools it starts. +A call with no gateway identity behind it is never sent: the user is told why, and the proxy never sees a request it would have to refuse. + +| `BIOMNI_LLM_PROXY_SCHEMA` | Endpoint | Notes | +|---|---|---| +| `openai` | `POST /v1/chat/completions` | For reasoning models (`gpt-5*` and the o-series: `o1`, `o3`, `o4-mini` and so on), `temperature` and `stop` are left out, since those models reject them. | +| `anthropic` | `POST /v1/messages` | The token is sent as a bearer token, and no `x-api-key` header is sent. The system prompt carries a `cache_control` marker, so the provider can reuse it from one turn to the next. | + +If the proxy refuses requests that carry `cache_control`, set `BIOMNI_ENABLE_PROMPT_CACHING=false`. +Biomni-AD does not use `/v1/embeddings`. +The one tool that would reach a model provider directly, the Claude web search, is withdrawn whenever the proxy is configured, and the agent searches with its other tools instead. + +When the proxy refuses a call, the user reads why in plain words rather than a stack trace, for example: + +- HTTP 429: "The language model usage limit has been reached (HTTP 429). The platform limits model usage per user and per workspace. Try again later." +- HTTP 403: "The platform's LLM proxy refused the request (HTTP 403)", followed by the proxy's own explanation. +- HTTP 401: the token was rejected, and the message names `BIOMNI_LLM_PROXY_API_KEY`. +- HTTP 404: the model is not served, and the message names `BIOMNI_LLM`. + +The proxy's explanation is taken from its JSON error body (`error.message`, `error`, `message` or `detail`). +A body that is not JSON, such as a gateway's HTML error page, is left out. +A refusal ends the run where it is, with what it had produced so far, instead of the run carrying on without a model. + +### Network access + +Besides the proxy, the app makes outbound HTTPS requests for its data: + +- `https://biomni-release.s3.amazonaws.com`, the public data lake, fetched one dataset at a time as questions need them; +- `https://www.google.com`, for the agent's web search while the proxy is configured; +- public biomedical services its tools query, such as NCBI (PubMed), UniProt, Ensembl and the GWAS Catalog; +- the hosts in the Alzheimer's dataset catalogues AD1 draws on, chiefly NIAGADS (`st1.niagads.org`, `dss.niagads.org`) and Zenodo (`zenodo.org`), from which the agent's code downloads a dataset when a question needs it. + +The analysis code the agent writes can also fetch other public URLs a question or a catalogue points it to. +An egress allowlist therefore limits which questions can be answered, and a question that needs a blocked host fails at that step with the connection error. + +## Storage + +The app runs as uid 57439 and uses these locations: + +| Location | Set by | Holds | Needs | +|---|---|---|---| +| `/app/user-data` | `BIOMNI_USER_DATA_PATH` | the user's workspace: their input data | read, and write too if results are to go into it | +| a results directory | `BIOMNI_OUTPUT_ROOT` | a folder per run: figures, tables and the code that produced them | write | +| `/app/state` | `BIOMNI_STATE_DIR` | each user's settings, run records and chat history | write | +| a reference-data directory | `BIOMNI_DATA_LAKE_PATH` | reference datasets, downloaded the first time a question needs them | write; inside the container at `/app/data/biomni_data/data_lake` unless pointed at a volume | + +On GRIP today, `/app/runs` and `/app/state` are Azure Disks, and `/app/user-data` is an Azure Files share mounted over SMB. +The two kinds of volume are given to uid 57439 in different ways. + +### Azure Disk: `fsGroup` + +A new disk mounts as `root:root 0755`. +`securityContext.fsGroup: 57439` on the pod, as in the manifest, makes the kubelet give it to the app's group. + +### Azure Files over SMB: mount options + +SMB has no Unix ownership of its own. +The owner and mode the container sees come only from the mount options, and `fsGroup` is ignored. +Mounted without them, the share appears as `root:root 0755`: readable, but not writable by uid 57439. +Set these on the StorageClass or the PersistentVolume: + +```yaml +mountOptions: + - uid=57439 + - gid=57439 + - dir_mode=0770 + - file_mode=0770 + - mfsymlinks +``` + +Two things decide whether changed options actually reach the container: + +1. A StorageClass's options are copied into each PersistentVolume when it is created, so changing the StorageClass does not change a volume that already exists. + Change the PersistentVolume's `spec.mountOptions` too. +2. Options take effect when the node mounts the share, and the node keeps it mounted while any pod there uses the volume, so a rolling restart can start the new pod on the old mount. + Scale the deployment to zero, wait for its pods to be gone, then scale it back up. + +A partly applied set of options is worse than none. +With `dir_mode=0770` in effect but not `uid=57439` and `gid=57439`, the share belongs to root and grants nothing to anyone else, so the app cannot even read the workspace. + +A share mounted with `cifsacl` or `modefromsid` is different: the file server's ACLs decide the owner and mode, and those are what must grant uid 57439 access. + +### Checking what the container sees + +```bash +kubectl exec deploy/biomni-ad -- sh -c 'id; grep " /app/user-data " /proc/self/mounts; ls -ldn /app/user-data' +``` + +Mounted correctly, this shows uid 57439, mount options that include `uid=57439` and `gid=57439`, and the directory as `drwxrwx--- ... 57439 57439`. + +The app makes the same check itself when it starts, and logs one `storage_check` line per location with the live mount options. +When it cannot read or write a location, the line says why and what to change: + +```bash +kubectl logs deploy/biomni-ad | grep storage_check +``` + +For a share mounted without ownership options, the reason reads: + +```text +/app/user-data is owned by uid 0, gid 0 with mode 0755 on an SMB share (/). Its owner and mode come only from the mount options, which are uid=0,gid=0,dir_mode=0755 on the mount this container sees. Mount it with uid=57439,gid=57439,dir_mode=0770,file_mode=0770; new options only apply once the share is mounted afresh on the node +``` + +The share is named without its server, so the storage account does not appear in logs or in the app. +Users see the same explanation in the app: at the start of each chat when the workspace cannot be read, and in the Settings dialog when results cannot go where the deployment meant them to. + +### Where results go + +Results go to the first of these that the app can write to: the folder the user chose in Settings, `BIOMNI_OUTPUT_ROOT`, `/biomni-outputs`, and `./runs` in the app's working directory. +A location that cannot be written to is passed over, and the app says so in the log and in Settings. + +- `BIOMNI_OUTPUT_ROOT=/app/user-data/biomni-outputs` puts results in the workspace, next to the user's other files, once the share is writable by uid 57439. +- `BIOMNI_OUTPUT_ROOT=/app/runs`, on an Azure Disk, keeps results durable in the meantime. + Users download each run's results from the app. + +Avoid pointing `BIOMNI_OUTPUT_ROOT` at the workspace root itself, since every run's folder would then appear at the top level of the workspace. + +## Status and health + +| Endpoint | Answers | +|---|---| +| `GET /status` | The platform's activity check, below. Also usable as a health check. | +| `GET /healthz` | Liveness: the process is serving. | +| `GET /readyz` | Readiness: the workspace at `BIOMNI_USER_DATA_PATH` exists, and a model can be reached. With the proxy configured, it reports not ready (HTTP 503) until `BIOMNI_LLM_PROXY_API_KEY` is set and `BIOMNI_TRUST_AUTH_HEADERS` is on. The response body names the check that failed. | + +`GET /status` returns: + +```json +{"active": true, "last_activity": 1790440933, "schema_version": 1, "background_jobs": 1, "open_sessions": 1, "internal_processes": {"active_threads": 5, "database_connections": 3}} +``` + +- `active` is true while any run is in progress, however long it takes, and for 4 hours after the last user activity (`BIOMNI_STATUS_INACTIVITY_SECONDS`). +- While a browser session is connected, the idle window is 12 hours instead (`BIOMNI_STATUS_OPEN_SESSION_INACTIVITY_SECONDS`), because reclaiming the pod would discard the in-memory state of that user's analysis. +- Both windows fall inside the platform's range of 4 to 24 hours. +- Health checks and `/status` requests themselves never count as activity. +- `last_activity` is a Unix timestamp in seconds. + +## Resources + +| | Minimum | Recommended | +|---|---|---| +| CPU | 2 vCPU | 4 vCPU | +| Memory | 8 GiB | 16 GiB | +| GPU | none | none | +| Disk for reference data | 20 GiB | 20 GiB | + +These come from a measured run of the "Multi-omics AD risk gene portrait" starter, a full AD1 analysis drawing on GWAS, eQTL, proteomic and CRISPR datasets: 11 minutes, 23 code steps, 31 model calls. + +- **No GPU.** + Every model runs behind the LLM proxy, and the analysis code the agent writes is ordinary CPU work: pandas, SciPy, R. +- **CPU.** + The run spent most of its time waiting for the model, and used at most one core, during code steps. + The minimum keeps a core free for the interface while a run computes; the recommendation lets several people's runs compute at once. +- **Memory.** + The run peaked at about 1 GB. + What an analysis loads matters more: most large reference datasets, such as GeneBass's variant tables, are 1 to 1.7 GB on disk and several GB once loaded, and each open chat keeps what it loaded in memory until the chat is evicted (`BIOMNI_MAX_REPL_SESSIONS`). + The minimum fits one person working with those, the recommendation several. + The exception is BindingDB (`BindingDB_All_202409.tsv`): 6.25 GB on disk and about 13 GiB once loaded whole, more than the minimum and most of the recommendation. + An analysis that loads all of it can exhaust the pod's memory, and the container then restarts for everyone using it. +- **Disk.** + Reference datasets are downloaded the first time a question needs them, about 16.5 GB if every one is used: 15 GB of Biomni's data lake, 6.25 GB of it a single BindingDB table, and 1.4 GB of AD1's catalogue data. + By default they are written inside the container, where they count against the node's ephemeral storage and are downloaded again after every restart. + Point `BIOMNI_DATA_LAKE_PATH` at a directory on a volume, as the manifest does. + A download cut short, by a pod restart say, leaves no half-written dataset behind: the file is fetched again the next time a question needs it. + +The same run sent about 800,000 input and 20,000 output tokens through the proxy, and at its busiest exceeded 100,000 input tokens a minute. +Per-user limits below that slow long runs down with retries, and eventually stop them with the usage-limit message above. + +## Smoke test + +After deploying, sign in through the platform and check, in order: + +1. The menu under the avatar at the top right shows your name, and no warning appears when the chat opens. +2. Copy [`samples/synthetic_gwas_hits.csv`](../samples/synthetic_gwas_hits.csv) into the workspace and ask: "Plot the -log10 p-values in synthetic_gwas_hits.csv as a bar chart." + Approve the plan it proposes. + The run should read the file, draw the chart, and save it in a new run folder under `BIOMNI_OUTPUT_ROOT`. +3. The proxy's log shows each model call with the user's `sub` as `X-User-Id` and the workspace `uuid` as `X-Workspace-Id`. +4. Settings, reopened after a page reload, still show what was saved, and the conversation is listed on the left. + +## Troubleshooting + +| What you see | Cause | Fix | +|---|---|---| +| "The assistant cannot answer in this session" when a chat opens, and Settings say "Settings apply to this session only: sign-in details from the authentication gateway were received but ignored, because BIOMNI_TRUST_AUTH_HEADERS is not enabled" | Trust in the gateway's headers is off | Set `BIOMNI_TRUST_AUTH_HEADERS=true` | +| The page says "Login to access the app" and "Unable to sign in", and the log says "rejecting a session with no gateway identity headers (BIOMNI_TRUST_AUTH_HEADERS is on)" | The request reached the app without `Ai-App-User-Context`, with neither `sub` nor `email` in it, or with a header ignored as malformed (a warning just before says which) | Check that every route to the app passes through the gateway, and that the gateway writes valid CSV | +| "The assistant cannot answer in this session", and Settings say "Settings apply to this session only: the authentication gateway did not identify this session" | The chat connected without the gateway's headers, and the browser's sign-in could not stand in for them | Reload the page, which signs in again | +| "This browser is still signed in to Biomni as a different user than the one the platform reports" | Someone else used the app in this browser, and their session cookie is still there | Reload the page, which signs in as the user the gateway names | +| Startup warning "`BIOMNI_AUTH_..._HEADER=...` has no effect" | A single-value header setting points at a context header | Remove the variable | +| "This app cannot read your workspace" | The workspace share's mount options do not give uid 57439 access | [Azure Files over SMB](#azure-files-over-smb-mount-options) | +| Settings say "The configured location could not be used, so results go to `` instead" | `BIOMNI_OUTPUT_ROOT`, or the folder chosen in Settings, is not writable by uid 57439 | The reason shown names the fix; see [Storage](#storage) | +| "The language model usage limit has been reached (HTTP 429)" | The platform's per-user or per-workspace limit | Wait, or raise the limit in the proxy | +| An HTTP 401, 403 or 404 from "the platform's LLM proxy" | The token, the workspace's access to the model, or the model name | `BIOMNI_LLM_PROXY_API_KEY`, the proxy's policy, or `BIOMNI_LLM` | +| `/readyz` returns 503 | See `checks` in the response body | Usually the workspace mount, `BIOMNI_LLM_PROXY_API_KEY`, or `BIOMNI_TRUST_AUTH_HEADERS` | diff --git a/public/biomni.css b/public/biomni.css new file mode 100644 index 000000000..f677f1aac --- /dev/null +++ b/public/biomni.css @@ -0,0 +1,12 @@ +/* + * Paragraph breaks in the chat settings dialog. + * + * Chainlit renders a widget's description as plain text in a block that + * collapses newlines, so the output-directory help - a sentence, then each + * warning on its own - arrived as one unbroken wall of text. pre-line keeps the + * line breaks and still folds runs of spaces, so a description written without + * breaks renders exactly as before. + */ +[role="dialog"] .text-muted-foreground { + white-space: pre-line; +} diff --git a/public/logo_dark.png b/public/logo_dark.png new file mode 100644 index 000000000..51477e177 Binary files /dev/null and b/public/logo_dark.png differ diff --git a/public/logo_light.png b/public/logo_light.png new file mode 100644 index 000000000..51477e177 Binary files /dev/null and b/public/logo_light.png differ diff --git a/pyproject.toml b/pyproject.toml index c3a205d1c..c7db7ece3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,22 +37,28 @@ Upstream = "https://github.com/snap-stanford/biomni" [project.optional-dependencies] # LLM provider extras - install the ones you need. -anthropic = ["langchain-anthropic>=0.2,<0.4"] +# 0.3.1 is the first to build its client from _client_params, which is where the +# LLM proxy's bearer token replaces the API key header (biomni/llm.py). +anthropic = ["langchain-anthropic>=0.3.1,<0.4"] bedrock = ["langchain-aws>=0.2,<0.4"] ollama = ["langchain-ollama>=0.2,<0.4"] # UI extras gradio = ["gradio>=5.0,<6.0"] -# The UI persists chat threads through Chainlit's SQLAlchemy data layer: -# "custom-data" pulls in SQLAlchemy, aiosqlite drives the default SQLite store. -chainlit = ["chainlit[custom-data]>=2.8,<3", "aiosqlite>=0.20,<1.0"] +# The UI persists chat threads through Chainlit's SQLAlchemy data layer, with +# aiosqlite for the default SQLite store and asyncpg for Postgres. Chainlit is +# pinned to the version the app is tested against, as in the Dockerfile: the app +# hooks its login dependency and styles its settings dialog. +chainlit = ["chainlit==2.11.1", "sqlalchemy>=2.0.28,<3", "aiosqlite>=0.20,<1.0", "asyncpg>=0.30,<1"] # Convenience meta-extra all = [ - "langchain-anthropic>=0.2,<0.4", + "langchain-anthropic>=0.3.1,<0.4", "langchain-aws>=0.2,<0.4", "langchain-ollama>=0.2,<0.4", "gradio>=5.0,<6.0", - "chainlit[custom-data]>=2.8,<3", + "chainlit==2.11.1", + "sqlalchemy>=2.0.28,<3", "aiosqlite>=0.20,<1.0", + "asyncpg>=0.30,<1", ] [tool.setuptools] diff --git a/samples/README.md b/samples/README.md new file mode 100644 index 000000000..1af4fcaef --- /dev/null +++ b/samples/README.md @@ -0,0 +1,15 @@ +# Sample input + +Biomni-AD needs no sample input: it works on whatever is in the user's workspace, and on the public AD data lake, which it downloads as questions need it. + +`synthetic_gwas_hits.csv` exists to check a new deployment end to end. +Its gene and chromosome columns are real, but its p-values are invented, so it is only fit for testing. + +Copy it into a workspace, open Biomni-AD from the platform, and ask: + +> Plot the -log10 p-values in synthetic_gwas_hits.csv as a bar chart. + +Approve the plan it proposes. +A working deployment then reads the file from the workspace, runs the plotting code, shows the chart, and saves it with the code that drew it in a new run folder under the output directory. +Each step exercises one part of the deployment: the workspace mount, the model calls through the LLM proxy, and the output volume. +See [docs/grip_deployment.md](../docs/grip_deployment.md#smoke-test) for the full check. diff --git a/samples/synthetic_gwas_hits.csv b/samples/synthetic_gwas_hits.csv new file mode 100644 index 000000000..1ceadee1c --- /dev/null +++ b/samples/synthetic_gwas_hits.csv @@ -0,0 +1,11 @@ +gene,chromosome,p_value +APOE,19,1e-300 +BIN1,2,3e-60 +CLU,8,2e-25 +PICALM,11,5e-24 +CR1,1,8e-22 +TREM2,6,4e-20 +SORL1,11,6e-18 +ABCA7,19,9e-16 +MS4A6A,11,3e-14 +CD33,19,2e-11 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..0d148fdb3 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,19 @@ +"""Process-wide state that one test must not leak into the next.""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor + +import pytest + + +@pytest.fixture(autouse=True) +def _restore_thread_pool_submit(): + """Undo ``carry_context_into_thread_pools``, which running generated code installs. + + The app installs it once and for good. In a test run that would change how + every later test's thread pools behave, depending on which test ran first. + """ + submit = ThreadPoolExecutor.submit + yield + ThreadPoolExecutor.submit = submit # type: ignore[method-assign] diff --git a/tests/login_guard_app.py b/tests/login_guard_app.py new file mode 100644 index 000000000..48c843a9f --- /dev/null +++ b/tests/login_guard_app.py @@ -0,0 +1,113 @@ +"""Sign in to the real app through Chainlit's HTTP API; run by test_login_guard.py. + +A separate interpreter because chainlit_app reads its settings when imported and +registers its callbacks on Chainlit's module-global server, so it cannot be set +up twice, with different settings, in one test process. + +Prints one line, ``LOGIN_GUARD_RESULT {"checks": [[label, got, want], ...]}``. +""" + +import asyncio +import json +import os +import sys +from pathlib import Path + +import dotenv + +# The environment is the test's to define: a developer's .env, which the app +# loads over it, could otherwise switch the gateway trust off. +dotenv.load_dotenv = lambda *args, **kwargs: False + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +os.chdir(os.environ["CHAINLIT_APP_ROOT"]) + +import chainlit_app # noqa: F401 (registers the callbacks and installs the guard) +from chainlit.auth import create_jwt, get_current_user +from chainlit.data import get_data_layer +from chainlit.server import app +from chainlit.user import User +from starlette.testclient import TestClient + +checks: list[list] = [] + + +def check(label, got, want): + checks.append([label, got, want]) + + +def gateway(sub, workspace="ws-1"): + """The headers the platform's gateway sends for ``sub`` in ``workspace``.""" + return { + "Ai-App-User-Context": f"email={sub}@example.org,given_name={sub.title()},family_name=X,sub={sub}", + "Ai-App-Workspace-Context": f"uuid={workspace}", + } + + +def cookie_for(user): + client = TestClient(app) + client.cookies.set("access_token", create_jwt(user)) + return client + + +check("the guard is installed", get_current_user in app.dependency_overrides, True) + +alice = TestClient(app) +check("alice signs in", alice.post("/auth/header", headers=gateway("alice")).status_code, 200) +me = alice.get("/user", headers=gateway("alice")) +check("alice's cookie with alice's headers", me.status_code, 200) +check("alice's cookie carries her subject", me.json()["metadata"]["user_id"], "alice") +check("alice's cookie with bob's headers", alice.get("/user", headers=gateway("bob")).status_code, 401) +check("alice's cookie in another workspace", alice.get("/user", headers=gateway("alice", "ws-2")).status_code, 401) +check( + "alice's cookie on /project/settings as bob", + alice.get("/project/settings", headers=gateway("bob")).status_code, + 401, +) +check("health checks need no sign-in", alice.get("/healthz", headers=gateway("bob")).status_code, 200) + +# Cookies the previous release issued while the gateway was not trusted: one +# shared identity for every browser, and no user_id in the metadata. +released = {"email": "", "workspace_id": "", "given_name": "", "family_name": ""} +for identity in ("local", "anonymous"): + stale = cookie_for(User(identifier=identity, display_name=identity, metadata={"source": identity, **released})) + check(f"a released '{identity}' cookie as alice", stale.get("/user", headers=gateway("alice")).status_code, 401) + listing = stale.post("/project/threads", headers=gateway("carol"), json={"pagination": {"first": 20}, "filter": {}}) + check(f"a released '{identity}' cookie listing threads as carol", listing.status_code, 401) + +# A gateway cookie from the previous release: the right person, but it cannot +# say which subject it was issued for. +old_bob = cookie_for( + User( + identifier="bob-81b637d8fc@ws-1", + display_name="Bob X", + metadata={"source": "headers", "email": "bob@example.org", "workspace_id": "ws-1", "given_name": "Bob"}, + ) +) +check("a released gateway cookie for bob, as bob", old_bob.get("/user", headers=gateway("bob")).status_code, 401) +check("bob signs in again", old_bob.post("/auth/header", headers=gateway("bob")).status_code, 200) +check("bob's new cookie", old_bob.get("/user", headers=gateway("bob")).status_code, 200) + +# A client-sent copy of the context header under an underscore name. +forged = TestClient(app) +forged.post("/auth/header", headers={**gateway("carol"), "Ai_App_User_Context": "sub=alice"}) +check( + "an underscore copy of the header does not sign in as alice", + forged.get("/user", headers=gateway("carol")).json()["metadata"]["user_id"], + "carol", +) + + +async def write_alices_thread(): + await get_data_layer().update_thread("t-alice-1", name="alice: APOE4 carriers", user_id=me.json()["id"]) + + +asyncio.run(write_alices_thread()) +carol = TestClient(app) +carol.post("/auth/header", headers=gateway("carol")) +listed = carol.post("/project/threads", headers=gateway("carol"), json={"pagination": {"first": 20}, "filter": {}}) +check("carol's thread list", [t.get("name") for t in listed.json().get("data", [])], []) +opened = carol.get("/project/thread/t-alice-1", headers=gateway("carol")).status_code +check("carol cannot open alice's thread", opened in (401, 403, 404), True) + +print("LOGIN_GUARD_RESULT " + json.dumps({"checks": checks}), flush=True) diff --git a/tests/test_agent_prompt_tools.py b/tests/test_agent_prompt_tools.py new file mode 100644 index 000000000..7ffe796ad --- /dev/null +++ b/tests/test_agent_prompt_tools.py @@ -0,0 +1,67 @@ +"""The agent's system prompt must name only tools the agent can actually call. + +It tells the model which functions to reach for first, on every question. A +name that does not exist - three of the old ones did not - or a tool this +deployment withholds is a failed step the user watches, every time. +""" + +from __future__ import annotations + +import re + +import pytest + +try: + from biomni.agent.a1 import A1 + + HAS_AGENT = True +except Exception: # pragma: no cover - depends on the installed extras + HAS_AGENT = False + +pytestmark = pytest.mark.skipif(not HAS_AGENT, reason="agent stack (pandas/langchain/langgraph) not installed") + + +class _StubAgent: + def __init__(self, module2api) -> None: + self.module2api = module2api + + _advertises_tool = A1._advertises_tool if HAS_AGENT else None + _tool_priority_instructions = A1._tool_priority_instructions if HAS_AGENT else None + + +def _full_catalogue() -> dict[str, list[dict]]: + import importlib + + from biomni.utils import read_module2api + + # Every tool that exists, whatever this machine's deployment would withhold. + fields = [module.rsplit(".", 1)[1] for module in read_module2api()] + return { + f"biomni.tool.{field}": importlib.import_module(f"biomni.tool.tool_description.{field}").description + for field in fields + } + + +def _named_functions(text: str) -> set[str]: + return set(re.findall(r"\b([a-z_][a-z0-9_]*)\(\)", text)) + + +def test_every_function_the_prompt_names_exists(): + catalogue = _full_catalogue() + known = {api["name"] for apis in catalogue.values() for api in apis} + named = _named_functions(_StubAgent(catalogue)._tool_priority_instructions()) + assert named, "the prompt should name the tools to use first" + assert named <= known, f"named but not a tool: {sorted(named - known)}" + + +def test_the_claude_web_search_is_named_only_when_advertised(): + catalogue = _full_catalogue() + assert "advanced_web_search_claude()" in _StubAgent(catalogue)._tool_priority_instructions() + + withheld = { + module: [api for api in apis if api["name"] != "advanced_web_search_claude"] + for module, apis in catalogue.items() + } + text = _StubAgent(withheld)._tool_priority_instructions() + assert "advanced_web_search_claude" not in text + assert "search_google()" in text diff --git a/tests/test_data_lake_download.py b/tests/test_data_lake_download.py new file mode 100644 index 000000000..bd649d7ce --- /dev/null +++ b/tests/test_data_lake_download.py @@ -0,0 +1,195 @@ +"""Tests for fetching data lake files - biomni.utils.check_and_download_s3_files. + +The data lake counts a file that exists as downloaded, so what matters is that +a file only ever exists complete: a download cut short must leave nothing at +the target, or every later query reads the truncated dataset and none fetches +it again. +""" + +from __future__ import annotations + +import os +import threading +import time + +import pytest +import requests +from biomni import utils + +BUCKET = "https://bucket.example" +NAME = "dataset.tsv" +PAYLOAD = [b"gene\tscore\n", b"APOE\t9.1\n", b"TREM2\t7.4\n"] + + +class FakeResponse: + """The slice of ``requests.Response`` the downloader uses.""" + + def __init__(self, chunks, *, content_length=None, fail_at=None, error=None, status=200): + self.chunks = chunks + size = sum(len(c) for c in chunks) if content_length is None else content_length + self.headers = {"content-length": str(size)} + self.fail_at = fail_at + self.error = error + self.status = status + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + def raise_for_status(self): + if self.status >= 400: + raise requests.HTTPError(f"{self.status} Client Error") + + def iter_content(self, chunk_size): + for index, chunk in enumerate(self.chunks): + if index == self.fail_at: + raise self.error + yield chunk + + +def serve(monkeypatch, response_for): + """Route the downloader's requests to ``response_for(url)``, recording each url.""" + calls = [] + + def fake_get(url, **kwargs): + calls.append(url) + assert kwargs.get("timeout"), "a download without a timeout can hang the chat turn forever" + return response_for(url) + + monkeypatch.setattr(utils.requests, "get", fake_get) + return calls + + +def fetch(lake): + return utils.check_and_download_s3_files(BUCKET, str(lake), [NAME], folder="data_lake") + + +def test_a_completed_download_is_renamed_into_place(tmp_path, monkeypatch): + calls = serve(monkeypatch, lambda url: FakeResponse(PAYLOAD)) + + assert fetch(tmp_path) == {NAME: True} + assert calls == [f"{BUCKET}/data_lake/{NAME}"] + assert (tmp_path / NAME).read_bytes() == b"".join(PAYLOAD) + assert os.listdir(tmp_path) == [NAME] + + +def test_a_connection_closed_early_leaves_nothing_behind(tmp_path, monkeypatch): + serve(monkeypatch, lambda url: FakeResponse(PAYLOAD[:1], content_length=1000)) + + assert fetch(tmp_path) == {NAME: False} + assert os.listdir(tmp_path) == [] + + +def test_an_error_mid_download_leaves_nothing_behind(tmp_path, monkeypatch): + serve(monkeypatch, lambda url: FakeResponse(PAYLOAD, fail_at=2, error=requests.ConnectionError("reset"))) + + assert fetch(tmp_path) == {NAME: False} + assert os.listdir(tmp_path) == [] + + +def test_an_http_error_leaves_nothing_behind(tmp_path, monkeypatch): + serve(monkeypatch, lambda url: FakeResponse([b"NoSuchKey"], status=404)) + + assert fetch(tmp_path) == {NAME: False} + assert os.listdir(tmp_path) == [] + + +def test_an_interrupted_download_leaves_nothing_behind(tmp_path, monkeypatch): + # KeyboardInterrupt is not an Exception: the old in-place download left its + # partial file at the target, where it passed for the finished dataset. + serve(monkeypatch, lambda url: FakeResponse(PAYLOAD, fail_at=1, error=KeyboardInterrupt())) + + with pytest.raises(KeyboardInterrupt): + fetch(tmp_path) + assert os.listdir(tmp_path) == [] + + +def test_the_file_does_not_exist_until_every_byte_has_arrived(tmp_path, monkeypatch): + # A process killed at any point - a pod evicted, the OOM killer - runs no + # cleanup at all, so the target must simply not exist before the end. + target = tmp_path / NAME + seen = [] + + class WatchedResponse(FakeResponse): + def iter_content(self, chunk_size): + for chunk in super().iter_content(chunk_size): + seen.append(target.exists()) + yield chunk + + serve(monkeypatch, lambda url: WatchedResponse(PAYLOAD)) + + assert fetch(tmp_path) == {NAME: True} + assert seen == [False] * len(PAYLOAD) + assert target.read_bytes() == b"".join(PAYLOAD) + + +def test_what_a_killed_download_leaves_does_not_stop_the_next_fetch(tmp_path, monkeypatch): + # What a SIGKILL leaves: a partial file, which the next fetch must ignore. + (tmp_path / f".{NAME}.abc123.part").write_bytes(PAYLOAD[0]) + calls = serve(monkeypatch, lambda url: FakeResponse(PAYLOAD)) + + assert fetch(tmp_path) == {NAME: True} + assert len(calls) == 1 + assert (tmp_path / NAME).read_bytes() == b"".join(PAYLOAD) + + +def test_abandoned_partial_downloads_are_swept_and_live_ones_kept(tmp_path, monkeypatch): + abandoned = tmp_path / f".{NAME}.old.part" + live = tmp_path / f".{NAME}.new.part" + other_file = tmp_path / ".other.tsv.old.part" + for partial in (abandoned, live, other_file): + partial.write_bytes(b"x") + long_ago = time.time() - utils._STALE_PARTIAL_S - 60 + os.utime(abandoned, (long_ago, long_ago)) + os.utime(other_file, (long_ago, long_ago)) + serve(monkeypatch, lambda url: FakeResponse(PAYLOAD)) + + assert fetch(tmp_path) == {NAME: True} + assert not abandoned.exists() + assert live.exists() # another process may still be writing it + assert other_file.exists() # swept when that file is next fetched + + +def test_a_file_already_present_is_not_fetched(tmp_path, monkeypatch): + (tmp_path / NAME).write_bytes(b"local copy") + calls = serve(monkeypatch, lambda url: FakeResponse(PAYLOAD)) + + assert fetch(tmp_path) == {NAME: True} + assert calls == [] + assert (tmp_path / NAME).read_bytes() == b"local copy" + + +def test_two_chats_needing_the_same_file_download_it_once(tmp_path, monkeypatch): + release = threading.Event() + first_started = threading.Event() + + class SlowResponse(FakeResponse): + def iter_content(self, chunk_size): + first_started.set() + assert release.wait(5) + yield from super().iter_content(chunk_size) + + calls = serve(monkeypatch, lambda url: SlowResponse(PAYLOAD)) + results = {} + seen = {} + + def chat(key): + results[key] = fetch(tmp_path) + # What the chat's code step would read the moment the fetch returns. + seen[key] = (tmp_path / NAME).read_bytes() if (tmp_path / NAME).exists() else None + + first = threading.Thread(target=chat, args=("first",)) + first.start() + assert first_started.wait(5) + second = threading.Thread(target=chat, args=("second",)) + second.start() + time.sleep(0.2) # the second is now waiting for the first's download + release.set() + first.join(5) + second.join(5) + + assert results == {"first": {NAME: True}, "second": {NAME: True}} + assert seen == {"first": b"".join(PAYLOAD), "second": b"".join(PAYLOAD)} + assert len(calls) == 1 diff --git a/tests/test_graph_stream.py b/tests/test_graph_stream.py new file mode 100644 index 000000000..1a44b155a --- /dev/null +++ b/tests/test_graph_stream.py @@ -0,0 +1,136 @@ +"""Tests for chainlit_ui.graph_stream - a LangGraph run read from the event loop. + +The graph is a stand-in with LangGraph's ``stream`` signature, so these run +without LangGraph and exercise exactly what crosses the thread boundary: states, +a failure, the order to stop, and the caller's context. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import threading + +import pytest +from biomni.identity import UserIdentity, bound_identity, current_identity +from chainlit_ui.graph_stream import stream_langgraph + + +class _Graph: + """Yields the given states, then raises ``error`` if one is given.""" + + def __init__(self, states, error: Exception | None = None, gate: threading.Event | None = None) -> None: + self.states = states + self.error = error + self.gate = gate + self.produced = 0 + self.seen_identity: list[object] = [] + self.finished = threading.Event() + + def stream(self, inputs, stream_mode, config): + assert stream_mode == "values" + try: + for state in self.states: + self.seen_identity.append(current_identity()) + self.produced += 1 + yield state + if self.gate is not None: + self.gate.wait(timeout=5) + if self.error is not None: + raise self.error + finally: + self.finished.set() + + +async def _collect(graph) -> list: + return [state async for state in stream_langgraph(graph, {"messages": []}, {})] + + +def test_states_arrive_in_order(): + assert asyncio.run(_collect(_Graph([1, 2, 3]))) == [1, 2, 3] + + +def test_a_failed_run_raises_in_the_reader_instead_of_ending_quietly(): + """It used to end like a finished run, and the user was told "Run complete".""" + graph = _Graph([1], error=PermissionError("refused by the LLM proxy")) + received: list = [] + + async def read(): + async for state in stream_langgraph(graph, {}, {}): + received.append(state) + + with pytest.raises(PermissionError, match="refused by the LLM proxy"): + asyncio.run(read()) + assert received == [1] + + +def test_a_run_that_fails_before_its_first_state_still_raises(): + with pytest.raises(RuntimeError, match="boom"): + asyncio.run(_collect(_Graph([], error=RuntimeError("boom")))) + + +def test_the_graph_runs_as_the_session_that_started_it(): + who = UserIdentity(user_id="u1", workspace_id="ws-1", source="headers") + graph = _Graph([1, 2]) + + async def read(): + with bound_identity(who): + return await _collect(graph) + + asyncio.run(read()) + assert graph.seen_identity == [who, who] + + +def test_stopping_the_reader_stops_the_graph_after_its_current_step(): + """Stop cancels the reading task; the thread running the graph must follow. + + It used to keep calling the model for as long as the whole plan took. + """ + gate = threading.Event() + graph = _Graph(list(range(100)), gate=gate) + + async def main(): + first = asyncio.Event() + + async def read(): + async for _state in stream_langgraph(graph, {}, {}): + first.set() + + task = asyncio.create_task(read()) + await first.wait() + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + gate.set() # let the node in progress finish + assert await asyncio.to_thread(graph.finished.wait, 5) + + asyncio.run(main()) + # The node that was running when Stop came in, and not one more. + assert graph.produced == 2 + + +class _SharedListGraph: + """Appends to one messages list in place and yields it, as A1's nodes do.""" + + def stream(self, inputs, stream_mode, config): + state = {"messages": [], "next_step": None} + for i in range(5): + state["messages"].append(f"message {i}") + yield state + + +def test_each_state_arrives_as_it_was_when_streamed(): + """A slow reader must not see messages from nodes that ran after the state it is reading. + + It did: a code step went missing from the transcript and the observation + after it was shown twice. + """ + + async def read_slowly(): + seen = [] + async for state in stream_langgraph(_SharedListGraph(), {}, {}): + await asyncio.sleep(0.02) # the graph runs ahead meanwhile + seen.append(state["messages"][-1]) + return seen + + assert asyncio.run(read_slowly()) == [f"message {i}" for i in range(5)] diff --git a/tests/test_health.py b/tests/test_health.py index d3f54a01f..97ff8cd51 100644 --- a/tests/test_health.py +++ b/tests/test_health.py @@ -39,6 +39,10 @@ "GIT_SHA", "BIOMNI_GIT_REF", "GIT_REF", + "BIOMNI_LLM_PROXY_URL", + "BIOMNI_LLM_PROXY_SCHEMA", + "BIOMNI_LLM_PROXY_API_KEY", + "BIOMNI_TRUST_AUTH_HEADERS", ) @@ -68,6 +72,48 @@ def test_build_info_reads_git_env(monkeypatch): # --------------------------------------------------------------------------- +def test_behind_the_llm_proxy_its_token_is_the_credential(monkeypatch, tmp_path): + monkeypatch.setenv("BIOMNI_PATH", str(tmp_path)) + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + monkeypatch.setenv("BIOMNI_LLM_PROXY_API_KEY", "proxy-token") + monkeypatch.setenv("BIOMNI_TRUST_AUTH_HEADERS", "true") + ok, report = health.readiness_report() + assert ok is True + assert report["checks"]["llm_credential"] == {"ok": True, "source": "BIOMNI_LLM_PROXY_URL"} + + +def test_a_proxy_that_will_refuse_every_call_is_not_ready(monkeypatch, tmp_path): + """Without the gateway's ids the proxy refuses every request, so the pod cannot serve.""" + monkeypatch.setenv("BIOMNI_PATH", str(tmp_path)) + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + monkeypatch.setenv("BIOMNI_LLM_PROXY_API_KEY", "proxy-token") + ok, report = health.readiness_report() + assert ok is False + assert "BIOMNI_TRUST_AUTH_HEADERS" in report["checks"]["llm_credential"]["detail"] + + +def test_a_leftover_provider_key_does_not_make_a_tokenless_proxy_ready(monkeypatch, tmp_path): + """Every model call goes through the proxy, so without its token none can succeed.""" + monkeypatch.setenv("BIOMNI_PATH", str(tmp_path)) + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-whatever") + ok, report = health.readiness_report() + assert ok is False + check = report["checks"]["llm_credential"] + assert check["ok"] is False + assert "BIOMNI_LLM_PROXY_API_KEY" in check["detail"] + + +def test_a_misspelt_proxy_schema_makes_the_pod_unready(monkeypatch, tmp_path): + monkeypatch.setenv("BIOMNI_PATH", str(tmp_path)) + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + monkeypatch.setenv("BIOMNI_LLM_PROXY_API_KEY", "proxy-token") + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", "openia") + ok, report = health.readiness_report() + assert ok is False + assert "BIOMNI_LLM_PROXY_SCHEMA" in report["checks"]["llm_credential"]["detail"] + + def test_readiness_ok_when_data_and_credential_present(monkeypatch, tmp_path): monkeypatch.setenv("BIOMNI_PATH", str(tmp_path)) monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-whatever") @@ -86,12 +132,22 @@ def test_readiness_fails_without_credential(monkeypatch, tmp_path): assert report["checks"]["llm_credential"]["ok"] is False -def test_readiness_fails_without_data_dir(monkeypatch, tmp_path): - monkeypatch.setenv("BIOMNI_PATH", str(tmp_path / "does-not-exist")) +@pytest.mark.parametrize("env", ["BIOMNI_USER_DATA_PATH", "BIOMNI_DATA_PATH"]) +def test_readiness_waits_for_a_configured_workspace(monkeypatch, tmp_path, env): + monkeypatch.setenv(env, str(tmp_path / "does-not-exist")) monkeypatch.setenv("OPENAI_API_KEY", "sk-whatever") ok, report = health.readiness_report() assert ok is False - assert report["checks"]["data_path"]["ok"] is False + assert report["checks"]["data_path"] == {"ok": False, "path": str(tmp_path / "does-not-exist")} + + +def test_readiness_does_not_wait_for_the_apps_own_data_directory(monkeypatch, tmp_path): + """The image sets BIOMNI_PATH to a directory only a chat creates - and an unready pod gets no chats.""" + monkeypatch.setenv("BIOMNI_PATH", str(tmp_path / "created-on-first-use")) + monkeypatch.setenv("OPENAI_API_KEY", "sk-whatever") + ok, report = health.readiness_report() + assert ok is True + assert report["checks"]["data_path"] == {"ok": True, "path": None} def test_readiness_accepts_custom_base_url(monkeypatch, tmp_path): diff --git a/tests/test_identity.py b/tests/test_identity.py index 637a25182..353f6084d 100644 --- a/tests/test_identity.py +++ b/tests/test_identity.py @@ -8,13 +8,18 @@ from __future__ import annotations +import asyncio +import threading + import pytest from biomni import identity +from biomni.observability import capture_context _ENV_TO_CLEAR = ( "BIOMNI_AUTH_USER_ID_HEADER", "BIOMNI_AUTH_EMAIL_HEADER", "BIOMNI_AUTH_WORKSPACE_HEADER", + "BIOMNI_AUTH_ISSUER_HEADER", "BIOMNI_AUTH_USER_CONTEXT_HEADER", "BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER", "BIOMNI_TRUST_AUTH_HEADERS", @@ -71,6 +76,37 @@ def test_extract_headers_from_environ_with_nested_asgi_scope(): assert identity.extract_headers(environ)["x-user-id"] == "nested-1" +def test_the_raw_asgi_headers_win_over_the_environ_keys_derived_from_them(): + """socket.io folds "Ai_App_User_Context" into the key of the real header. + + A gateway that strips client copies of ``Ai-App-User-Context`` by name lets + the underscore spelling through, so the environ's HTTP_* keys are not the + gateway's word - the raw header list next to them is. + """ + environ = { + "HTTP_AI_APP_USER_CONTEXT": "sub=victim", + "asgi.scope": { + "headers": [(b"ai_app_user_context", b"sub=victim"), (b"ai-app-user-context", b"sub=real")], + }, + } + headers = identity.extract_headers(environ) + assert headers["ai-app-user-context"] == "sub=real" + who = identity.identity_from_headers(environ) + assert who is not None and who.user_id == "real" + + +def test_an_asgi_scope_yields_only_its_headers(): + scope = {"type": "websocket", "path": "/ws/socket.io/", "headers": [(b"x-user-id", b"u1")]} + assert identity.extract_headers(scope) == {"x-user-id": "u1"} + + +def test_a_repeated_header_resolves_to_its_last_value(): + """A gateway that appends its own copy puts it last; both carriers agree.""" + scope = {"headers": [(b"ai-app-user-context", b"sub=client"), (b"ai-app-user-context", b"sub=gateway")]} + who = identity.identity_from_headers(scope) + assert who is not None and who.user_id == "gateway" + + def test_extract_headers_tolerates_junk(): assert identity.extract_headers(None) == {} assert identity.extract_headers(object()) == {} @@ -117,6 +153,28 @@ def test_parse_context_header_survives_junk(raw): assert identity.parse_context_header(raw) == {} +def test_a_whole_quoted_field_is_unquoted(): + """How a standard CSV writer quotes a field that contains a comma.""" + fields = identity.parse_context_header('"family_name=Smith, Jr.",sub=s1') + assert fields == {"family_name": "Smith, Jr.", "sub": "s1"} + + +def test_an_unterminated_quote_rejects_the_whole_header(): + """An unescaped value must not swallow - or supply - the fields after it.""" + assert identity.parse_context_header('email=a@b.org,family_name="M,sub=victim') == {} + + +def test_an_identity_key_given_twice_rejects_the_whole_header(): + """What an unescaped family name smuggling in a ``sub`` of its own looks like.""" + assert identity.parse_context_header('family_name=M,sub=victim,x=",sub=real') == {} + assert identity.parse_context_header("sub=a,sub=b") == {} + assert identity.parse_context_header("uuid=w1,uuid=w2") == {} + + +def test_a_display_field_given_twice_keeps_the_identity(): + assert identity.parse_context_header("sub=s1,given_name=A,given_name=B") == {"sub": "s1", "given_name": "B"} + + def test_parse_context_header_keeps_the_good_fields_from_a_malformed_header(): """One broken fragment must not cost us the identity.""" fields = identity.parse_context_header("garbage,sub=s1,=,email=a@b.org") @@ -192,14 +250,34 @@ def test_context_headers_win_over_single_value_headers(): assert (who.user_id, who.workspace_id) == ("grip-1", "grip-ws") -def test_single_value_headers_fill_gaps_the_context_left(): +def test_single_value_headers_never_fill_gaps_the_context_left(): + """The gateway strips its own headers, not every name another gateway uses. + + Anything it left out of the context header therefore comes from the client + if read from elsewhere - an issuer that re-keys storage, or a workspace that + another workspace's quota pays for. + """ who = identity.identity_from_headers( - {"ai-app-user-context": "sub=grip-1", "x-auth-request-email": "a@b.org", "x-workspace-id": "proxy-ws"} + { + "ai-app-user-context": "sub=grip-1", + "x-auth-request-email": "a@b.org", + "x-auth-issuer": "https://attacker.example", + "x-workspace-id": "proxy-ws", + } ) assert who is not None assert who.user_id == "grip-1" - assert who.email == "a@b.org" - assert who.workspace_id == "proxy-ws" + assert (who.email, who.issuer, who.workspace_id) == (None, None, None) + + +def test_a_workspace_context_alone_still_silences_the_single_value_headers(): + who = identity.identity_from_headers({"ai-app-workspace-context": "uuid=w1", "x-user-id": "victim"}) + assert who is None + + +def test_a_malformed_context_header_does_not_fall_back_to_single_value_headers(): + who = identity.identity_from_headers({"ai-app-user-context": "sub=a,sub=b", "x-user-id": "victim"}) + assert who is None def test_context_header_names_are_configurable(monkeypatch): @@ -528,3 +606,281 @@ def test_password_identity_rejects_an_unusable_username(): def test_password_identity_is_stable_across_calls(): assert identity.password_identity("alice").scoped_key() == identity.password_identity("ALICE").scoped_key() + + +# --------------------------------------------------------------------------- # +# The GRIP team's configuration: one CSV header per context +# --------------------------------------------------------------------------- # + + +def test_a_composite_header_named_as_a_single_value_header_is_still_parsed(monkeypatch): + """Pointing the single-value overrides at the context header must not break anything. + + That is what the GRIP team configured. With ``sub`` present it was harmless; + with ``sub`` missing, the whole ``email=...,given_name=...`` string used to + become the user id. + """ + monkeypatch.setenv("BIOMNI_AUTH_USER_ID_HEADER", "Ai-App-User-Context") + monkeypatch.setenv("BIOMNI_AUTH_EMAIL_HEADER", "Ai-App-User-Context") + monkeypatch.setenv("BIOMNI_AUTH_WORKSPACE_HEADER", "Ai-App-Workspace-Context") + who = identity.identity_from_headers( + { + "Ai-App-User-Context": "email=jane.doe@grip.org,given_name=Jane, family_name=Doe,sub=", + "Ai-App-Workspace-Context": GRIP_WORKSPACE_CONTEXT, + } + ) + assert who is not None + assert who.user_id is None + assert who.email == "jane.doe@grip.org" + assert who.workspace_id == "ws-7a3d-4e21" + + +def test_the_overrides_the_grip_team_set_are_reported_as_moot(monkeypatch): + monkeypatch.setenv("BIOMNI_AUTH_USER_ID_HEADER", "Ai-App-User-Context") + monkeypatch.setenv("BIOMNI_AUTH_WORKSPACE_HEADER", "X-Workspace-Id, ai-app-workspace-context") + warnings = identity.auth_config_warnings() + assert len(warnings) == 2 + assert all("has no effect" in w and "Remove the variable" in w for w in warnings) + assert "BIOMNI_AUTH_USER_ID_HEADER" in warnings[0] + + +def test_ordinary_overrides_are_not_reported(monkeypatch): + monkeypatch.setenv("BIOMNI_AUTH_USER_ID_HEADER", "X-Remote-User") + assert identity.auth_config_warnings() == [] + + +def test_a_utf8_name_decoded_as_latin1_is_repaired(): + """Starlette decodes header bytes as latin-1; a gateway sends UTF-8.""" + mangled = "given_name=João,sub=s1".encode().decode("latin-1") + who = identity.identity_from_headers({"Ai-App-User-Context": mangled}) + assert who is not None + assert who.given_name == "João" + + +def test_a_utf8_name_in_raw_asgi_bytes_is_repaired(): + scope = {"headers": [(b"ai-app-user-context", "given_name=Zoë,family_name=Brontë,sub=s1".encode())]} + who = identity.identity_from_headers(scope) + assert who is not None + assert (who.given_name, who.family_name) == ("Zoë", "Brontë") + + +def test_genuine_latin1_text_is_left_alone(): + # "Müller" in latin-1 is not valid UTF-8, so it must come through untouched. + assert identity._repair_utf8("Müller") == "Müller" + assert identity._repair_utf8("plain ascii") == "plain ascii" + + +def test_describe_auth_config_names_the_headers_in_force(monkeypatch): + monkeypatch.setenv("BIOMNI_TRUST_AUTH_HEADERS", "true") + report = identity.describe_auth_config() + assert report["trust_auth_headers"] is True + assert report["user_context_headers"] == ["ai-app-user-context"] + assert report["workspace_context_headers"] == ["ai-app-workspace-context"] + assert report["anonymous_persistence"] is False + + +# --------------------------------------------------------------------------- # +# Carrying an identity through a login +# --------------------------------------------------------------------------- # + + +def _grip_identity(**overrides) -> identity.UserIdentity: + who = identity.identity_from_headers(GRIP_HEADERS) + assert who is not None + return identity.UserIdentity(**{**who.__dict__, **overrides}) + + +def test_metadata_round_trips_every_field(): + who = _grip_identity(issuer="https://idp.example") + assert identity.UserIdentity.from_metadata(who.to_metadata()) == who + + +@pytest.mark.parametrize( + "metadata", + [ + None, + "not a mapping", + # Written before user_id was carried: reading it would key the user by + # e-mail instead of subject and show them an empty history. + {"source": "headers", "email": "jane.doe@grip.org", "workspace_id": "ws"}, + {"source": "somewhere-else", "user_id": "u1"}, + ], +) +def test_metadata_that_is_not_an_identity_record_is_rejected(metadata): + assert identity.UserIdentity.from_metadata(metadata) is None + + +# --------------------------------------------------------------------------- # +# One identity per session +# --------------------------------------------------------------------------- # + + +def _headers_for(sub: str, workspace: str = "ws-7a3d-4e21") -> dict[str, str]: + return { + "Ai-App-User-Context": f"email={sub}@grip.org,given_name={sub.title()},sub={sub}", + "Ai-App-Workspace-Context": f"uuid={workspace}", + } + + +def _login_for(sub: str, workspace: str = "ws-7a3d-4e21") -> identity.Login: + who = identity.identity_from_headers(_headers_for(sub, workspace)) + assert who is not None + return identity.Login(who.scoped_key(), who.to_metadata()) + + +@pytest.fixture +def trusted(monkeypatch): + monkeypatch.setenv("BIOMNI_TRUST_AUTH_HEADERS", "true") + + +def test_behind_the_gateway_its_headers_decide(trusted): + session = identity.resolve_session_identity(None, _headers_for("alice"), session_id="s1") + assert session.identity.user_id == "alice" + assert session.identity.source == "headers" + assert not session.conflict and session.notice is None + + +def test_a_login_that_agrees_with_the_gateway_is_accepted(trusted): + session = identity.resolve_session_identity(_login_for("alice"), _headers_for("alice"), session_id="s1") + assert session.identity.user_id == "alice" and not session.conflict + + +def test_a_login_for_someone_else_is_a_conflict(trusted): + """A shared browser: the platform now names bob, the old cookie names alice.""" + session = identity.resolve_session_identity(_login_for("alice"), _headers_for("bob"), session_id="s1") + assert session.conflict + assert not session.identity.is_authenticated + assert "Bob" in (session.notice or "") + assert "reload the page" in (session.notice or "") + + +def test_a_login_for_another_workspace_is_a_conflict(trusted): + session = identity.resolve_session_identity(_login_for("alice", "ws-1"), _headers_for("alice", "ws-2")) + assert session.conflict + + +def test_a_login_from_before_the_gateway_was_trusted_is_a_conflict(trusted): + anonymous = identity.anonymous_identity("old") + stale = identity.Login(anonymous.scoped_key(), anonymous.to_metadata()) + assert identity.resolve_session_identity(stale, _headers_for("alice")).conflict + + +# The metadata an earlier release wrote to its login cookies, which carries no +# user_id. The shared "local" login is what every browser of a single-user +# deployment holds; honoured, it would open every user's session onto one history. +_RELEASED_LOCAL_LOGIN = identity.Login( + "local", {"source": "local", "email": "", "workspace_id": "", "given_name": "", "family_name": ""} +) + + +def test_a_login_cookie_from_the_previous_release_is_a_conflict(trusted): + assert identity.resolve_session_identity(_RELEASED_LOCAL_LOGIN, _headers_for("alice")).conflict + assert not identity.login_is_current(_RELEASED_LOCAL_LOGIN, _headers_for("alice")) + + +def test_a_gateway_login_from_the_previous_release_is_a_conflict(trusted): + """Same person, but the cookie cannot say so: it predates the raw subject id.""" + current = _login_for("alice") + released = identity.Login( + current.identifier, {k: v for k, v in current.metadata.items() if k not in {"user_id", "issuer"}} + ) + assert identity.resolve_session_identity(released, _headers_for("alice")).conflict + + +def test_a_login_whose_identifier_and_metadata_disagree_is_a_conflict(trusted): + bob = _login_for("bob") + forged = identity.Login(_login_for("alice").identifier, bob.metadata) + assert identity.resolve_session_identity(forged, _headers_for("bob")).conflict + assert identity.resolve_session_identity(identity.Login(None, bob.metadata), _headers_for("bob")).conflict + + +def test_a_login_that_cannot_say_who_it_is_does_not_stand_in(trusted): + session = identity.resolve_session_identity(_RELEASED_LOCAL_LOGIN, None, session_id="s1") + assert not session.identity.is_authenticated + assert session.notice == identity.GATEWAY_SILENT_NOTICE + + +def test_the_login_stands_in_when_the_websocket_carries_no_headers(trusted): + """Not every gateway forwards its headers on the websocket upgrade.""" + session = identity.resolve_session_identity(_login_for("alice"), None, session_id="s1") + assert session.identity.user_id == "alice" + assert session.identity.workspace_id == "ws-7a3d-4e21" + + +def test_no_headers_and_no_login_is_anonymous_with_a_reason(trusted): + session = identity.resolve_session_identity(None, None, session_id="s1") + assert not session.identity.is_authenticated + assert session.notice == identity.GATEWAY_SILENT_NOTICE + + +def test_without_trust_gateway_headers_are_ignored_and_the_reason_is_given(): + """The GRIP symptom: headers arrive, nobody is detected, and now it says why.""" + session = identity.resolve_session_identity(None, GRIP_HEADERS, session_id="s1") + assert not session.identity.is_authenticated + assert session.notice == identity.HEADERS_IGNORED_NOTICE + assert "BIOMNI_TRUST_AUTH_HEADERS" in session.notice + + +def test_without_a_gateway_a_password_login_is_the_identity(): + jane = identity.password_identity("Jane") + login = identity.Login(jane.scoped_key(), jane.to_metadata()) + session = identity.resolve_session_identity(login, None, session_id="s1") + assert session.identity.source == "password" + assert session.identity.is_authenticated + + +def test_without_a_gateway_a_header_login_is_not_trusted(): + session = identity.resolve_session_identity(_login_for("alice"), None, session_id="s1") + assert not session.identity.is_authenticated + assert session.notice == identity.NO_GATEWAY_NOTICE + + +def test_single_user_mode_gives_the_shared_local_identity(monkeypatch): + monkeypatch.setenv("BIOMNI_ALLOW_ANONYMOUS_PERSISTENCE", "true") + session = identity.resolve_session_identity(None, None, session_id="s1") + assert session.identity.source == "local" + + +def test_login_is_current_only_checks_behind_a_trusted_gateway(monkeypatch): + assert identity.login_is_current(_login_for("alice"), _headers_for("bob")) + monkeypatch.setenv("BIOMNI_TRUST_AUTH_HEADERS", "true") + assert not identity.login_is_current(_login_for("alice"), _headers_for("bob")) + assert identity.login_is_current(_login_for("alice"), _headers_for("alice")) + # A request the gateway did not annotate cannot contradict anything. + assert identity.login_is_current(_login_for("alice"), {}) + + +# --------------------------------------------------------------------------- # +# The bound identity follows the work into threads +# --------------------------------------------------------------------------- # + + +def test_bound_identity_restores_what_was_bound_before(): + outer, inner = _grip_identity(), _grip_identity(user_id="someone-else") + with identity.bound_identity(outer): + with identity.bound_identity(inner): + assert identity.current_identity() == inner + assert identity.current_identity() == outer + assert identity.current_identity() is None + + +def test_the_identity_reaches_a_worker_thread_through_capture_context(): + """How run_in_executor and run_with_timeout hand work to threads.""" + seen = [] + who = _grip_identity() + with identity.bound_identity(who): + ctx = capture_context() + thread = threading.Thread(target=ctx.run, args=(lambda: seen.append(identity.current_identity()),)) + thread.start() + thread.join() + assert seen == [who] + + +def test_the_identity_reaches_asyncio_to_thread(): + who = _grip_identity() + + async def main(): + with identity.bound_identity(who): + return await asyncio.to_thread(identity.current_identity) + + assert asyncio.run(main()) == who diff --git a/tests/test_live_runs.py b/tests/test_live_runs.py index dda64bfe9..2ca0a45d8 100644 --- a/tests/test_live_runs.py +++ b/tests/test_live_runs.py @@ -11,13 +11,19 @@ from chainlit_ui.live_runs import LiveRuns +class FakeUser: + def __init__(self, identifier: str) -> None: + self.identifier = identifier + + class FakeSession: - """A stand-in for chainlit's WebsocketSession - only emit matters here.""" + """A stand-in for chainlit's WebsocketSession: its emitters and its login.""" - def __init__(self, name: str) -> None: + def __init__(self, name: str, user: str = "alice") -> None: self.name = name self.emit = f"emit-{name}" self.emit_call = f"emit_call-{name}" + self.user = FakeUser(user) @pytest.fixture @@ -114,6 +120,14 @@ def __setattr__(self, name: str, value: object) -> None: assert runs.is_live("t1") +def test_a_run_is_never_handed_to_another_users_connection(runs: LiveRuns) -> None: + running = FakeSession("running", user="alice") + runs.register("t1", running) + + assert runs.reattach("t1", FakeSession("intruder", user="bob")) is False + assert running.emit == "emit-running" + + def test_reattach_needs_both_a_thread_and_a_session(runs: LiveRuns) -> None: runs.register("t1", FakeSession("running")) assert runs.reattach(None, FakeSession("reopened")) is False @@ -142,16 +156,25 @@ def test_cancel_stops_the_task_that_is_doing_the_work(runs: LiveRuns) -> None: task = FakeTask() runs.register("t1", FakeSession("running"), task) - assert runs.cancel("t1") is True + assert runs.cancel("t1", FakeSession("reopened")) is True assert task.cancelled +def test_only_the_user_who_started_a_run_can_stop_it(runs: LiveRuns) -> None: + task = FakeTask() + runs.register("t1", FakeSession("running", user="alice"), task) + + assert runs.cancel("t1", FakeSession("intruder", user="bob")) is False + assert runs.cancel("t1", None) is False + assert not task.cancelled + + def test_cancel_is_false_with_nothing_to_stop(runs: LiveRuns) -> None: - assert runs.cancel("t1") is False - assert runs.cancel(None) is False + assert runs.cancel("t1", FakeSession("a")) is False + assert runs.cancel(None, FakeSession("a")) is False runs.register("t2", FakeSession("no-task")) - assert runs.cancel("t2") is False + assert runs.cancel("t2", FakeSession("a")) is False def test_a_finished_task_is_not_cancelled_again(runs: LiveRuns) -> None: @@ -159,7 +182,7 @@ def test_a_finished_task_is_not_cancelled_again(runs: LiveRuns) -> None: task = FakeTask(done=True) runs.register("t1", FakeSession("running"), task) - assert runs.cancel("t1") is False + assert runs.cancel("t1", FakeSession("reopened")) is False assert not task.cancelled diff --git a/tests/test_llm_callsite_resilience.py b/tests/test_llm_callsite_resilience.py index 4c8f0bd00..031e28928 100644 --- a/tests/test_llm_callsite_resilience.py +++ b/tests/test_llm_callsite_resilience.py @@ -51,8 +51,11 @@ def test_write_python_code_routes_through_get_llm(monkeypatch: pytest.MonkeyPatc assert "print(1 + 1)" in out assert len(calls) == 1, "write_python_code must build its model via get_llm exactly once" args, kwargs = calls[0]["args"], calls[0]["kwargs"] - assert args[0] == "claude-3-5-sonnet-20240620" - assert kwargs["source"] == "Anthropic" + # No pinned model or provider: the deployment's configured model (BIOMNI_LLM) + # comes from the config. A pinned id is eventually retired, and one the + # deployment never set up is one its key or LLM proxy cannot reach. + assert not args + assert "model" not in kwargs and "source" not in kwargs assert isinstance(kwargs["config"], BiomniConfig) @@ -94,8 +97,9 @@ def test_retriever_fallback_routes_through_get_llm(monkeypatch: pytest.MonkeyPat assert len(calls) == 1, "the llm=None fallback must build its model via get_llm" args, kwargs = calls[0]["args"], calls[0]["kwargs"] - assert args[0] == "gpt-4o" - assert kwargs["source"] == "OpenAI" + # The configured model, not a pinned provider the deployment may have no key for. + assert not args + assert "model" not in kwargs and "source" not in kwargs assert isinstance(kwargs["config"], BiomniConfig) # Routing through get_llm must not change selection behavior. assert selected["tools"] == [{"name": "t0", "description": "d0"}] diff --git a/tests/test_llm_factory.py b/tests/test_llm_factory.py index e98d81a6f..0f078b480 100644 --- a/tests/test_llm_factory.py +++ b/tests/test_llm_factory.py @@ -29,6 +29,10 @@ def _clear_llm_env(monkeypatch: pytest.MonkeyPatch) -> None: "AZURE_OPENAI_API_KEY", "GEMINI_API_KEY", "GROQ_API_KEY", + # A configured proxy outranks every other source (tests/test_llm_proxy.py). + "BIOMNI_LLM_PROXY_URL", + "BIOMNI_LLM_PROXY_SCHEMA", + "BIOMNI_LLM_PROXY_API_KEY", ): monkeypatch.delenv(k, raising=False) @@ -155,6 +159,23 @@ def test_get_llm_inherits_config_resilience(monkeypatch: pytest.MonkeyPatch) -> assert kw["default_request_timeout"] == 7.5 +@pytest.mark.parametrize( + ("model", "env_name"), [("gemini-1.5-pro", "GEMINI_API_KEY"), ("some-groq-mixtral", "GROQ_API_KEY")] +) +def test_an_openai_compatible_provider_never_gets_the_openai_key(monkeypatch, model, env_name) -> None: + """Without its own key ChatOpenAI would read OPENAI_API_KEY - and send it to Google or Groq.""" + captured = _install_fake_provider_module(monkeypatch, "langchain_openai", "ChatOpenAI") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai") + + with pytest.raises(ValueError, match=env_name): + get_llm(model) + assert not captured + + monkeypatch.setenv(env_name, "provider-key") + get_llm(model) + assert captured[-1]["api_key"].get_secret_value() == "provider-key" + + def test_get_llm_config_none_timeout_disables(monkeypatch: pytest.MonkeyPatch) -> None: """Explicit None on config means disable per-call timeout (vs. default-fill).""" captured = _install_fake_provider_module(monkeypatch, "langchain_anthropic", "ChatAnthropic") diff --git a/tests/test_llm_failures.py b/tests/test_llm_failures.py new file mode 100644 index 000000000..3a303cc6a --- /dev/null +++ b/tests/test_llm_failures.py @@ -0,0 +1,170 @@ +"""Tests for chainlit_ui.llm_failures - what a user reads when a model call fails. + +Most exceptions are stand-ins shaped like the provider SDKs' (a ``status_code`` +and a JSON ``body``), since the module is duck-typed on exactly that. The last +section raises the real thing, through each SDK's own client, for the bodies +proxies actually send - skipped for an SDK that is not installed. +""" + +from __future__ import annotations + +import pytest +from biomni.llm_proxy import LLMProxyIdentityError +from chainlit_ui.llm_failures import describe_llm_failure, is_access_failure + + +class APIStatusError(Exception): + def __init__(self, status_code: int, message: str = "", body: object = None) -> None: + super().__init__(f"Error code: {status_code} - {body!r}") + self.status_code = status_code + self.message = message + self.body = body + + +class APIConnectionError(Exception): + pass + + +class ConnectTimeout(APIConnectionError): + pass + + +@pytest.fixture +def behind_proxy(monkeypatch): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + + +@pytest.fixture(autouse=True) +def _direct(monkeypatch): + monkeypatch.delenv("BIOMNI_LLM_PROXY_URL", raising=False) + monkeypatch.delenv("BIOMNI_LLM_PROXY_SCHEMA", raising=False) + + +def test_a_missing_identity_is_explained_in_its_own_words(): + exc = LLMProxyIdentityError("the LLM proxy needs the user and workspace every model call is made for, and x") + assert ( + describe_llm_failure(exc) == "The LLM proxy needs the user and workspace every model call is made for, and x." + ) + assert is_access_failure(exc) + + +def test_the_servers_message_is_shown_rather_than_the_raw_body(behind_proxy): + exc = APIStatusError(403, body={"error": {"message": "workspace ws-1 has no access to model gpt-4.1"}}) + text = describe_llm_failure(exc) + assert ( + text + == "The platform's LLM proxy refused the request (HTTP 403): workspace ws-1 has no access to model gpt-4.1." + ) + assert "{" not in text + + +def test_a_rejected_token_names_the_setting_behind_the_proxy(behind_proxy): + assert "BIOMNI_LLM_PROXY_API_KEY" in describe_llm_failure(APIStatusError(401)) + + +def test_a_rejected_key_without_the_proxy_names_the_api_key(): + text = describe_llm_failure(APIStatusError(401)) + assert text.startswith("The language model service rejected") + assert "API key" in text + + +def test_a_usage_limit_explains_the_platforms_per_user_limits(behind_proxy): + text = describe_llm_failure(APIStatusError(429)) + assert "per user and per workspace" in text and "Try again later" in text + + +def test_an_unknown_model_points_at_the_model_setting(behind_proxy): + assert "BIOMNI_LLM" in describe_llm_failure(APIStatusError(404, body={"message": "no such deployment"})) + + +def test_a_server_error_suggests_trying_again(): + assert describe_llm_failure(APIStatusError(502)).endswith("Try again shortly.") + + +def test_a_connection_failure_is_recognised_through_its_base_class(behind_proxy): + exc = ConnectTimeout("timed out") + assert describe_llm_failure(exc) == "Could not reach the platform's LLM proxy. Try again shortly." + assert is_access_failure(exc) + + +def test_anything_else_still_says_what_was_raised(): + assert describe_llm_failure(KeyError("messages")) == "KeyError: 'messages'." + assert describe_llm_failure(RuntimeError()) == "RuntimeError." + + +def test_a_very_long_server_message_is_truncated(): + exc = APIStatusError(400, body={"error": {"message": "x" * 5000}}) + assert len(describe_llm_failure(exc)) < 400 + + +@pytest.mark.parametrize( + "status,unreachable", [(400, False), (401, True), (403, True), (404, True), (429, True), (500, True)] +) +def test_which_failures_end_the_run(status, unreachable): + assert is_access_failure(APIStatusError(status)) is unreachable + + +def test_an_ordinary_bug_is_not_an_access_failure(): + assert not is_access_failure(ValueError("bad plan")) + + +# --------------------------------------------------------------------------- # +# The real SDKs' exceptions +# --------------------------------------------------------------------------- # + + +def _raised_by(sdk_name: str, status: int, content: bytes, content_type: str) -> Exception: + """What the SDK raises when the proxy answers with this response.""" + sdk = pytest.importorskip(sdk_name) + httpx = pytest.importorskip("httpx") + + def respond(request): + return httpx.Response(status, content=content, headers={"content-type": content_type}) + + http_client = httpx.Client(transport=httpx.MockTransport(respond)) + messages = [{"role": "user", "content": "hi"}] + with pytest.raises(sdk.APIStatusError) as info: + if sdk_name == "openai": + client = sdk.OpenAI( + api_key="t", base_url="https://proxy.example/v1", http_client=http_client, max_retries=0 + ) + client.chat.completions.create(model="m", messages=messages) + else: + client = sdk.Anthropic( + api_key="t", base_url="https://proxy.example", http_client=http_client, max_retries=0 + ) + client.messages.create(model="m", max_tokens=1, messages=messages) + return info.value + + +_SDKS = ["openai", "anthropic"] + + +@pytest.mark.parametrize("sdk_name", _SDKS) +def test_a_detail_body_is_read(behind_proxy, sdk_name): + exc = _raised_by(sdk_name, 403, b'{"detail": "workspace ws-1 is over its limit"}', "application/json") + assert describe_llm_failure(exc) == ( + "The platform's LLM proxy refused the request (HTTP 403): workspace ws-1 is over its limit." + ) + + +@pytest.mark.parametrize("sdk_name", _SDKS) +def test_a_plain_string_error_body_is_read(behind_proxy, sdk_name): + exc = _raised_by(sdk_name, 400, b'{"error": "model m is not enabled"}', "application/json") + assert ( + describe_llm_failure(exc) == "The platform's LLM proxy rejected the request (HTTP 400): model m is not enabled." + ) + + +@pytest.mark.parametrize("sdk_name", _SDKS) +def test_an_html_error_page_is_not_pasted_into_the_chat(behind_proxy, sdk_name): + page = b"502 Bad Gateway
nginx
" + exc = _raised_by(sdk_name, 502, page, "text/html") + assert describe_llm_failure(exc) == "The platform's LLM proxy failed to answer (HTTP 502). Try again shortly." + + +@pytest.mark.parametrize("sdk_name", _SDKS) +def test_the_providers_own_error_shape_is_read(behind_proxy, sdk_name): + body = b'{"type": "error", "error": {"type": "invalid_request_error", "message": "prompt is too long"}}' + exc = _raised_by(sdk_name, 400, body, "application/json") + assert describe_llm_failure(exc) == "The platform's LLM proxy rejected the request (HTTP 400): prompt is too long." diff --git a/tests/test_llm_proxy.py b/tests/test_llm_proxy.py new file mode 100644 index 000000000..062a4e478 --- /dev/null +++ b/tests/test_llm_proxy.py @@ -0,0 +1,332 @@ +"""Tests for the platform LLM proxy: settings, the identity rule, and the wire. + +The first half is pure Python (``biomni.llm_proxy``). The second half runs the +real chat-model classes against a local stub of the proxy and asserts on what +actually crosses the wire - the headers the proxy refuses to go without, the +bearer token, and the absence of any direct provider credential - because that +is the contract the platform states, and mocks of the SDKs would only restate +our assumptions about them. +""" + +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import pytest +from biomni import llm_proxy +from biomni.identity import UserIdentity, bound_identity, password_identity +from biomni.llm_proxy import LLMProxyIdentityError + +_PROXY_ENV = ("BIOMNI_LLM_PROXY_URL", "BIOMNI_LLM_PROXY_SCHEMA", "BIOMNI_LLM_PROXY_API_KEY") + +GRIP_USER = UserIdentity( + user_id="6f1c0e9e-1111-2222", + email="jane.doe@grip.org", + workspace_id="ws-7a3d-4e21", + given_name="Jane", + family_name="Doe", + source="headers", +) + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch: pytest.MonkeyPatch) -> None: + for name in (*_PROXY_ENV, "LLM_SOURCE", "BIOMNI_SOURCE", "BIOMNI_TRUST_AUTH_HEADERS"): + monkeypatch.delenv(name, raising=False) + + +# --------------------------------------------------------------------------- # +# Settings +# --------------------------------------------------------------------------- # + + +def test_no_proxy_url_means_no_proxy(): + assert llm_proxy.llm_proxy_settings() is None + assert llm_proxy.describe_llm_proxy() == {"enabled": False} + + +@pytest.mark.parametrize( + "url", + ["https://proxy.example", "https://proxy.example/", "https://proxy.example/v1", "https://proxy.example/v1/"], +) +def test_the_root_and_the_v1_base_mean_the_same_proxy(monkeypatch, url): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", url) + settings = llm_proxy.llm_proxy_settings() + assert settings is not None + assert settings.root_url == "https://proxy.example" + assert settings.openai_base_url == "https://proxy.example/v1" + assert settings.anthropic_base_url == "https://proxy.example" + + +def test_the_schema_defaults_to_openai_and_accepts_anthropic(monkeypatch): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example") + assert llm_proxy.llm_proxy_settings().schema == "openai" + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", " Anthropic ") + assert llm_proxy.llm_proxy_settings().schema == "anthropic" + + +def test_an_unknown_schema_fails_loudly(monkeypatch): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example") + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", "azure") + with pytest.raises(ValueError, match="BIOMNI_LLM_PROXY_SCHEMA"): + llm_proxy.llm_proxy_settings() + assert "error" in llm_proxy.describe_llm_proxy() + + +def test_describe_reports_the_token_without_revealing_it(monkeypatch): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + monkeypatch.setenv("BIOMNI_LLM_PROXY_API_KEY", "sk-secret-value") + report = llm_proxy.describe_llm_proxy() + assert report == {"enabled": True, "url": "https://proxy.example", "schema": "openai", "token_configured": True} + assert "sk-secret-value" not in json.dumps(report) + + +# --------------------------------------------------------------------------- # +# Who a call is for +# --------------------------------------------------------------------------- # + + +def test_a_gateway_identity_names_the_user_and_workspace(): + with bound_identity(GRIP_USER): + assert llm_proxy.proxy_identity_headers() == { + "X-User-Id": "6f1c0e9e-1111-2222", + "X-Workspace-Id": "ws-7a3d-4e21", + } + + +def test_a_call_outside_any_session_is_refused(): + with pytest.raises(LLMProxyIdentityError, match="not being made for any chat session") as info: + llm_proxy.proxy_identity_headers() + # Generated code reads this. It must be told to stay in its session, not + # handed the means to make up an identity of its own. + assert "bound_identity" not in str(info.value) + assert "thread pool" in str(info.value) + + +def test_behind_a_trusted_gateway_an_unidentified_session_is_told_to_sign_in_again(monkeypatch): + """Telling it to set the flag that is already set sends the operator in circles.""" + monkeypatch.setenv("BIOMNI_TRUST_AUTH_HEADERS", "true") + problem = llm_proxy.proxy_identity_problem(UserIdentity(user_id="session-1", source="anonymous")) + assert problem is not None + assert "did not identify this session" in problem and "Reload the page" in problem + assert "BIOMNI_TRUST_AUTH_HEADERS" not in problem + + +@pytest.mark.parametrize( + "who,expected", + [ + (password_identity("jane"), "not signed in through the platform's authentication gateway"), + (UserIdentity(source="anonymous"), "identity source: anonymous"), + (UserIdentity(user_id="u1", source="headers"), "the workspace id (uuid in Ai-App-Workspace-Context)"), + ( + UserIdentity(email="a@b.org", workspace_id="ws", source="headers"), + "the user id (sub in Ai-App-User-Context)", + ), + (UserIdentity(user_id="Jöhn", workspace_id="ws", source="headers"), "cannot be sent as HTTP headers"), + (UserIdentity(user_id="u1\r\nX-Evil: 1", workspace_id="ws", source="headers"), "cannot be sent"), + ], +) +def test_identities_the_proxy_cannot_bill_are_refused_with_the_reason(who, expected): + problem = llm_proxy.proxy_identity_problem(who) + assert problem is not None and expected in problem + with bound_identity(who), pytest.raises(LLMProxyIdentityError): + llm_proxy.proxy_identity_headers() + + +def test_a_complete_gateway_identity_has_no_problem(): + assert llm_proxy.proxy_identity_problem(GRIP_USER) is None + + +# --------------------------------------------------------------------------- # +# On the wire +# --------------------------------------------------------------------------- # + +_OPENAI_REPLY = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "stub", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_ANTHROPIC_REPLY = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "stub", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +class _StubProxy: + """A local HTTP server that answers like the proxy and records each request.""" + + def __init__(self) -> None: + self.requests: list[dict] = [] + recorded = self.requests + + class Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + body = json.loads(self.rfile.read(int(self.headers.get("Content-Length") or 0)) or b"{}") + recorded.append( + {"path": self.path, "headers": {k.lower(): v for k, v in self.headers.items()}, "body": body} + ) + reply = _ANTHROPIC_REPLY if self.path.endswith("/messages") else _OPENAI_REPLY + payload = json.dumps(reply).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args) -> None: + pass + + self._server = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self.url = f"http://127.0.0.1:{self._server.server_address[1]}" + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + + def __enter__(self) -> _StubProxy: + self._thread.start() + return self + + def __exit__(self, *exc) -> None: + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def proxy(monkeypatch): + """The stub proxy, configured the way GRIP configures the real one. + + Direct provider keys are set too, to prove none of them ever leaves the pod. + """ + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-direct-must-not-leak") + monkeypatch.setenv("OPENAI_API_KEY", "sk-openai-direct-must-not-leak") + monkeypatch.setenv("BIOMNI_LLM_PROXY_API_KEY", "proxy-token") + with _StubProxy() as stub: + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", f"{stub.url}/v1") + yield stub + + +def _assert_no_direct_credential(request: dict) -> None: + sent = json.dumps(request["headers"]) + assert "must-not-leak" not in sent + assert "x-api-key" not in request["headers"] + + +def test_openai_schema_calls_chat_completions_with_both_headers(proxy): + pytest.importorskip("langchain_openai") + from biomni.llm import get_llm + + llm = get_llm("gpt-4.1", max_retries=0) + with bound_identity(GRIP_USER): + assert llm.invoke("hi").content == "hello" + + (request,) = proxy.requests + assert request["path"] == "/v1/chat/completions" + assert request["headers"]["authorization"] == "Bearer proxy-token" + assert request["headers"]["x-user-id"] == "6f1c0e9e-1111-2222" + assert request["headers"]["x-workspace-id"] == "ws-7a3d-4e21" + assert request["body"]["model"] == "gpt-4.1" + _assert_no_direct_credential(request) + + +def test_openai_schema_async_calls_carry_the_headers_too(proxy): + pytest.importorskip("langchain_openai") + import asyncio + + from biomni.llm import get_llm + + llm = get_llm("gpt-4.1", max_retries=0) + + async def call(): + with bound_identity(GRIP_USER): + return await llm.ainvoke("hi") + + assert asyncio.run(call()).content == "hello" + assert proxy.requests[0]["headers"]["x-user-id"] == "6f1c0e9e-1111-2222" + + +def test_one_model_instance_names_whichever_session_is_calling(proxy): + """Instances outlive requests; the ids must come from the call, never the instance.""" + pytest.importorskip("langchain_openai") + from biomni.llm import get_llm + + llm = get_llm("gpt-4.1", max_retries=0) + other = UserIdentity(user_id="bob", workspace_id="ws-2", source="headers") + for who in (GRIP_USER, other): + with bound_identity(who): + llm.invoke("hi") + assert [r["headers"]["x-user-id"] for r in proxy.requests] == ["6f1c0e9e-1111-2222", "bob"] + assert [r["headers"]["x-workspace-id"] for r in proxy.requests] == ["ws-7a3d-4e21", "ws-2"] + + +def test_a_reasoning_model_is_sent_no_stop_or_temperature(proxy): + pytest.importorskip("langchain_openai") + from biomni.llm import get_llm + + llm = get_llm("gpt-5-mini", stop_sequences=[""], temperature=0.2, max_retries=0) + with bound_identity(GRIP_USER): + llm.invoke("hi") + body = proxy.requests[0]["body"] + assert "stop" not in body and "temperature" not in body + + +def test_a_call_with_no_session_is_refused_before_anything_is_sent(proxy): + pytest.importorskip("langchain_openai") + from biomni.llm import get_llm + + llm = get_llm("gpt-4.1", max_retries=0) + with pytest.raises(LLMProxyIdentityError): + llm.invoke("hi") + with bound_identity(password_identity("jane")), pytest.raises(LLMProxyIdentityError): + llm.invoke("hi") + assert proxy.requests == [] + + +def test_anthropic_schema_calls_messages_with_a_bearer_token_and_both_headers(proxy, monkeypatch): + pytest.importorskip("langchain_anthropic") + from biomni.llm import get_llm + + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", "anthropic") + llm = get_llm("claude-sonnet-4-5", stop_sequences=[""], max_retries=0) + with bound_identity(GRIP_USER): + assert llm.invoke("hi").content == "hello" + + (request,) = proxy.requests + assert request["path"] == "/v1/messages" + assert request["headers"]["authorization"] == "Bearer proxy-token" + assert request["headers"]["x-user-id"] == "6f1c0e9e-1111-2222" + assert request["headers"]["x-workspace-id"] == "ws-7a3d-4e21" + assert request["body"]["model"] == "claude-sonnet-4-5" + assert request["body"]["stop_sequences"] == [""] + _assert_no_direct_credential(request) + + +def test_the_proxy_outranks_any_provider_a_call_site_asks_for(proxy): + from biomni.llm import resolve_source + + assert resolve_source("claude-sonnet-4-5", "Anthropic", None) == "LLMProxy" + assert resolve_source("gpt-4o", None, "http://elsewhere/v1") == "LLMProxy" + + +def test_a_proxy_without_its_token_is_refused_at_construction(proxy, monkeypatch): + pytest.importorskip("langchain_openai") + from biomni.llm import get_llm + + monkeypatch.delenv("BIOMNI_LLM_PROXY_API_KEY") + with pytest.raises(ValueError, match="BIOMNI_LLM_PROXY_API_KEY"): + get_llm("gpt-4.1") + + +def test_llmproxy_as_a_source_needs_the_proxy_url(): + from biomni.llm import resolve_source + + with pytest.raises(ValueError, match="BIOMNI_LLM_PROXY_URL"): + resolve_source("gpt-4o", "LLMProxy", None) diff --git a/tests/test_llm_resilience.py b/tests/test_llm_resilience.py index 80998bd87..2f5705e2f 100644 --- a/tests/test_llm_resilience.py +++ b/tests/test_llm_resilience.py @@ -39,6 +39,19 @@ def test_supports_prompt_caching(source, expected) -> None: assert supports_prompt_caching(source) is expected +@pytest.mark.parametrize(("schema", "expected"), [("anthropic", True), ("openai", False), ("bogus", False)]) +def test_prompt_caching_behind_the_llm_proxy_follows_its_schema(monkeypatch, schema, expected) -> None: + """The proxy meters every input token, so a cacheable prompt must be cached.""" + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example") + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", schema) + assert supports_prompt_caching("LLMProxy") is expected + + +def test_the_llm_proxy_source_without_a_proxy_does_not_cache(monkeypatch) -> None: + monkeypatch.delenv("BIOMNI_LLM_PROXY_URL", raising=False) + assert supports_prompt_caching("LLMProxy") is False + + # --------------------------------------------------------------------------- # prepare_messages_for_cache # --------------------------------------------------------------------------- diff --git a/tests/test_login_guard.py b/tests/test_login_guard.py new file mode 100644 index 000000000..4ef713061 --- /dev/null +++ b/tests/test_login_guard.py @@ -0,0 +1,60 @@ +"""The login guard, in the real app: a cookie is only good for the user the gateway names. + +Chainlit signs a browser in once, from the gateway's headers, and then believes +its cookie for 15 days. The guard (chainlit_app.gateway_checked_user) checks the +cookie against the headers on every request, so a browser someone else used, or +a cookie the previous release issued to every visitor under one shared name, +cannot open another user's settings and conversations. + +Driven through Chainlit's own HTTP API by login_guard_app.py, in an interpreter +of its own. Skipped where the chainlit extra is not installed, as in CI. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("chainlit") +pytest.importorskip("aiosqlite") # the chat-history store the conversation checks use + +REPO = Path(__file__).resolve().parents[1] +DRIVER = Path(__file__).with_name("login_guard_app.py") +RESULT = "LOGIN_GUARD_RESULT " + + +@pytest.fixture(scope="module") +def checks(tmp_path_factory: pytest.TempPathFactory) -> list[list]: + root = tmp_path_factory.mktemp("login_guard") + approot = root / "approot" + shutil.copytree(REPO / ".chainlit", approot / ".chainlit") + for name in ("workspace", "outputs", "state"): + (root / name).mkdir() + env = {k: v for k, v in os.environ.items() if not k.startswith(("BIOMNI_", "CHAINLIT_"))} + env.update( + CHAINLIT_APP_ROOT=str(approot), + BIOMNI_TRUST_AUTH_HEADERS="true", + BIOMNI_DATA_PATH=str(root / "workspace"), + BIOMNI_OUTPUT_ROOT=str(root / "outputs"), + BIOMNI_STATE_DIR=str(root / "state"), + BIOMNI_LOG_FORMAT="text", + LOG_LEVEL="ERROR", + ) + done = subprocess.run( + [sys.executable, str(DRIVER)], env=env, cwd=approot, capture_output=True, text=True, timeout=300 + ) + results = [line for line in done.stdout.splitlines() if line.startswith(RESULT)] + assert done.returncode == 0 and results, f"the app did not run:\n{done.stdout[-3000:]}\n{done.stderr[-3000:]}" + return json.loads(results[-1][len(RESULT) :])["checks"] + + +def test_a_cookie_is_only_good_for_the_user_the_gateway_names(checks: list[list]) -> None: + failed = [f"{label}: got {got!r}, want {want!r}" for label, got, want in checks if got != want] + assert not failed, "\n".join(failed) + assert len(checks) == 18, "a check stopped running" diff --git a/tests/test_observability.py b/tests/test_observability.py index 94fa065fd..3b816b352 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -142,6 +142,20 @@ def test_json_includes_exception_traceback(): assert "ValueError: boom" in line["exc"] +def test_identity_header_names_are_not_treated_as_secrets(): + """They contain "AUTH", and scrubbing them hid the header names from every warning about them.""" + r = obs.Redactor.from_environ( + { + "BIOMNI_AUTH_USER_ID_HEADER": "Ai-App-User-Context", + "BIOMNI_AUTH_WORKSPACE_CONTEXT_HEADER": "Ai-App-Workspace-Context", + "BIOMNI_TRUST_AUTH_HEADERS": "yes-trust", + "CHAINLIT_AUTH_SECRET": "a-real-signing-secret", + } + ) + line = "Ai-App-User-Context and Ai-App-Workspace-Context, signed with a-real-signing-secret, yes-trust" + assert r.redact(line) == "Ai-App-User-Context and Ai-App-Workspace-Context, signed with [REDACTED], yes-trust" + + def test_json_redacts_secret_in_message(monkeypatch): monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-secretkey-ABCDEFGHIJKL") stream, _ = _capture() @@ -262,6 +276,35 @@ def test_capture_context_propagates_into_worker_thread(): assert replayed == {"session_id": "S9", "run_id": "R9"} +def test_thread_pool_tasks_run_in_the_context_that_submitted_them(): + """What generated code gets: its own pools stay in its session.""" + obs.carry_context_into_thread_pools() + with obs.bind_run(session_id="S1", run_id="R1"), ThreadPoolExecutor(max_workers=2) as ex: + submitted = ex.submit(obs.get_context).result() + mapped = list(ex.map(lambda _: obs.get_context(), range(3))) + assert submitted == {"session_id": "S1", "run_id": "R1"} + assert mapped == [{"session_id": "S1", "run_id": "R1"}] * 3 + + +def test_a_pool_thread_never_carries_one_sessions_context_into_the_next(): + """Pool threads outlive the task that started them and serve every session.""" + obs.carry_context_into_thread_pools() + with ThreadPoolExecutor(max_workers=1) as ex: + with obs.bind_run(session_id="alice"): + first = ex.submit(obs.get_context).result() + with obs.bind_run(session_id="bob"): + second = ex.submit(obs.get_context).result() + third = ex.submit(obs.get_context).result() + assert (first, second, third) == ({"session_id": "alice"}, {"session_id": "bob"}, {}) + + +def test_carrying_context_into_thread_pools_is_installed_once(): + obs.carry_context_into_thread_pools() + installed = ThreadPoolExecutor.submit + obs.carry_context_into_thread_pools() + assert ThreadPoolExecutor.submit is installed + + def test_set_run_id_and_session_id_with_token_reset(): sid_token = obs.set_session_id("S") rid_token = obs.set_run_id("R") diff --git a/tests/test_planning.py b/tests/test_planning.py index 5768081bc..cd1361b24 100644 --- a/tests/test_planning.py +++ b/tests/test_planning.py @@ -314,3 +314,123 @@ def test_plan_action_names_are_stable() -> None: planning = _import_planning() assert [name for name, _label in planning.PLAN_ACTIONS] == ["approve", "revise", "cancel"] assert "Nothing was run" in planning.STALE_PLAN_ACTION_NOTE + + +# --------------------------------------------------------------------------- # +# Model calls: whose they are, and what happens when they fail +# --------------------------------------------------------------------------- # + + +class _FakeStep: + """Stands in for cl.Step: an async context manager with an ``output``.""" + + last: _FakeStep | None = None + + def __init__(self, **_kwargs) -> None: + self.output = None + + async def __aenter__(self) -> _FakeStep: + _FakeStep.last = self + return self + + async def __aexit__(self, *exc) -> bool: + return False + + +class _RecordingLLM: + """Records the identity bound when it is invoked, then answers or raises.""" + + def __init__(self, reply: str = "NEEDS_PLAN", error: Exception | None = None) -> None: + self.seen: list[object] = [] + self.reply = reply + self.error = error + + def invoke(self, _messages): + from biomni.identity import current_identity + + self.seen.append(current_identity()) + if self.error is not None: + raise self.error + return types.SimpleNamespace(content=self.reply) + + +def _agent_with(llm: _RecordingLLM) -> _Agent: + agent = _Agent() + agent.llm = llm + return agent + + +def _grip_user(): + from biomni.identity import UserIdentity + + return UserIdentity(user_id="u1", workspace_id="ws-1", source="headers") + + +def test_triage_runs_its_model_call_as_the_session(monkeypatch) -> None: + """The LLM proxy bills each call to the session's user; a thread hop must not lose it.""" + import asyncio + + from biomni.identity import bound_identity + + planning = _import_planning() + monkeypatch.delenv("BIOMNI_ALWAYS_PLAN", raising=False) + llm = _RecordingLLM(reply="It is in studyA.") + who = _grip_user() + + async def ask(): + with bound_identity(who): + return await planning.quick_answer(_agent_with(llm), "where is my file?") + + assert asyncio.run(ask()) == "It is in studyA." + assert llm.seen == [who] + + +def test_planning_runs_its_model_call_as_the_session(monkeypatch) -> None: + import asyncio + + from biomni.identity import bound_identity + + planning = _import_planning() + monkeypatch.setattr(planning, "cl", types.SimpleNamespace(Step=_FakeStep)) + # A failure that is not about access ends planning with the original prompt, + # which is all this test needs after the call has been observed. + llm = _RecordingLLM(error=ValueError("malformed")) + who = _grip_user() + + async def plan(): + with bound_identity(who): + return await planning.interactive_planning(_agent_with(llm), "analyse studyA", agent_type="ad1") + + assert asyncio.run(plan()) == "analyse studyA" + assert llm.seen == [who] + assert "Could not generate a plan: ValueError: malformed." in _FakeStep.last.output + + +def test_an_unreachable_model_stops_the_run_instead_of_carrying_on_without_a_plan(monkeypatch) -> None: + """Proceeding would only fail again, identically, at the agent's first step.""" + import asyncio + + from biomni.llm_proxy import LLMProxyIdentityError + + planning = _import_planning() + monkeypatch.setattr(planning, "cl", types.SimpleNamespace(Step=_FakeStep)) + error = LLMProxyIdentityError("the LLM proxy needs the user and workspace every model call is made for, and ...") + llm = _RecordingLLM(error=error) + + with pytest.raises(LLMProxyIdentityError): + asyncio.run(planning.interactive_planning(_agent_with(llm), "analyse studyA", agent_type="ad1")) + # Closed with a readable line, not the raw exception text Chainlit would + # have written had the error escaped from inside the step. + assert _FakeStep.last.output == "⚠️ Could not reach the language model." + + +def test_the_a1_planning_prompt_names_only_tools_that_exist() -> None: + """A tool the plan names but the agent lacks is a failed first step.""" + import re + + from biomni.tool.tool_description import literature + + planning = _import_planning() + named = set(re.findall(r"\b(?:search|query)_[a-z_]+\b", planning.PLANNING_SYSTEM_PROMPT)) + known = {api["name"] for api in literature.description} + assert named and named <= known diff --git a/tests/test_repl_isolation.py b/tests/test_repl_isolation.py index 13e050ee7..3aabd9fff 100644 --- a/tests/test_repl_isolation.py +++ b/tests/test_repl_isolation.py @@ -231,6 +231,31 @@ def test_credentials_are_restored_when_generated_code_raises(monkeypatch: pytest assert os.environ["ANTHROPIC_API_KEY"] == "sk-ant-secret" +_POOL_SNIPPET = """ +from concurrent.futures import ThreadPoolExecutor +from biomni.identity import current_identity + +def who(_): + print("from the pool") + return current_identity().user_id + +with ThreadPoolExecutor(max_workers=2) as pool: + seen = list(pool.map(who, range(3))) +print(seen) +""" + + +def test_a_thread_pool_in_generated_code_stays_in_its_session() -> None: + """Behind the LLM proxy a pool task outside the session cannot call a model at all.""" + from biomni.identity import UserIdentity, bound_identity + + who = UserIdentity(user_id="u1", workspace_id="w1", source="headers") + with bind_run(session_id="chat-1"), bound_identity(who): + out = run_python_repl(_POOL_SNIPPET) + # The pool's prints are captured into the step that ran it, too. + assert out == "from the pool\n" * 3 + "['u1', 'u1', 'u1']\n" + + def test_output_printed_before_a_failure_is_kept() -> None: """The agent uses it to work out which step broke.""" out = run_python_repl("print('step one done')\nraise ValueError('boom')") diff --git a/tests/test_tool_availability.py b/tests/test_tool_availability.py new file mode 100644 index 000000000..10785da3b --- /dev/null +++ b/tests/test_tool_availability.py @@ -0,0 +1,100 @@ +"""Tests for biomni.tool.availability - tools a deployment must not advertise. + +``advanced_web_search_claude`` calls Anthropic directly. Behind the platform's +LLM proxy that would go around the per-user metering the platform requires, so +there it is neither offered to the agent nor callable. +""" + +from __future__ import annotations + +import pytest +from biomni.tool import availability +from biomni.tool.availability import CLAUDE_WEB_SEARCH + + +@pytest.fixture(autouse=True) +def _direct_claude(monkeypatch): + """A deployment where the Claude web search works, unless a test says otherwise.""" + from biomni.config import default_config + + for name in ("BIOMNI_LLM_PROXY_URL", "BIOMNI_LLM_PROXY_SCHEMA", "BIOMNI_LLM_PROXY_API_KEY"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-test") + monkeypatch.setattr(default_config, "llm", "claude-sonnet-4-5") + + +def test_available_with_a_claude_model_and_a_key(): + assert availability.claude_web_search_problem() is None + assert availability.unavailable_tools() == {} + + +@pytest.mark.parametrize("schema", [None, "anthropic", "not-a-schema"]) +def test_withheld_behind_the_llm_proxy_even_with_a_key(monkeypatch, schema): + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + if schema: + monkeypatch.setenv("BIOMNI_LLM_PROXY_SCHEMA", schema) + problem = availability.unavailable_tools()[CLAUDE_WEB_SEARCH] + assert "LLM proxy" in problem and "search_google()" in problem + + +def test_withheld_without_a_claude_model(monkeypatch): + from biomni.config import default_config + + monkeypatch.setattr(default_config, "llm", "gpt-4.1") + assert "needs a Claude model" in availability.claude_web_search_problem() + + +def test_withheld_without_an_anthropic_key(monkeypatch): + monkeypatch.delenv("ANTHROPIC_API_KEY") + assert "ANTHROPIC_API_KEY" in availability.claude_web_search_problem() + + +def _catalogue_names() -> set[str]: + from biomni.utils import read_module2api + + return {api["name"] for apis in read_module2api().values() for api in apis} + + +def test_the_catalogue_leaves_out_what_cannot_run(monkeypatch): + assert CLAUDE_WEB_SEARCH in _catalogue_names() + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + names = _catalogue_names() + assert CLAUDE_WEB_SEARCH not in names + # Only that one tool: the rest of the literature tools are still offered. + assert {"search_google", "query_pubmed", "query_arxiv"} <= names + + +def test_the_tool_refuses_behind_the_proxy_without_calling_anthropic(monkeypatch): + literature = pytest.importorskip("biomni.tool.literature") + anthropic = pytest.importorskip("anthropic") + + def no_client(*_args, **_kwargs): + raise AssertionError("a direct Anthropic client must not be created behind the proxy") + + monkeypatch.setattr(anthropic, "Anthropic", no_client) + monkeypatch.setenv("BIOMNI_LLM_PROXY_URL", "https://proxy.example/v1") + with pytest.raises(RuntimeError, match="LLM proxy"): + literature.advanced_web_search_claude("amyloid PET tracers") + + +def test_the_tool_sends_only_the_anthropic_key(monkeypatch): + """Never the custom-endpoint key: that would hand one provider's credential to another.""" + literature = pytest.importorskip("biomni.tool.literature") + anthropic = pytest.importorskip("anthropic") + from biomni.config import default_config + + monkeypatch.setattr(default_config, "api_key", "custom-endpoint-key") + created = {} + + class _Client: + def __init__(self, api_key=None, **_kwargs): + created["api_key"] = api_key + self.messages = self + + def create(self, **_kwargs): + raise RuntimeError("stop here") + + monkeypatch.setattr(anthropic, "Anthropic", _Client) + monkeypatch.setattr(literature.time, "sleep", lambda _s: None) + literature.advanced_web_search_claude("amyloid PET tracers", max_retries=1) + assert created["api_key"] == "sk-ant-test" diff --git a/tests/test_workspace_panel.py b/tests/test_workspace_panel.py index 2b37f9ccd..6366a67ef 100644 --- a/tests/test_workspace_panel.py +++ b/tests/test_workspace_panel.py @@ -9,6 +9,8 @@ from __future__ import annotations +import os + import pytest from biomni import fs_scan from biomni.workspace_prefs import WorkspacePrefs, resolve_scope @@ -94,6 +96,20 @@ def test_inventory_without_a_selection_lists_folders_and_forbids_assuming_empty( assert "a1.csv" not in text +@pytest.mark.skipif(os.geteuid() == 0, reason="root can list any directory") +def test_an_unreadable_workspace_is_described_as_unreadable_not_empty(workspace): + """Otherwise the agent tells a user whose data is all there that there is none.""" + workspace.chmod(0o300) + try: + scope = resolve_scope(WorkspacePrefs(), str(workspace)) + text = panel.build_scope_inventory(scope, str(workspace)) + finally: + workspace.chmod(0o755) + assert "NOT permitted to read the workspace" in text + assert "Do not guess at what the workspace contains." in text + assert "No top-level folders" not in text + + def test_inventory_without_a_selection_does_not_walk_the_workspace(workspace, monkeypatch): """The whole point of the default path: no recursive scan on a cold session.""" diff --git a/tests/test_workspace_prefs.py b/tests/test_workspace_prefs.py index 426580687..c59049399 100644 --- a/tests/test_workspace_prefs.py +++ b/tests/test_workspace_prefs.py @@ -11,6 +11,7 @@ import json import logging import os +import stat import pytest from biomni import workspace_prefs as wp @@ -248,7 +249,7 @@ def test_unwritable_preference_falls_through_to_next_candidate(tmp_path): prefs = wp.WorkspacePrefs(output_dir=str(ro / "nope")) target = wp.resolve_output_dir(prefs, workspace_root=str(tmp_path)) assert target.source == "workspace" - assert target.reason and "preference" in target.reason + assert target.reason and target.reason.startswith("The output directory you configured") finally: ro.chmod(0o700) @@ -435,3 +436,349 @@ def test_delete_forgets_stored_preferences(tmp_path): def test_null_store_delete_is_a_no_op(): assert wp.NullPrefsStore().delete("k") is False + + +# --------------------------------------------------------------------------- # +# Explaining an unusable directory +# --------------------------------------------------------------------------- # +# +# GRIP runs the app as uid 57439 against an Azure Files (SMB) share. What the +# process meets there is decided by the mount table and by who the process is, +# so both are substituted: a mount table written for the test, and a process +# identity that is not the owner of the test's directories. + +_GRIP_UID = 57439 + + +def _mount_table(tmp_path, monkeypatch, *lines: str) -> None: + table = tmp_path / "mounts" + table.write_text("".join(f"{line}\n" for line in lines)) + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(table)) + + +def _as_grip_uid(monkeypatch) -> None: + monkeypatch.setattr(wp, "_process_ids", lambda: (_GRIP_UID, _GRIP_UID, frozenset({_GRIP_UID}))) + + +def _share(tmp_path, mode: int = 0o755): + share = tmp_path / "user-data" + share.mkdir() + share.chmod(mode) + return share + + +def test_mount_info_reads_escaped_paths_and_options(tmp_path, monkeypatch): + _mount_table( + tmp_path, + monkeypatch, + "overlay / overlay rw,relatime 0 0", + "//acct.file.core.windows.net/share /app/user\\040data cifs rw,vers=3.1.1,uid=0,gid=0,dir_mode=0777 0 0", + ) + info = wp.mount_info("/app/user data/studyA") + assert info is not None + assert (info.mountpoint, info.fstype, info.source) == ( + "/app/user data", + "cifs", + "//acct.file.core.windows.net/share", + ) + assert info.option("uid") == "0" and info.option("dir_mode") == "0777" + assert not info.read_only + + +def test_the_last_mount_on_a_mountpoint_is_the_one_in_effect(tmp_path, monkeypatch): + """Mounting again over the same point hides the first mount.""" + _mount_table( + tmp_path, + monkeypatch, + "//acct/share /app/user-data cifs rw,uid=0,gid=0 0 0", + "//acct/share /app/user-data cifs rw,uid=57439,gid=57439 0 0", + ) + assert wp.mount_info("/app/user-data").option("uid") == "57439" + + +def test_a_mount_point_matches_whole_path_components_only(tmp_path, monkeypatch): + """/app/user-data is not a parent of /app/user.""" + _mount_table(tmp_path, monkeypatch, "overlay / overlay rw 0 0", "//acct/share /app/user-data cifs rw 0 0") + assert wp.mount_info("/app/user").mountpoint == "/" + assert wp.mount_info("/app/user-data/x").mountpoint == "/app/user-data" + + +def test_a_symlink_onto_a_share_is_on_the_share(tmp_path, monkeypatch): + share = _share(tmp_path) + link = tmp_path / "workspace-link" + link.symlink_to(share) + _mount_table(tmp_path, monkeypatch, "overlay / overlay rw 0 0", f"//acct/share {share} cifs rw 0 0") + assert wp.mount_info(str(link)).fstype == "cifs" + assert wp.mount_info(str(link / "study")).mountpoint == str(share) + + +def test_a_mount_point_that_is_not_utf8_does_not_break_the_lookup(tmp_path, monkeypatch): + table = tmp_path / "mounts" + table.write_bytes(b"overlay / overlay rw 0 0\n/dev/sdc /mnt/caf\xe9 ext4 rw 0 0\n") + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(table)) + assert wp.mount_info("/app/runs").mountpoint == "/" + + +def test_no_mount_table_means_cannot_tell(tmp_path, monkeypatch): + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(tmp_path / "absent")) + assert wp.mount_info("/app") is None + + +def test_an_smb_share_is_explained_by_the_options_it_is_actually_mounted_with(tmp_path, monkeypatch): + """The GRIP case: new mount options set, but not the ones this container sees.""" + share = _share(tmp_path) + _mount_table( + tmp_path, monkeypatch, f"//acct.file.core.windows.net/share {share} cifs rw,uid=0,gid=0,dir_mode=0755 0 0" + ) + _as_grip_uid(monkeypatch) + + reason = wp.explain_unwritable(str(share / "biomni-outputs")) + assert f"{share} is owned by uid" in reason + # The share, never the storage account that serves it. + assert "on an SMB share (/share)" in reason and "acct" not in reason + assert "which are uid=0,gid=0,dir_mode=0755 on the mount this container sees" in reason + assert "Mount it with uid=57439,gid=57439,dir_mode=0770,file_mode=0770" in reason + assert "mounted afresh on the node" in reason + + +def test_an_smb_share_with_no_ownership_options_says_they_are_unset(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"//acct/share {share} cifs rw,vers=3.1.1 0 0") + _as_grip_uid(monkeypatch) + assert "which are unset (uid=0,gid=0)" in wp.explain_unwritable(str(share)) + + +def test_an_smb_share_under_server_acls_is_not_blamed_on_the_mount_options(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"//acct/share {share} cifs rw,cifsacl,uid=0,gid=0 0 0") + _as_grip_uid(monkeypatch) + reason = wp.explain_unwritable(str(share)) + assert "mounted with cifsacl, so its owner and mode come from the file server's ACLs" in reason + assert "grant uid 57439 write access in the share's ACLs" in reason + assert "come only from the mount options" not in reason + + +def test_a_share_the_process_cannot_list_is_explained_for_reading(tmp_path, monkeypatch): + """dir_mode=0770 applied but uid/gid not: the workspace becomes unreadable, not just read-only.""" + share = _share(tmp_path, 0o770) + try: + _mount_table(tmp_path, monkeypatch, f"//acct/share {share} cifs rw,uid=0,gid=0,dir_mode=0770 0 0") + _as_grip_uid(monkeypatch) + reason = wp.explain_unreadable(str(share)) + assert "with mode 0770 on an SMB share" in reason + assert "uid=0,gid=0,dir_mode=0770" in reason + assert "Mount it with uid=57439,gid=57439" in reason + finally: + share.chmod(0o755) + + +def test_a_read_only_mount_is_named_as_such(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"/dev/sdb {share} ext4 ro,relatime 0 0") + _as_grip_uid(monkeypatch) + + def no_statvfs(_path): + raise OSError("unsupported") + + # The mount table is the fallback witness when statvfs cannot answer. + monkeypatch.setattr(os, "statvfs", no_statvfs) + reason = wp.explain_unwritable(str(share)) + assert reason == f"{share} is on a read-only mount (ext4 mount of /dev/sdb); mount the volume read-write" + + +def test_a_read_only_container_image_is_not_told_to_remount(tmp_path, monkeypatch): + """readOnlyRootFilesystem with the volume missing: there is nothing to remount.""" + image = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"overlay {image} overlay ro,relatime 0 0") + _as_grip_uid(monkeypatch) + monkeypatch.setattr(os, "statvfs", lambda _path: (_ for _ in ()).throw(OSError("unsupported"))) + reason = wp.explain_unwritable(str(image / "runs")) + assert reason == ( + f"{image} is inside the container image, which is mounted read-only; mount a writable volume at " + f"{image / 'runs'}" + ) + + +def test_an_nfs_export_is_explained_as_ignoring_fsgroup(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"nfs.example:/export {share} nfs4 rw 0 0") + _as_grip_uid(monkeypatch) + reason = wp.explain_unwritable(str(share)) + assert "on an NFS export (/export), which ignores fsGroup" in reason + assert "grant uid 57439 write access on the NFS server" in reason + + +def test_a_fuse_mount_is_not_given_the_fsgroup_remedy(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"blobfuse2 {share} fuse rw,nosuid,nodev 0 0") + _as_grip_uid(monkeypatch) + reason = wp.explain_unwritable(str(share)) + assert "on a FUSE mount (blobfuse2), which ignores fsGroup" in reason + assert "allow_other" in reason + assert "securityContext.fsGroup" not in reason + + +def test_an_ordinary_volume_gets_the_fsgroup_remedy(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"/dev/sdb {share} ext4 rw 0 0") + _as_grip_uid(monkeypatch) + assert wp.explain_unwritable(str(share)).endswith( + "(ext4 mount of /dev/sdb); set securityContext.fsGroup: 57439 on the pod (it does not apply to hostPath " + "volumes), or chown it to uid 57439" + ) + + +def _dir_stat(mode: int, uid: int, gid: int): + return os.stat_result((stat.S_IFDIR | mode, 0, 0, 0, uid, gid, 0, 0, 0, 0)) + + +def test_the_owner_class_alone_decides_for_the_owner(): + """0o577: group and others may write, the owner may not - and the owner is denied.""" + st = _dir_stat(0o577, uid=1000, gid=1000) + assert not wp._mode_allows(st, 1000, frozenset({1000}), write=True) + assert wp._mode_allows(st, 1000, frozenset({1000}), write=False) + assert wp._mode_allows(st, 2000, frozenset({1000}), write=True) # group member + assert wp._mode_allows(st, 3000, frozenset({3000}), write=True) # anyone else + + +@pytest.mark.skipif(os.geteuid() == 0, reason="root is never denied by mode bits") +def test_a_directory_that_denies_its_own_owner_is_not_answered_with_chown(tmp_path, monkeypatch): + share = _share(tmp_path, 0o500) + try: + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(tmp_path / "absent")) + reason = wp.explain_unwritable(str(share)) + assert f"which denies its own owner (uid {os.getuid()})" in reason + assert "chmod u+rwx" in reason + assert "chown" not in reason + finally: + share.chmod(0o755) + + +def test_a_mode_that_permits_access_points_at_what_else_can_refuse_it(tmp_path, monkeypatch): + share = _share(tmp_path) + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(tmp_path / "absent")) + uid = os.stat(share).st_uid + monkeypatch.setattr(wp, "_process_ids", lambda: (uid, 0, frozenset({0}))) + assert "which permits writing, so something else refuses it" in wp.explain_unwritable(str(share)) + + +def test_storage_facts_carry_what_a_log_reader_needs(tmp_path, monkeypatch): + share = _share(tmp_path) + _mount_table( + tmp_path, + monkeypatch, + f"//acct/share {share} cifs rw,vers=3.1.1,username=acct,addr=10.0.0.9,uid=0,gid=0,dir_mode=0755,mfsymlinks 0 0", + ) + _as_grip_uid(monkeypatch) + facts = wp.storage_facts(str(share / "biomni-outputs")) + assert facts["path"] == str(share / "biomni-outputs") + assert facts["checked"] == str(share) + assert facts["fstype"] == "cifs" and facts["source"] == "/share" + assert facts["mode"] == "0755" + # Only the options that decide access - no account name, no server address. + assert facts["options"] == "rw,vers=3.1.1,uid=0,gid=0,dir_mode=0755" + assert not any("acct" in str(value) or "10.0.0.9" in str(value) for value in facts.values()) + assert {"readable", "writable"} <= facts.keys() + + +def test_storage_facts_explain_what_the_process_cannot_do(tmp_path, monkeypatch): + share = _share(tmp_path) + monkeypatch.setattr(wp, "_MOUNTS_FILE", str(tmp_path / "absent")) + monkeypatch.setattr(wp, "is_writable_dir", lambda _path: False) + monkeypatch.setattr(wp, "is_readable_dir", lambda _path: False) + _as_grip_uid(monkeypatch) + facts = wp.storage_facts(str(share)) + assert facts["write_reason"].startswith(f"{share} is owned by uid") + assert facts["read_reason"].startswith(f"{share} is owned by uid") + + +def test_candidates_failing_on_the_same_directory_are_explained_once(tmp_path, monkeypatch): + """BIOMNI_OUTPUT_ROOT and the workspace default usually fail together, on one mount.""" + workspace = _share(tmp_path, 0o500) + try: + monkeypatch.setenv("BIOMNI_OUTPUT_ROOT", str(workspace)) + target = wp.resolve_output_dir(wp.WorkspacePrefs(), workspace_root=str(workspace)) + assert target.source == "cwd-fallback" + assert target.reason is not None + assert target.reason.startswith( + f"BIOMNI_OUTPUT_ROOT ({workspace}) and the workspace default ({workspace}/biomni-outputs) are not writable" + ) + assert target.reason.count(f"{workspace} is owned by uid") == 1 + finally: + workspace.chmod(0o755) + + +@pytest.mark.skipif(os.geteuid() == 0, reason="root is never denied by mode bits") +def test_three_locations_failing_together_read_as_a_list(tmp_path, monkeypatch): + workspace = _share(tmp_path, 0o500) + try: + monkeypatch.setenv("BIOMNI_OUTPUT_ROOT", str(workspace / "root")) + prefs = wp.WorkspacePrefs(output_dir=str(workspace / "mine")) + target = wp.resolve_output_dir(prefs, workspace_root=str(workspace)) + assert target.reason is not None + assert target.reason.startswith( + f"The output directory you configured ({workspace}/mine), BIOMNI_OUTPUT_ROOT ({workspace}/root) and the " + f"workspace default ({workspace}/biomni-outputs) are not writable" + ) + finally: + workspace.chmod(0o755) + + +@pytest.mark.skipif(os.geteuid() == 0, reason="root is never denied by mode bits") +def test_an_output_root_that_is_the_workspace_default_is_reported_once(tmp_path, monkeypatch): + """The GRIP setup: BIOMNI_OUTPUT_ROOT=/biomni-outputs.""" + workspace = _share(tmp_path, 0o500) + try: + monkeypatch.setenv("BIOMNI_OUTPUT_ROOT", str(workspace / "biomni-outputs")) + target = wp.resolve_output_dir(wp.WorkspacePrefs(), workspace_root=str(workspace)) + assert target.reason is not None + assert target.reason.startswith(f"BIOMNI_OUTPUT_ROOT ({workspace}/biomni-outputs) is not writable") + assert "the workspace default" not in target.reason + finally: + workspace.chmod(0o755) + + +@pytest.mark.skipif(os.geteuid() == 0, reason="root is never denied by mode bits") +def test_each_passed_over_location_starts_a_sentence_of_its_own(tmp_path, monkeypatch): + chosen = _share(tmp_path, 0o500) + workspace = tmp_path / "workspace" + workspace.mkdir(mode=0o500) + try: + monkeypatch.delenv("BIOMNI_OUTPUT_ROOT", raising=False) + target = wp.resolve_output_dir(wp.WorkspacePrefs(output_dir=str(chosen)), workspace_root=str(workspace)) + assert target.reason is not None + assert target.reason.startswith(f"The output directory you configured ({chosen}) is not writable") + assert f". The workspace default ({workspace}/biomni-outputs) is not writable" in target.reason + finally: + chosen.chmod(0o755) + workspace.chmod(0o755) + + +def test_a_directory_inside_the_image_is_answered_with_a_volume(tmp_path, monkeypatch): + """A setting that points where no volume is mounted: opening up the image would keep nothing.""" + image = _share(tmp_path) + _mount_table(tmp_path, monkeypatch, f"overlay {image} overlay rw,relatime 0 0") + _as_grip_uid(monkeypatch) + assert wp.explain_unwritable(str(image / "outputs")).endswith( + f"inside the container image rather than on a mounted volume; mount a writable volume at {image}/outputs" + ) + + +def test_an_unwritable_location_is_logged_once_not_on_every_resolution(tmp_path, monkeypatch, caplog): + ro = _share(tmp_path, 0o500) + try: + monkeypatch.setenv("BIOMNI_OUTPUT_ROOT", str(ro)) + with caplog.at_level(logging.WARNING, logger="biomni.workspace_prefs"): + for _ in range(3): + wp.resolve_output_dir(wp.WorkspacePrefs(), workspace_root=None) + assert sum("BIOMNI_OUTPUT_ROOT" in r.getMessage() for r in caplog.records) == 1 + finally: + ro.chmod(0o755) + + +def test_is_readable_dir(tmp_path): + assert wp.is_readable_dir(str(tmp_path)) + assert not wp.is_readable_dir(None) + assert not wp.is_readable_dir(str(tmp_path / "absent")) + (tmp_path / "file").write_text("x") + assert not wp.is_readable_dir(str(tmp_path / "file"))