mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-23 17:58:35 +08:00
Merge branch 'release-1.11.0' into feat/a11y
This commit is contained in:
Viktor Avelino
121
docs/docs/Components/bundles-nextplaid.mdx
Normal file
121
docs/docs/Components/bundles-nextplaid.mdx
Normal file
@ -0,0 +1,121 @@
|
||||
---
|
||||
title: NextPlaid
|
||||
slug: /bundles-nextplaid
|
||||
---
|
||||
|
||||
import Icon from "@site/src/components/icon";
|
||||
import PartialParams from '@site/docs/_partial-hidden-params.mdx';
|
||||
import PartialVectorSearchResults from '@site/docs/_partial-vector-search-results.mdx';
|
||||
import PartialVectorStoreInstance from '@site/docs/_partial-vector-store-instance.mdx';
|
||||
|
||||
<Icon name="Blocks" aria-hidden="true" /> [**Bundles**](/components-bundle-components) contain custom components that support specific third-party integrations with Langflow.
|
||||
|
||||
The **NextPlaid** bundle provides multi-vector ColBERT-style retrieval for Langflow through a running [NextPlaid](https://github.com/meetdoshi90/next-plaid) server.
|
||||
NextPlaid stores each document as a matrix of token embeddings, which enables higher retrieval quality on semantic search tasks through ColBERT-style late interaction with MaxSim scoring.
|
||||
For more information on multi-vector embeddings, see the [langchain-plaid](https://github.com/meetdoshi90/langchain-plaid) package that this bundle is built on.
|
||||
|
||||
## Install the NextPlaid bundle
|
||||
|
||||
The bundle includes two components:
|
||||
|
||||
* **NextPlaid**: vector store component backed by a running NextPlaid server.
|
||||
* **vLLM Multivector Embeddings**: generates the multi-vector token embeddings required by NextPlaid.
|
||||
|
||||
The **NextPlaid** bundle is included in the `lfx-nextplaid` Extension bundle, which is installed automatically as part of `uv pip install langflow`.
|
||||
|
||||
If you need to install it separately, run:
|
||||
|
||||
```bash
|
||||
uv pip install lfx-nextplaid
|
||||
uv run langflow run
|
||||
```
|
||||
|
||||
To verify the bundle is loaded in your environment:
|
||||
|
||||
```bash
|
||||
lfx extension list
|
||||
```
|
||||
|
||||
## NextPlaid
|
||||
|
||||
The **NextPlaid** component reads from and writes to a [NextPlaid](https://github.com/meetdoshi90/next-plaid) multi-vector search server.
|
||||
|
||||
NextPlaid is a Rust-based server that implements ColBERT-style late interaction retrieval with MaxSim scoring.
|
||||
Each document is stored as a matrix of token embeddings rather than a single dense vector, giving better retrieval quality on semantic search tasks.
|
||||
The component supports both text retrieval with ColBERT models and image retrieval with ColPali models.
|
||||
|
||||
<details>
|
||||
<summary>About vector store instances</summary>
|
||||
|
||||
<PartialVectorStoreInstance />
|
||||
|
||||
</details>
|
||||
|
||||
<PartialVectorSearchResults />
|
||||
|
||||
### Use the NextPlaid component in a flow
|
||||
|
||||
Connect a **vLLM Multivector Embeddings** component to the **Embedding (Multivector)** input. Standard single-vector embedding components are not compatible with NextPlaid.
|
||||
|
||||
The **NextPlaid** component can be used for both writes and reads:
|
||||
|
||||
* When writing, it ingests documents from an attached data source, computes multi-vector embeddings with the connected **vLLM Multivector Embeddings** component, and then loads them into the NextPlaid index.
|
||||
To trigger writes, click <Icon name="Play" aria-hidden="true"/> **Run component** on the **NextPlaid** component.
|
||||
|
||||
* When reading, the **NextPlaid** component uses chat input to perform a MaxSim similarity search against the index and returns the top results.
|
||||
|
||||
### NextPlaid parameters
|
||||
|
||||
<PartialParams />
|
||||
|
||||
| Name | Type | Description |
|
||||
|------|------|-------------|
|
||||
| **Server URL** (`url`) | String | Input parameter. Base URL of the running NextPlaid server. Default: `http://localhost:8080`. |
|
||||
| **Index Name** (`index_name`) | String | Input parameter. Name of the index to create or connect to. Default: `langflow`. |
|
||||
| **Ingest Data** (`ingest_data`) | Data | Input parameter. Documents or images to write to the vector store. Only relevant for writes. |
|
||||
| **Search Query** (`search_query`) | String | Input parameter. The query string to use for similarity search. Only relevant for reads. |
|
||||
| **Embedding (Multivector)** (`embedding`) | Embeddings | Input parameter. Connect a **vLLM Multivector Embeddings** component to generate token-level embeddings. Required for both reads and writes. |
|
||||
| **Index Batch Size** (`index_batch_size`) | Integer | Input parameter. Number of documents per indexing request. PLAID builds its initial cluster centroids from the first batch — larger batches produce better retrieval quality. Default: `500`. |
|
||||
| **Number of Results** (`number_of_results`) | Integer | Input parameter. Number of results to return from similarity search. Default: `4`. |
|
||||
| **Quantization Bits** (`nbits`) | Dropdown | Input parameter. Bit-width for PLAID quantization. `4` gives better quality; `2` uses less memory. Options: `2`, `4`. Default: `4`. |
|
||||
| **Create Index If Not Exists** (`create_index_if_not_exists`) | Boolean | Input parameter. If `true`, creates the index on the NextPlaid server if it does not already exist. Default: `true`. |
|
||||
| **Write Timeout** (`write_timeout`) | Float | Input parameter. Seconds to wait for each indexing batch to finish. Set to `0` for async indexing (search may return empty results on the first run). Recommended: `30` or higher when ingesting and searching in the same flow run. Default: `30.0`. |
|
||||
|
||||
## vLLM Multivector Embeddings
|
||||
|
||||
The **vLLM Multivector Embeddings** component generates multi-vector token embeddings by calling vLLM's `/pooling` endpoint with `task: token_embed`.
|
||||
|
||||
The output is a multi-vector `Embeddings` object that returns a matrix of token embeddings per document, which is required by the **NextPlaid** vector store. It is not compatible with standard single-vector stores.
|
||||
|
||||
For more information about using embedding model components in flows, see [Embedding model components](/components-embedding-models).
|
||||
|
||||
Your vLLM server must be started with a ColBERT- or ColPali-compatible model and the pooling runner enabled:
|
||||
|
||||
```bash
|
||||
vllm serve <model> --runner pooling --pooler-config '{"task": "token_embed"}'
|
||||
```
|
||||
|
||||
Compatible models include:
|
||||
* **Text (ColBERT):** `answerdotai/answerai-colbert-small-v1`
|
||||
* **Text + Images (ColPali):** `ModernVBERT/colmodernvbert`
|
||||
|
||||
For more information on the vLLM server, see the [vLLM documentation](https://docs.vllm.ai/).
|
||||
|
||||
### vLLM Multivector Embeddings parameters
|
||||
|
||||
<PartialParams />
|
||||
|
||||
| Name | Type | Description |
|
||||
|------|------|-------------|
|
||||
| **Model Name** (`model_name`) | String | Input parameter. The multi-vector model name served by vLLM. Default: `answerdotai/answerai-colbert-small-v1`. |
|
||||
| **vLLM API Base** (`api_base`) | String | Input parameter. Base URL of the vLLM server (without the `/v1` suffix). Default: `http://localhost:8000`. |
|
||||
| **API Key** (`api_key`) | SecretString | Input parameter. API key for the vLLM server. Leave empty for local servers. Optional. |
|
||||
| **Request Timeout** (`request_timeout`) | Float | Input parameter. Timeout in seconds for each request to the vLLM API. Default: `60.0`. Advanced. |
|
||||
| **Max Retries** (`max_retries`) | Integer | Input parameter. Number of times to retry a failed request before raising an error. Default: `3`. Advanced. |
|
||||
|
||||
## See also
|
||||
|
||||
* [**Embedding model** components](/components-embedding-models)
|
||||
* [NextPlaid server](https://github.com/meetdoshi90/next-plaid)
|
||||
* [langchain-plaid](https://github.com/meetdoshi90/langchain-plaid)
|
||||
* [vLLM documentation](https://docs.vllm.ai/)
|
||||
@ -56,8 +56,19 @@ To use the vLLM component, you need to have a vLLM server running. Here are the
|
||||
|
||||
For more detailed setup instructions, see the [vLLM documentation](https://docs.vllm.ai/en/latest/getting_started/quickstart.html).
|
||||
|
||||
## vLLM multi-vector embeddings
|
||||
|
||||
For multi-vector ColBERT or ColPali-style retrieval with vLLM, use the **vLLM Multivector Embeddings** component.
|
||||
|
||||
The **vLLM Multivector Embeddings** component ships with the **NextPlaid** bundle (`lfx-nextplaid`), since it is purpose-built to feed that vector store.
|
||||
|
||||
For more information on using the component in a flow, see [vLLM Multi-vector Embeddings](./bundles-nextplaid#vllm-multivector-embeddings).
|
||||
|
||||
For more information on multi-vector embeddings, see the [langchain-plaid repository](https://github.com/meetdoshi90/langchain-plaid).
|
||||
|
||||
## See also
|
||||
|
||||
* [**Agent** component](/components-agents)
|
||||
* [**Language Model** component](/components-models)
|
||||
* [**NextPlaid** bundle](./bundles-nextplaid)
|
||||
* [vLLM GitHub repository](https://github.com/vllm-project/vllm)
|
||||
|
||||
@ -46,211 +46,24 @@ To avoid the impact of potential breaking changes and test new versions, the Lan
|
||||
|
||||
If you made changes to your flows in the isolated installation, you might want to export and import those flows back to your upgraded primary installation so you don't have to repeat the component upgrade process.
|
||||
|
||||
## 1.10.x
|
||||
## 1.11.x
|
||||
|
||||
Highlights of this release include the following changes.
|
||||
For all changes, see the [Changelog](https://github.com/langflow-ai/langflow/releases).
|
||||
|
||||
|
||||
### Breaking changes
|
||||
|
||||
- Upcoming breaking change: bundle separation
|
||||
|
||||
Langflow 1.10.x introduces the Extension bundles model.
|
||||
In future releases, bundle separation will extend to more component providers.
|
||||
|
||||
Saved flows will continue to open normally, but the following may require manual updates before upgrading:
|
||||
|
||||
- Code or scripts that reference component class names or `lfx.components.<provider>` import paths.
|
||||
- External tooling that reads component type identifiers from raw flow JSON.
|
||||
- Custom deployment scripts that install or pin individual component packages.
|
||||
|
||||
For more information, see [Langflow Extensions overview](../Develop/extensions-overview.mdx).
|
||||
|
||||
### New features and enhancements
|
||||
|
||||
- **Langflow Assistant**: build complete flows
|
||||
- NextPlaid multi-vector bundle
|
||||
|
||||
**Langflow Assistant** can now build entire flows, and not only individual components.
|
||||
The **NextPlaid** bundle adds two new components for ColBERT-style multi-vector retrieval.
|
||||
|
||||
For more information, see [Build flows and components with Langflow Assistant](/langflow-assistant).
|
||||
The **NextPlaid** vector store, backed by a running [NextPlaid](https://github.com/meetdoshi90/next-plaid) server, stores each document as a matrix of token embeddings, and the **vLLM Multivector Embeddings** component generates the token-level multi-vector embeddings required by NextPlaid.
|
||||
|
||||
- Redis-backed job queue for multi-worker deployments
|
||||
For more information, see [NextPlaid bundle](./bundles-nextplaid).
|
||||
|
||||
Langflow now supports a Redis-backed job queue that allows flow build events to be shared across multiple Gunicorn/Uvicorn workers and across multiple replicas behind a load balancer.
|
||||
## 1.10.x
|
||||
|
||||
Single-process deployments are unaffected. The default `asyncio` in-memory queue is unchanged.
|
||||
|
||||
For setup instructions and configuration details, see [Deploy Langflow with multiple workers](./deployment-multi-worker).
|
||||
|
||||
- Memory bases
|
||||
|
||||
Memory bases are per-flow vector stores that automatically ingest conversation messages.
|
||||
The **Memory Base** component retrieves context from the vector store, offering long-term semantic memory for flows.
|
||||
For more information, see [Manage memory bases](../Develop/memory-bases.mdx).
|
||||
|
||||
- Python 3.14 support
|
||||
|
||||
Langflow now supports Python 3.10 through 3.14 on macOS, Linux, and Windows.
|
||||
The Langflow Docker images now use Python 3.14.
|
||||
|
||||
<details>
|
||||
<summary>Optional integrations not yet available on Python 3.14</summary>
|
||||
|
||||
The following optional integrations are excluded from installs on Python 3.14:
|
||||
|
||||
- [IBM watsonx](/bundles-ibm)
|
||||
- [IBM watsonx Orchestrate](/deployment-wxo)
|
||||
- [ALTK](/bundles-altk)
|
||||
- [CUGA](/bundles-cuga)
|
||||
- [LangWatch](/integrations-langwatch)
|
||||
- [Pinecone](/bundles-pinecone)
|
||||
- OpenDsStar
|
||||
|
||||
</details>
|
||||
|
||||
- **Agent** component: structured response output
|
||||
|
||||
The **Agent** component now includes a **Structured Response** (`structured_response`) output that returns the agent's reply as structured data according to an **Output Schema** you define.
|
||||
|
||||
For more information, see [Agent component output](/agents#agent-component-output).
|
||||
|
||||
- **Agent** component: default system prompt
|
||||
|
||||
The **Agent** component now uses a structured default system prompt that includes the current date and the active model name.
|
||||
Existing flows are not affected.
|
||||
|
||||
For more information, see [Agent instructions and input](/agents#agent-instructions-and-input).
|
||||
|
||||
- Extension bundles
|
||||
|
||||
Langflow 1.10 introduces Extension bundles: component providers that are packaged and versioned as standalone pip packages, independent of the core Langflow server.
|
||||
All bundles are still included in `uv pip install langflow`, so nothing changes for existing users.
|
||||
|
||||
The following Extension bundles ship with Langflow 1.10:
|
||||
|
||||
| Package | Bundle | Components |
|
||||
|---------|--------|------------|
|
||||
| `lfx-arxiv` | [arXiv](/bundles-arxiv) | arXiv search |
|
||||
| `lfx-docling` | [Docling](/bundles-docling) | Document parsing and chunking |
|
||||
| `lfx-duckduckgo` | [DuckDuckGo](/bundles-duckduckgo) | Web search |
|
||||
| `lfx-ibm` | [IBM](/bundles-ibm) | IBM watsonx.ai LLM and embeddings, IBM Db2 Vector Store |
|
||||
|
||||
To view all Extension bundles currently loaded in your environment:
|
||||
|
||||
```bash
|
||||
lfx extension list
|
||||
```
|
||||
|
||||
If a bundle is not installed, its components are absent from the canvas.
|
||||
To install a missing bundle, install its standalone package and restart Langflow:
|
||||
|
||||
```bash
|
||||
uv pip install lfx-duckduckgo
|
||||
uv run langflow run
|
||||
```
|
||||
|
||||
Components in Extension bundles use a namespaced identifier instead of a bare class name.
|
||||
For example, `DuckDuckGoSearchComponent` is now identified as:
|
||||
|
||||
```text
|
||||
ext:duckduckgo:DuckDuckGoSearchComponent@official
|
||||
```
|
||||
|
||||
You can see this identifier in the raw JSON of any saved flow that uses the component.
|
||||
Langflow migrates saved flow references to the new format automatically when you open a flow.
|
||||
|
||||
For more information, see [Langflow Extensions overview](../Develop/extensions-overview.mdx).
|
||||
|
||||
- File System component
|
||||
|
||||
The **File System** component gives agents sandboxed read/write access to files on disk.
|
||||
An optional **Read Only** mode restricts the agent to read and search operations only.
|
||||
|
||||
For more information, see [File System](../Components/file-system.mdx).
|
||||
|
||||
- Unified **Knowledge Base** component
|
||||
|
||||
The **Knowledge Ingestion** and **Knowledge Base** (legacy) components are merged into a single **Knowledge Base** component with a **Mode** selector.
|
||||
|
||||
Existing flows that use the legacy components continue to work.
|
||||
|
||||
For more information, see the [**Knowledge Base** component](/knowledge-base) and [Manage vector data](/knowledge).
|
||||
|
||||
- Database connectors for knowledge bases
|
||||
|
||||
Knowledge bases now support configurable vector database backends through **DB Providers** configured in **Settings → DB Providers**.
|
||||
Available providers include [Chroma (default)](https://docs.trychroma.com/), [Chroma Cloud](https://docs.trychroma.com/cloud/getting-started), and [OpenSearch](https://docs.opensearch.org/latest/about/).
|
||||
|
||||
For more information, see [Manage vector data](/knowledge).
|
||||
|
||||
- MCP server management lock
|
||||
|
||||
Set `LANGFLOW_MCP_SERVERS_LOCKED=true` to prevent non-superusers from adding, editing, or removing MCP server connections.
|
||||
|
||||
For more information, see [Restrict MCP server management to superusers](/mcp-server#restrict-mcp-server-management).
|
||||
|
||||
- Embedded mode UI flags
|
||||
|
||||
Set `LANGFLOW_EMBEDDED_MODE=true` to hide standalone UI elements when embedding the Langflow visual editor in another application.
|
||||
|
||||
These flags control UI visibility only. They do not block the underlying API endpoints.
|
||||
|
||||
For more information, see [Embedded mode](/environment-variables#embedded-mode).
|
||||
|
||||
- Custom component admin-only restriction
|
||||
|
||||
Set `LANGFLOW_CUSTOM_COMPONENT_ADMIN_ONLY=true` to restrict custom component creation and code editing to superusers.
|
||||
|
||||
For more information, see [Restrict custom component creation to superusers](/deployment-block-custom-components#restrict-custom-components-to-superusers).
|
||||
|
||||
- macOS support matrix
|
||||
|
||||
A new [macOS support](./macos-support-matrix.mdx) page documents feature availability across Apple Silicon and Intel Macs.
|
||||
|
||||
- IBM Db2 Vector Store component
|
||||
|
||||
The **IBM Db2 Vector Store** component is now available in the [IBM bundle](/bundles-ibm).
|
||||
|
||||
For more information, see [IBM Db2 Vector Store](/bundles-ibm#ibm-db2-vector-store).
|
||||
|
||||
- Internationalization
|
||||
|
||||
The Langflow interface is now available in multiple languages.
|
||||
To change the display language, click your **Profile Picture**, select **Settings**, and then select a language from the **Language** dropdown.
|
||||
|
||||
Available languages include English, French (Français), Spanish (Español), German (Deutsch), Portuguese (Português), Japanese (日本語), and Chinese (中文).
|
||||
|
||||
- Login endpoint rate limiting
|
||||
|
||||
Langflow now applies IP-based rate limiting to the `/login` endpoint to protect against brute-force attacks.
|
||||
|
||||
For more information, see [Login rate limiting](/api-keys-and-authentication#login-rate-limiting).
|
||||
|
||||
- Code Agents bundle (beta)
|
||||
|
||||
The **Code Agents** bundle adds new beta agent components for code generation with the `smolagents` and `OpenDsStar` libraries.
|
||||
|
||||
For more information, see [Code Agents bundle](./bundles-codeagents).
|
||||
|
||||
- File Processing bundle (beta)
|
||||
|
||||
The **File Processing** bundle adds components for ingesting and retrieving file content in agent workflows.
|
||||
|
||||
For more information, see [File Processing bundle](./bundles-files-ingestion).
|
||||
|
||||
### Deprecations
|
||||
|
||||
- **Text Input** and **Text Output** components are legacy
|
||||
|
||||
The **Text Input** and **Text Output** components are now legacy and may be removed in a future release.
|
||||
Replace them with the [**Chat Input** and **Chat Output** components](/chat-input-and-output).
|
||||
|
||||
- Voice mode is removed
|
||||
|
||||
The <Icon name="Mic" aria-hidden="true"/> **Microphone** button in the **Playground** now only enables speech-to-text, with no additional voice mode functionality.
|
||||
The bi-directional `/v1/voice/ws/flow_as_tool/{flow_id}` voice-to-voice endpoints are deprecated.
|
||||
For more information, see [Use voice mode](/concepts-voice-mode).
|
||||
For 1.10.x release notes, see the [1.10.x documentation](https://docs.langflow.org/1.10.0/release-notes).
|
||||
|
||||
## 1.9.x
|
||||
|
||||
|
||||
@ -470,6 +470,7 @@ module.exports = {
|
||||
"Components/bundles-milvus",
|
||||
"Components/bundles-mistralai",
|
||||
"Components/bundles-mongodb",
|
||||
"Components/bundles-nextplaid",
|
||||
"Components/bundles-notion",
|
||||
"Components/bundles-novita",
|
||||
"Components/bundles-nvidia",
|
||||
|
||||
@ -238,9 +238,16 @@ class DatabaseVariableService(VariableService, Service):
|
||||
for variable in variables:
|
||||
value = None
|
||||
if variable.type == GENERIC_TYPE:
|
||||
if not variable.value:
|
||||
await logger.awarning("Variable '%s' has no stored value — skipping.", variable.name)
|
||||
continue
|
||||
value = auth_utils.decrypt_api_key(variable.value)
|
||||
if not value:
|
||||
# If decryption fails (likely due to encryption by different key), skip this variable
|
||||
await logger.awarning(
|
||||
"Variable '%s' could not be decrypted — likely encrypted with a different "
|
||||
"LANGFLOW_SECRET_KEY. Skipping.",
|
||||
variable.name,
|
||||
)
|
||||
continue
|
||||
|
||||
# Model validate will set value to None if credential type
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
import secrets
|
||||
from datetime import datetime
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
@ -394,6 +394,67 @@ async def test_update_generic_variable_with_fernet_signature_fails(service, sess
|
||||
await service.update_variable(user_id, "TEST_VAR", "gAAAAABthis-looks-like-encrypted", session=session)
|
||||
|
||||
|
||||
async def test_get_all__empty_value_warns_and_skips(service, session: AsyncSession):
|
||||
"""get_all warns and skips GENERIC variables whose stored value is None or empty."""
|
||||
user_id = uuid4()
|
||||
|
||||
# Create a normal generic variable, then manually blank its value to simulate a bad DB row.
|
||||
var = await service.create_variable(user_id, "EMPTY_VAR", "initial", type_=GENERIC_TYPE, session=session)
|
||||
var.value = ""
|
||||
session.add(var)
|
||||
await session.flush()
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_logger.awarning = AsyncMock()
|
||||
with patch("langflow.services.variable.service.logger", mock_logger):
|
||||
result = await service.get_all(user_id, session=session)
|
||||
|
||||
# Variable should be excluded from results.
|
||||
assert not any(v.name == "EMPTY_VAR" for v in result)
|
||||
# Warning must name the variable and mention empty value.
|
||||
warning_calls = [str(c) for c in mock_logger.awarning.call_args_list]
|
||||
assert any("EMPTY_VAR" in c for c in warning_calls)
|
||||
assert any("no stored value" in c for c in warning_calls)
|
||||
|
||||
|
||||
async def test_get_all__decrypt_failure_warns_and_skips(service, session: AsyncSession):
|
||||
"""get_all warns with a key-mismatch message when decrypt returns empty for a non-empty value."""
|
||||
user_id = uuid4()
|
||||
|
||||
await service.create_variable(user_id, "MY_VAR", "real_value", type_=GENERIC_TYPE, session=session)
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_logger.awarning = AsyncMock()
|
||||
# Simulate decrypt returning "" (key mismatch) without touching the stored value.
|
||||
with (
|
||||
patch("langflow.services.variable.service.auth_utils.decrypt_api_key", return_value=""),
|
||||
patch("langflow.services.variable.service.logger", mock_logger),
|
||||
):
|
||||
result = await service.get_all(user_id, session=session)
|
||||
|
||||
# Variable should be excluded from results.
|
||||
assert not any(v.name == "MY_VAR" for v in result)
|
||||
# Warning must name the variable and mention SECRET_KEY.
|
||||
warning_calls = [str(c) for c in mock_logger.awarning.call_args_list]
|
||||
assert any("MY_VAR" in c for c in warning_calls)
|
||||
assert any("SECRET_KEY" in c for c in warning_calls)
|
||||
|
||||
|
||||
async def test_get_all__healthy_generic_variable_included(service, session: AsyncSession):
|
||||
"""get_all includes GENERIC variables that decrypt successfully — no warnings emitted."""
|
||||
user_id = uuid4()
|
||||
|
||||
await service.create_variable(user_id, "GOOD_VAR", "good_value", type_=GENERIC_TYPE, session=session)
|
||||
|
||||
mock_logger = MagicMock()
|
||||
mock_logger.awarning = AsyncMock()
|
||||
with patch("langflow.services.variable.service.logger", mock_logger):
|
||||
result = await service.get_all(user_id, session=session)
|
||||
|
||||
assert any(v.name == "GOOD_VAR" for v in result)
|
||||
mock_logger.awarning.assert_not_called()
|
||||
|
||||
|
||||
async def test_create_credential_variable_with_fernet_signature_succeeds(service, session: AsyncSession):
|
||||
"""Test that CREDENTIAL variables can have values that look like Fernet tokens (they get encrypted anyway)."""
|
||||
user_id = uuid4()
|
||||
|
||||
@ -44,6 +44,7 @@ dependencies = [
|
||||
"filelock>=3.20.1,<4.0.0",
|
||||
"pypdf>=6.10.0,<7.0.0",
|
||||
"cryptography>=46.0.7",
|
||||
"pyjwt>=2.10.0,<3.0.0",
|
||||
"ag-ui-protocol>=0.1.10",
|
||||
"langflow-sdk>=0.1.0",
|
||||
"markitdown>=0.1.4,<2.0.0",
|
||||
|
||||
@ -177,6 +177,51 @@ def register(app: typer.Typer) -> None:
|
||||
"instead of reading from the process environment."
|
||||
),
|
||||
),
|
||||
identity_mode: str = typer.Option(
|
||||
"off",
|
||||
"--identity-mode",
|
||||
help=(
|
||||
"Per-user identity forwarding mode. 'off' (default) ignores caller identity. "
|
||||
"'jwt' verifies a signed JWT (Authorization: Bearer or --identity-jwt-header) and "
|
||||
"threads its claim as user_id. 'header' trusts a plain gateway header verbatim "
|
||||
"(weaker; safe only when network policy guarantees the gateway is the sole caller)."
|
||||
),
|
||||
),
|
||||
identity_jwt_issuer: str | None = typer.Option(
|
||||
None,
|
||||
"--identity-jwt-issuer",
|
||||
help="Expected JWT 'iss' (exact match). Required when --identity-mode=jwt.",
|
||||
),
|
||||
identity_jwt_audience: str | None = typer.Option(
|
||||
None,
|
||||
"--identity-jwt-audience",
|
||||
help="Expected JWT 'aud' (the token's audience must contain it). Required when --identity-mode=jwt.",
|
||||
),
|
||||
identity_jwt_jwks_url: str | None = typer.Option(
|
||||
None,
|
||||
"--identity-jwt-jwks-url",
|
||||
help="JWKS URL for signing keys. Optional; defaults to OIDC discovery from the issuer.",
|
||||
),
|
||||
identity_claim: str = typer.Option(
|
||||
"sub",
|
||||
"--identity-claim",
|
||||
help="JWT claim used as the caller identity (user_id). Defaults to 'sub'.",
|
||||
),
|
||||
identity_jwt_header: str = typer.Option(
|
||||
"Authorization",
|
||||
"--identity-jwt-header",
|
||||
help="Header the JWT is read from in jwt mode (a 'Bearer ' prefix is stripped). Defaults to Authorization.",
|
||||
),
|
||||
identity_header: str = typer.Option(
|
||||
"X-Consumer-Username",
|
||||
"--identity-header",
|
||||
help="Plain header trusted as identity in header mode. Defaults to X-Consumer-Username.",
|
||||
),
|
||||
identity_jwt_allow_insecure_http: bool = typer.Option(
|
||||
False,
|
||||
"--identity-jwt-allow-insecure-http",
|
||||
help="Allow plaintext http:// for the JWKS/issuer (dev/test only). HTTPS is required by default.",
|
||||
),
|
||||
) -> None:
|
||||
"""Serve LFX flows as a web API (lazy-loaded)."""
|
||||
from pathlib import Path
|
||||
@ -200,4 +245,12 @@ def register(app: typer.Typer) -> None:
|
||||
check_variables=check_variables,
|
||||
upgrade_flow=upgrade_flow,
|
||||
no_env_fallback=no_env_fallback,
|
||||
identity_mode=identity_mode,
|
||||
identity_jwt_issuer=identity_jwt_issuer,
|
||||
identity_jwt_audience=identity_jwt_audience,
|
||||
identity_jwt_jwks_url=identity_jwt_jwks_url,
|
||||
identity_claim=identity_claim,
|
||||
identity_jwt_header=identity_jwt_header,
|
||||
identity_header=identity_header,
|
||||
identity_jwt_allow_insecure_http=identity_jwt_allow_insecure_http,
|
||||
)
|
||||
|
||||
@ -279,6 +279,18 @@ def serve_command(
|
||||
"instead of reading from the process environment."
|
||||
),
|
||||
),
|
||||
# Plain defaults (not typer.Option): the CLI options live on serve_command_wrapper,
|
||||
# which passes these through. Plain defaults keep clean values even if a direct
|
||||
# caller omits them — a typer.Option default would leak an OptionInfo object
|
||||
# (e.g. mode would arrive as <typer.models.OptionInfo>, not "off").
|
||||
identity_mode: str = "off",
|
||||
identity_jwt_issuer: str | None = None,
|
||||
identity_jwt_audience: str | None = None,
|
||||
identity_jwt_jwks_url: str | None = None,
|
||||
identity_claim: str = "sub",
|
||||
identity_jwt_header: str = "Authorization",
|
||||
identity_header: str = "X-Consumer-Username",
|
||||
identity_jwt_allow_insecure_http: bool = False,
|
||||
upgrade_flow: str | None = None,
|
||||
) -> None:
|
||||
"""Serve LFX flows as a web API.
|
||||
@ -325,6 +337,23 @@ def serve_command(
|
||||
typer.echo("Error: --workers must be at least 1.", err=True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
from lfx.cli.serve_identity import IdentityConfig, IdentityConfigError
|
||||
|
||||
try:
|
||||
identity_config = IdentityConfig(
|
||||
mode=identity_mode, # type: ignore[arg-type]
|
||||
jwt_issuer=identity_jwt_issuer,
|
||||
jwt_audience=identity_jwt_audience,
|
||||
jwt_jwks_url=identity_jwt_jwks_url,
|
||||
claim=identity_claim,
|
||||
jwt_header=identity_jwt_header,
|
||||
trusted_header=identity_header,
|
||||
allow_insecure_http=identity_jwt_allow_insecure_http,
|
||||
)
|
||||
except IdentityConfigError as e:
|
||||
typer.echo(f"Error: {e}", err=True)
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
os.environ["LANGFLOW_PRETTY_LOGS"] = "false"
|
||||
configure(log_level=log_level)
|
||||
|
||||
@ -443,6 +472,11 @@ def serve_command(
|
||||
|
||||
os.environ[_SERVE_FLOW_DIR_ENV] = str(flow_dir) if flow_dir else ""
|
||||
os.environ[_SERVE_NO_ENV_FALLBACK_ENV] = "1" if no_env_fallback else "0"
|
||||
# Round-trip identity config into the worker processes (each worker keeps
|
||||
# its own JWKS cache — a few KB; no cross-worker sharing).
|
||||
identity_env = identity_config.to_env()
|
||||
for identity_key, identity_value in identity_env.items():
|
||||
os.environ[identity_key] = identity_value
|
||||
|
||||
# When flow_dir is set, startup flows are already in the store (written by
|
||||
# _build_serve_registry above) so workers load them via warm_from_store().
|
||||
@ -466,10 +500,15 @@ def serve_command(
|
||||
finally:
|
||||
# Only remove the keys we set above — a prefix sweep would also delete
|
||||
# any LFX_SERVE_* var the operator intentionally exported before launch.
|
||||
for k in (_SERVE_FLOW_DIR_ENV, _SERVE_NO_ENV_FALLBACK_ENV, _SERVE_STARTUP_PATHS_ENV):
|
||||
for k in (
|
||||
_SERVE_FLOW_DIR_ENV,
|
||||
_SERVE_NO_ENV_FALLBACK_ENV,
|
||||
_SERVE_STARTUP_PATHS_ENV,
|
||||
*identity_env,
|
||||
):
|
||||
os.environ.pop(k, None)
|
||||
else:
|
||||
serve_app = create_multi_serve_app(registry=registry)
|
||||
serve_app = create_multi_serve_app(registry=registry, identity_config=identity_config)
|
||||
uvicorn.run(serve_app, host=host, port=port, workers=1, log_level=log_level)
|
||||
except KeyboardInterrupt:
|
||||
verbose_print("\nServer stopped")
|
||||
|
||||
@ -307,7 +307,9 @@ def prepare_graph(graph, verbose_print):
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
|
||||
async def execute_graph_with_capture(graph, input_value: str | None, session_id: str | None = None):
|
||||
async def execute_graph_with_capture(
|
||||
graph, input_value: str | None, session_id: str | None = None, user_id: str | None = None
|
||||
):
|
||||
"""Execute a graph and capture output.
|
||||
|
||||
Args:
|
||||
@ -317,6 +319,13 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id:
|
||||
message-store paths (which validate session_id) succeed; an empty or
|
||||
whitespace-only string is rejected with ``ValueError`` to surface
|
||||
shell/env-var typos (see ``lfx.run._defaults.validate_provided_id``).
|
||||
user_id: Optional verified caller identity (e.g. forwarded by an edge
|
||||
gateway via a verified JWT — see ``lfx.cli.serve_identity``). ``None``
|
||||
keeps any ``user_id`` already pinned on the graph, and auto-generates
|
||||
a throwaway UUID when the graph has none (components require a
|
||||
non-empty user_id, but lfx's variable service is env-backed so the
|
||||
value is not used for scoping). A non-``None`` value overwrites the
|
||||
graph's value (the verified identity is authoritative).
|
||||
|
||||
Returns:
|
||||
Tuple of (results, captured_logs)
|
||||
@ -325,9 +334,11 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id:
|
||||
Exception: Re-raises any exception that occurs during graph execution
|
||||
"""
|
||||
# Apply session_id, user_id, and Memory-vertex propagation defaults via the
|
||||
# shared helper (same logic as run_flow). user_id is not exposed in this
|
||||
# entry point, so any pre-existing graph.user_id is preserved.
|
||||
apply_run_defaults(graph, session_id=session_id, user_id=None, overwrite_user_id=False)
|
||||
# shared helper (same logic as run_flow). A supplied (verified) user_id
|
||||
# overwrites any graph-pinned value; when None (off mode / CLI runs) the
|
||||
# graph's existing user_id is kept, or a UUID auto-generated if it has none —
|
||||
# the same result as calling this function before user_id was a parameter.
|
||||
apply_run_defaults(graph, session_id=session_id, user_id=user_id, overwrite_user_id=user_id is not None)
|
||||
|
||||
# Create input request
|
||||
inputs = InputValueRequest(input_value=input_value) if input_value else None
|
||||
|
||||
@ -25,10 +25,11 @@ import uuid
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Annotated, Any
|
||||
|
||||
from fastapi import Depends, FastAPI, HTTPException, Response, Security
|
||||
from fastapi import Depends, FastAPI, HTTPException, Request, Response, Security
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from fastapi.security import APIKeyHeader, APIKeyQuery
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
from rich.console import Console
|
||||
|
||||
from lfx.cli.common import (
|
||||
execute_graph_with_capture,
|
||||
@ -36,6 +37,7 @@ from lfx.cli.common import (
|
||||
get_api_key,
|
||||
)
|
||||
from lfx.cli.runtime_variables import apply_global_vars_to_graph
|
||||
from lfx.cli.serve_identity import IdentityConfig, build_identity_verifier
|
||||
from lfx.load import load_flow_from_json
|
||||
from lfx.log.logger import logger
|
||||
from lfx.utils.flow_validation import validate_flow_for_current_settings
|
||||
@ -46,6 +48,12 @@ if TYPE_CHECKING:
|
||||
from lfx.cli.flow_store import FlowStore
|
||||
from lfx.graph import Graph
|
||||
|
||||
# Operator-facing startup notices go to stderr via Rich (same mechanism as the
|
||||
# CLI startup banner). The structlog ``logger`` is not reliably surfaced on the
|
||||
# ``lfx serve`` stdout path under uvicorn (single- and multi-worker alike), so
|
||||
# security-critical startup warnings must use a guaranteed channel.
|
||||
_startup_console = Console(stderr=True, soft_wrap=True)
|
||||
|
||||
# Security - use the same pattern as Langflow main API
|
||||
API_KEY_NAME = "x-api-key"
|
||||
|
||||
@ -490,6 +498,7 @@ async def run_flow_generator_for_serve(
|
||||
flow_id: str,
|
||||
event_manager,
|
||||
client_consumed_queue: asyncio.Queue,
|
||||
user_id: str | None = None,
|
||||
) -> None:
|
||||
"""Executes a flow asynchronously and manages event streaming to the client.
|
||||
|
||||
@ -502,6 +511,10 @@ async def run_flow_generator_for_serve(
|
||||
flow_id (str): The ID of the flow being executed
|
||||
event_manager: Manages the streaming of events to the client
|
||||
client_consumed_queue (asyncio.Queue): Tracks client consumption of events
|
||||
user_id (str | None): Verified caller identity threaded into execution as
|
||||
the graph's ``user_id``. ``None`` means no verified identity — the
|
||||
graph's existing ``user_id`` is used, or a UUID auto-generated if it
|
||||
has none.
|
||||
|
||||
Events Generated:
|
||||
- "add_message": Sent when new messages are added during flow execution
|
||||
@ -520,7 +533,7 @@ async def run_flow_generator_for_serve(
|
||||
# Note: This is a simplified version. In a full implementation, you might want
|
||||
# to integrate with the full LFX streaming pipeline from endpoints.py
|
||||
results, logs = await execute_graph_with_capture(
|
||||
graph, input_request.input_value, session_id=input_request.session_id
|
||||
graph, input_request.input_value, session_id=input_request.session_id, user_id=user_id
|
||||
)
|
||||
result_data = extract_result_data(results, logs)
|
||||
|
||||
@ -542,11 +555,16 @@ async def run_flow_generator_for_serve(
|
||||
def create_multi_serve_app(
|
||||
*,
|
||||
registry: FlowRegistry,
|
||||
identity_config: IdentityConfig | None = None,
|
||||
) -> FastAPI:
|
||||
"""Create a FastAPI app exposing LFX flows via a mutable registry.
|
||||
|
||||
Routes dispatch to ``registry`` at request time, so flows added after
|
||||
startup (via ``POST /flows/upload/``) are immediately reachable.
|
||||
|
||||
``identity_config`` configures the optional per-user identity layer (see
|
||||
:mod:`lfx.cli.serve_identity`). ``None`` (the default) means ``off`` mode —
|
||||
no identity is read and execution behaves exactly as before.
|
||||
"""
|
||||
app = FastAPI(
|
||||
title=f"LFX Multi-Flow Server ({len(registry)})",
|
||||
@ -559,6 +577,49 @@ def create_multi_serve_app(
|
||||
)
|
||||
app.state.registry = registry
|
||||
|
||||
identity_config = identity_config or IdentityConfig()
|
||||
app.state.identity_config = identity_config
|
||||
identity_verifier = build_identity_verifier(identity_config)
|
||||
app.state.identity_verifier = identity_verifier
|
||||
if identity_verifier is not None:
|
||||
if identity_config.mode == "jwt":
|
||||
# Prefetch so the first request never pays the JWKS round-trip. A failed
|
||||
# prefetch does not abort startup (requests fail closed with 401), but the
|
||||
# operator must be told, so warn on the guaranteed stderr channel as well.
|
||||
if not identity_verifier.prefetch():
|
||||
_startup_console.print(
|
||||
"[bold red]WARNING:[/bold red] JWKS prefetch failed — lfx serve started but JWT "
|
||||
"verification will reject all requests (401) until the JWKS endpoint recovers."
|
||||
)
|
||||
elif identity_config.mode == "header":
|
||||
warning = (
|
||||
f"lfx serve identity mode=header: trusting the plain {identity_config.trusted_header!r} header "
|
||||
"as caller identity. This is only safe when network policy guarantees the gateway is the sole "
|
||||
"caller — header trust rests entirely on topology, which fails open on non-enforcing CNIs."
|
||||
)
|
||||
logger.warning(warning)
|
||||
_startup_console.print(f"[bold yellow]WARNING:[/bold yellow] {warning}")
|
||||
|
||||
def resolve_identity(request: Request, _api_key: str = Depends(verify_api_key)) -> str | None:
|
||||
"""Resolve the verified caller identity for an authenticated request.
|
||||
|
||||
Sub-depends on ``verify_api_key`` so identity processing only ever runs
|
||||
on requests that already cleared the serve-key floor — identity annotates
|
||||
authenticated requests, it never weakens or replaces the serve key.
|
||||
Returns ``None`` in ``off`` mode (no per-user attribution).
|
||||
|
||||
Deliberately a sync ``def``: FastAPI runs sync dependencies in a worker
|
||||
thread, so a rare blocking JWKS fetch (on key rotation or an issuer blip)
|
||||
never stalls the event loop. The warm path is CPU-only and ~sub-millisecond.
|
||||
|
||||
The verifier is read from ``app.state`` at request time so it can be
|
||||
swapped (e.g. by tests) after app construction.
|
||||
"""
|
||||
verifier = request.app.state.identity_verifier
|
||||
if verifier is None:
|
||||
return None
|
||||
return verifier.authenticate(request.headers)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Global endpoints
|
||||
# ------------------------------------------------------------------
|
||||
@ -683,7 +744,15 @@ def create_multi_serve_app(
|
||||
summary="Execute flow",
|
||||
dependencies=[Depends(verify_api_key)],
|
||||
)
|
||||
async def run_flow(flow_id: str, request: RunRequest) -> RunResponse:
|
||||
async def run_flow(
|
||||
flow_id: str,
|
||||
request: RunRequest,
|
||||
# Depends() lives in the default (not Annotated) so FastAPI reads the live
|
||||
# closure-local resolver — ``from __future__ import annotations`` stringizes
|
||||
# annotations, and a closure-local name can't be resolved from module globals.
|
||||
# (Hence the FAST002 suppression: the Annotated form ruff wants would break this.)
|
||||
user_id: str | None = Depends(resolve_identity), # noqa: FAST002
|
||||
) -> RunResponse:
|
||||
graph, _ = _get_flow_or_404(flow_id)
|
||||
try:
|
||||
validate_flow_for_current_settings(graph)
|
||||
@ -692,7 +761,7 @@ def create_multi_serve_app(
|
||||
registry.stamp(graph_copy)
|
||||
apply_global_vars_to_graph(graph_copy, request.global_vars)
|
||||
results, logs = await execute_graph_with_capture(
|
||||
graph_copy, request.input_value, session_id=request.session_id
|
||||
graph_copy, request.input_value, session_id=request.session_id, user_id=user_id
|
||||
)
|
||||
result_data = extract_result_data(results, logs)
|
||||
|
||||
@ -739,7 +808,11 @@ def create_multi_serve_app(
|
||||
summary="Stream flow execution",
|
||||
dependencies=[Depends(verify_api_key)],
|
||||
)
|
||||
async def stream_flow(flow_id: str, request: StreamRequest) -> StreamingResponse:
|
||||
async def stream_flow(
|
||||
flow_id: str,
|
||||
request: StreamRequest,
|
||||
user_id: str | None = Depends(resolve_identity), # noqa: FAST002 - see run_flow note on Depends-in-default
|
||||
) -> StreamingResponse:
|
||||
graph, _ = _get_flow_or_404(flow_id)
|
||||
try:
|
||||
validate_flow_for_current_settings(graph)
|
||||
@ -760,6 +833,7 @@ def create_multi_serve_app(
|
||||
flow_id=flow_id,
|
||||
event_manager=event_manager,
|
||||
client_consumed_queue=asyncio_queue_client_consumed,
|
||||
user_id=user_id,
|
||||
)
|
||||
)
|
||||
|
||||
@ -788,10 +862,11 @@ def create_serve_app() -> FastAPI:
|
||||
"""ASGI app factory called by each uvicorn worker in multi-worker mode.
|
||||
|
||||
Workers cannot inherit the parent's in-memory app object. Instead, each
|
||||
worker calls this factory, which reads ``LFX_SERVE_FLOW_DIR`` and
|
||||
``LFX_SERVE_NO_ENV_FALLBACK`` from the environment, pre-warms its own
|
||||
in-memory cache from the shared ``FilesystemFlowStore``, and returns a
|
||||
ready FastAPI app.
|
||||
worker calls this factory, which reads ``LFX_SERVE_FLOW_DIR``,
|
||||
``LFX_SERVE_NO_ENV_FALLBACK`` and the ``LFX_SERVE_IDENTITY_*`` identity
|
||||
settings from the environment, pre-warms its own in-memory cache from the
|
||||
shared ``FilesystemFlowStore`` (and, in ``jwt`` identity mode, its own JWKS
|
||||
cache), and returns a ready FastAPI app.
|
||||
|
||||
The parent process must set those env vars **before** calling
|
||||
``uvicorn.run("lfx.cli.serve_app:create_serve_app", workers=N, ...)``.
|
||||
@ -805,6 +880,7 @@ def create_serve_app() -> FastAPI:
|
||||
flow_dir_str = os.environ.get(_SERVE_FLOW_DIR_ENV)
|
||||
no_env_fallback = os.environ.get(_SERVE_NO_ENV_FALLBACK_ENV, "0") == "1"
|
||||
startup_paths_json = os.environ.get(_SERVE_STARTUP_PATHS_ENV, "")
|
||||
identity_config = IdentityConfig.from_env(os.environ)
|
||||
|
||||
flow_dir = Path(flow_dir_str) if flow_dir_str else None
|
||||
flow_store = FilesystemFlowStore(flow_dir) if flow_dir else NullFlowStore()
|
||||
@ -847,4 +923,4 @@ def create_serve_app() -> FastAPI:
|
||||
registry = FlowRegistry(no_env_fallback=no_env_fallback, store=flow_store)
|
||||
|
||||
registry.warm_from_store()
|
||||
return create_multi_serve_app(registry=registry)
|
||||
return create_multi_serve_app(registry=registry, identity_config=identity_config)
|
||||
|
||||
476
src/lfx/src/lfx/cli/serve_identity.py
Normal file
476
src/lfx/src/lfx/cli/serve_identity.py
Normal file
@ -0,0 +1,476 @@
|
||||
"""Per-user identity forwarding for ``lfx serve``.
|
||||
|
||||
This module lets ``lfx serve`` *consume* an upstream identity and thread it
|
||||
into flow execution as ``user_id``.
|
||||
|
||||
Two modes carry identity, plus an off switch:
|
||||
|
||||
``off``
|
||||
Default. No identity is read; ``user_id`` stays ``None`` and behavior is
|
||||
byte-for-byte identical to a server without this module.
|
||||
|
||||
``jwt``
|
||||
A signed JWT (``Authorization: Bearer`` or a configured header) is verified
|
||||
against the issuer's JWKS before its identity claim is trusted. Missing or
|
||||
invalid tokens are rejected with 401. This is the recommended secure mode
|
||||
when identity forwarding is enabled (note: ``off`` is the actual default).
|
||||
|
||||
``header``
|
||||
A plain gateway header (e.g. ``X-Consumer-Username``) is trusted verbatim.
|
||||
This is the *weaker* mode — its safety rests entirely on network topology
|
||||
(the gateway being the only reachable caller), so it logs a loud startup
|
||||
warning naming that assumption. It exists because Kong OSS key-auth callers
|
||||
have no JWT.
|
||||
|
||||
The verifier owns a small, per-process, TTL'd JWKS cache. It is prefetched at
|
||||
startup so the first request never pays the fetch round-trip, and JWKS fetch
|
||||
failures degrade the *request* (401), never the server.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, ClassVar, Literal
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
from fastapi import HTTPException
|
||||
|
||||
from lfx.log.logger import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Mapping
|
||||
|
||||
IdentityMode = Literal["off", "jwt", "header"]
|
||||
|
||||
# Only asymmetric signatures. Pinning these kills the two classic JWT forgeries:
|
||||
# ``alg: none`` (no signature) and HMAC alg-confusion (signing with the public
|
||||
# key as an HMAC secret). HS256/384/512 and none are rejected before any crypto.
|
||||
ALLOWED_ALGORITHMS: tuple[str, ...] = ("RS256", "ES256")
|
||||
|
||||
# JWKS cache TTL and the minimum spacing between forced refreshes triggered by an
|
||||
# unknown ``kid``. The TTL keeps keys fresh through rotation; the refresh floor
|
||||
# stops a caller spamming bogus ``kid``s from hammering the issuer's JWKS endpoint.
|
||||
_JWKS_TTL_SECONDS = 3600.0
|
||||
_JWKS_MIN_REFRESH_INTERVAL_SECONDS = 10.0
|
||||
|
||||
_HTTP_TIMEOUT_SECONDS = 10.0
|
||||
|
||||
# Env var names used to round-trip identity config into uvicorn worker processes
|
||||
# (same mechanism as ``LFX_SERVE_NO_ENV_FALLBACK``).
|
||||
ENV_MODE = "LFX_SERVE_IDENTITY_MODE"
|
||||
ENV_JWT_ISSUER = "LFX_SERVE_IDENTITY_JWT_ISSUER"
|
||||
ENV_JWT_AUDIENCE = "LFX_SERVE_IDENTITY_JWT_AUDIENCE"
|
||||
ENV_JWT_JWKS_URL = "LFX_SERVE_IDENTITY_JWT_JWKS_URL"
|
||||
ENV_CLAIM = "LFX_SERVE_IDENTITY_CLAIM"
|
||||
ENV_JWT_HEADER = "LFX_SERVE_IDENTITY_JWT_HEADER"
|
||||
ENV_TRUSTED_HEADER = "LFX_SERVE_IDENTITY_HEADER"
|
||||
ENV_ALLOW_INSECURE_HTTP = "LFX_SERVE_IDENTITY_ALLOW_INSECURE_HTTP"
|
||||
|
||||
# Every identity env key — used by the parent to clean up after uvicorn.run().
|
||||
IDENTITY_ENV_KEYS: tuple[str, ...] = (
|
||||
ENV_MODE,
|
||||
ENV_JWT_ISSUER,
|
||||
ENV_JWT_AUDIENCE,
|
||||
ENV_JWT_JWKS_URL,
|
||||
ENV_CLAIM,
|
||||
ENV_JWT_HEADER,
|
||||
ENV_TRUSTED_HEADER,
|
||||
ENV_ALLOW_INSECURE_HTTP,
|
||||
)
|
||||
|
||||
|
||||
class IdentityConfigError(ValueError):
|
||||
"""Raised when identity options are incomplete or inconsistent."""
|
||||
|
||||
|
||||
def _fetch_json(url: str) -> dict:
|
||||
"""Fetch and parse a JSON document over http(s).
|
||||
|
||||
Only http/https URLs are allowed — the scheme guard prevents ``file://``/``ftp://``
|
||||
smuggling via a misconfigured issuer or JWKS URL.
|
||||
|
||||
Redirects are NOT followed. Following a redirect would let a misconfigured or
|
||||
malicious issuer downgrade ``https://`` → ``http://`` (or point off-host), and the
|
||||
plaintext follow-up request would already be on the wire before any scheme check
|
||||
could run. We reject the redirect instead so the fetch fails closed rather than
|
||||
silently downgrading the transport.
|
||||
"""
|
||||
if not url.lower().startswith(("http://", "https://")):
|
||||
msg = f"Refusing to fetch non-http(s) URL: {url!r}"
|
||||
raise IdentityConfigError(msg)
|
||||
with httpx.Client(timeout=_HTTP_TIMEOUT_SECONDS, follow_redirects=False) as client:
|
||||
resp = client.get(url)
|
||||
if resp.is_redirect:
|
||||
location = resp.headers.get("location", "")
|
||||
msg = f"Refusing to follow redirect from {url!r} to {location!r} (possible transport downgrade)."
|
||||
raise IdentityConfigError(msg)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IdentityConfig:
|
||||
"""Immutable identity-layer configuration resolved at serve startup."""
|
||||
|
||||
mode: IdentityMode = "off"
|
||||
jwt_issuer: str | None = None
|
||||
jwt_audience: str | None = None
|
||||
jwt_jwks_url: str | None = None
|
||||
claim: str = "sub"
|
||||
# Header the JWT is read from in ``jwt`` mode (Bearer prefix stripped if present).
|
||||
jwt_header: str = "Authorization"
|
||||
# Header trusted verbatim in ``header`` mode.
|
||||
trusted_header: str = "X-Consumer-Username"
|
||||
# Allow plaintext http:// for the JWKS/issuer (dev/test only). Off by default:
|
||||
# a MITM of a plaintext JWKS response forges identities, so HTTPS is required.
|
||||
allow_insecure_http: bool = False
|
||||
|
||||
def _require_fetch_scheme(self, label: str, url: str) -> None:
|
||||
"""Reject a fetched URL whose scheme violates the transport policy.
|
||||
|
||||
HTTPS is required by default; ``http://`` is permitted only when
|
||||
``allow_insecure_http`` is set. Non-http(s) schemes are always refused.
|
||||
"""
|
||||
allowed = ("http://", "https://") if self.allow_insecure_http else ("https://",)
|
||||
if not url.lower().startswith(allowed):
|
||||
scheme = "an http(s)" if self.allow_insecure_http else "an https"
|
||||
hint = "" if self.allow_insecure_http else " Pass --identity-jwt-allow-insecure-http for local/dev http://."
|
||||
msg = f"{label} must be {scheme} URL, got {url!r}.{hint}"
|
||||
raise IdentityConfigError(msg)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.mode not in ("off", "jwt", "header"):
|
||||
msg = f"Unknown identity mode {self.mode!r}; expected off, jwt, or header."
|
||||
raise IdentityConfigError(msg)
|
||||
if self.mode == "jwt":
|
||||
if not self.jwt_issuer:
|
||||
msg = "--identity-jwt-issuer is required when --identity-mode=jwt."
|
||||
raise IdentityConfigError(msg)
|
||||
if not self.jwt_audience:
|
||||
msg = "--identity-jwt-audience is required when --identity-mode=jwt."
|
||||
raise IdentityConfigError(msg)
|
||||
if not self.claim:
|
||||
msg = "--identity-claim must be non-empty when --identity-mode=jwt."
|
||||
raise IdentityConfigError(msg)
|
||||
# Validate the fetched URL scheme offline so a bad scheme (e.g. file://)
|
||||
# or an insecure http:// fails at startup rather than masquerading as a
|
||||
# transient fetch error later. An explicit jwks_url is fetched directly;
|
||||
# otherwise the issuer is fetched for OIDC discovery, so it must comply too.
|
||||
if self.jwt_jwks_url:
|
||||
self._require_fetch_scheme("--identity-jwt-jwks-url", self.jwt_jwks_url)
|
||||
else:
|
||||
self._require_fetch_scheme("--identity-jwt-issuer (used for OIDC discovery)", self.jwt_issuer)
|
||||
if self.mode == "header" and not self.trusted_header:
|
||||
msg = "--identity-header must be non-empty when --identity-mode=header."
|
||||
raise IdentityConfigError(msg)
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self.mode != "off"
|
||||
|
||||
def to_env(self) -> dict[str, str]:
|
||||
"""Serialize to env vars for the uvicorn worker round-trip.
|
||||
|
||||
Only ``off`` emits a lone mode key; richer modes emit their settings so a
|
||||
worker's :meth:`from_env` reconstructs an identical config.
|
||||
"""
|
||||
env = {ENV_MODE: self.mode}
|
||||
if self.mode == "off":
|
||||
return env
|
||||
if self.jwt_issuer:
|
||||
env[ENV_JWT_ISSUER] = self.jwt_issuer
|
||||
if self.jwt_audience:
|
||||
env[ENV_JWT_AUDIENCE] = self.jwt_audience
|
||||
if self.jwt_jwks_url:
|
||||
env[ENV_JWT_JWKS_URL] = self.jwt_jwks_url
|
||||
env[ENV_CLAIM] = self.claim
|
||||
env[ENV_JWT_HEADER] = self.jwt_header
|
||||
env[ENV_TRUSTED_HEADER] = self.trusted_header
|
||||
if self.allow_insecure_http:
|
||||
env[ENV_ALLOW_INSECURE_HTTP] = "1"
|
||||
return env
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, environ: Mapping[str, str]) -> IdentityConfig:
|
||||
"""Reconstruct config from worker env vars; absence ⇒ ``off``."""
|
||||
mode = environ.get(ENV_MODE, "off")
|
||||
if mode == "off":
|
||||
return cls(mode="off")
|
||||
defaults = cls()
|
||||
return cls(
|
||||
mode=mode, # type: ignore[arg-type]
|
||||
jwt_issuer=environ.get(ENV_JWT_ISSUER),
|
||||
jwt_audience=environ.get(ENV_JWT_AUDIENCE),
|
||||
jwt_jwks_url=environ.get(ENV_JWT_JWKS_URL),
|
||||
claim=environ.get(ENV_CLAIM, defaults.claim),
|
||||
jwt_header=environ.get(ENV_JWT_HEADER, defaults.jwt_header),
|
||||
trusted_header=environ.get(ENV_TRUSTED_HEADER, defaults.trusted_header),
|
||||
allow_insecure_http=environ.get(ENV_ALLOW_INSECURE_HTTP) == "1",
|
||||
)
|
||||
|
||||
|
||||
class _SigningKeyUnavailableError(Exception):
|
||||
"""Internal: no signing key for the presented ``kid`` (after a bounded refresh)."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class _JwksCache:
|
||||
"""A per-process, TTL'd cache of the issuer's signing keys.
|
||||
|
||||
Holds a parsed ``PyJWKSet``. Refreshes on TTL expiry and, at most once per
|
||||
cooldown window, when a request presents an unknown ``kid`` (key rotation).
|
||||
Fetch failures keep the existing keys and surface as a per-request error.
|
||||
|
||||
The stale/cold-cache and unknown-``kid`` fetch paths have independent
|
||||
cooldown floors (``_last_stale_fetch_at`` / ``_last_unknown_refresh_at``),
|
||||
each allowing one fetch per ``min_refresh_interval`` — so the worst case is
|
||||
~two fetches per window, not one, which still bounds a JWKS outage away from
|
||||
a per-request fetch storm. The floors are best-effort and lock-free: under
|
||||
concurrent requests two threads can both decide to fetch, so the cap is
|
||||
approximate, not exact (a redundant fetch is harmless; a lock around the
|
||||
blocking fetch would be worse).
|
||||
"""
|
||||
|
||||
fetch_jwks: Callable[[], dict]
|
||||
ttl_seconds: float = _JWKS_TTL_SECONDS
|
||||
min_refresh_interval: float = _JWKS_MIN_REFRESH_INTERVAL_SECONDS
|
||||
source_label: str = "the JWKS endpoint" # human-readable source for log messages
|
||||
_jwk_set: jwt.PyJWKSet | None = field(default=None, init=False)
|
||||
_fetched_at: float | None = field(default=None, init=False)
|
||||
_last_stale_fetch_at: float | None = field(default=None, init=False)
|
||||
_last_unknown_refresh_at: float | None = field(default=None, init=False)
|
||||
|
||||
def _fetch_into_cache(self) -> None:
|
||||
"""Best-effort refresh. On failure, keep the existing keys and log loudly.
|
||||
|
||||
A permanent configuration error (bad/undiscoverable JWKS source) is logged
|
||||
distinctly from a transient transport/parse failure, so the operator isn't
|
||||
told to wait for a recovery that can never come. Either way the cache is
|
||||
left intact and requests fail closed (401); serving never aborts here.
|
||||
Unexpected exceptions are intentionally NOT swallowed — they surface as a
|
||||
real error rather than a silent "keeping cached keys" line.
|
||||
"""
|
||||
try:
|
||||
data = self.fetch_jwks()
|
||||
# Guard the shape before PyJWKSet.from_dict: a non-object JWKS (e.g. a
|
||||
# JSON array) makes from_dict raise an uncaught AttributeError that
|
||||
# would abort startup or 500 a request instead of failing closed.
|
||||
if not isinstance(data, dict):
|
||||
msg = f"{self.source_label} returned a non-object JSON document (expected a JWKS)."
|
||||
raise IdentityConfigError(msg)
|
||||
self._jwk_set = jwt.PyJWKSet.from_dict(data)
|
||||
self._fetched_at = time.monotonic()
|
||||
except IdentityConfigError as exc:
|
||||
logger.error(
|
||||
f"JWKS source is misconfigured ({self.source_label}): {exc} "
|
||||
"This is a configuration error, not a transient outage — requests will be "
|
||||
"rejected (401) until it is fixed."
|
||||
)
|
||||
except (httpx.HTTPError, OSError, TimeoutError, json.JSONDecodeError, jwt.PyJWTError) as exc:
|
||||
logger.error(
|
||||
f"JWKS refresh from {self.source_label} failed ({type(exc).__name__}: {exc}); keeping cached keys."
|
||||
)
|
||||
|
||||
def prefetch(self) -> bool:
|
||||
"""Warm the cache at startup. Logs an error if it stays empty, never raises.
|
||||
|
||||
Returns ``True`` if the cache holds keys after the attempt, ``False`` if it
|
||||
stayed empty, so callers can surface a startup warning on a guaranteed channel.
|
||||
"""
|
||||
self._fetch_into_cache()
|
||||
if self._jwk_set is None:
|
||||
logger.error(
|
||||
"JWKS prefetch failed; serving will start but JWT verification will reject requests "
|
||||
"(401) until the JWKS endpoint recovers."
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
def _lookup(self, kid: str) -> jwt.PyJWK | None:
|
||||
if self._jwk_set is None:
|
||||
return None
|
||||
for key in self._jwk_set.keys:
|
||||
if key.key_id == kid:
|
||||
return key
|
||||
return None
|
||||
|
||||
def get_signing_key(self, kid: str) -> jwt.PyJWK:
|
||||
now = time.monotonic()
|
||||
# Stale/cold cache: refresh, but at most once per cooldown so a down issuer
|
||||
# cannot make every request pay its own blocking fetch. The floor is NOT
|
||||
# primed by prefetch(), so the first request after a failed prefetch still
|
||||
# refetches immediately and recovers.
|
||||
fetched_this_call = False
|
||||
stale = self._jwk_set is None or (self._fetched_at is not None and now - self._fetched_at > self.ttl_seconds)
|
||||
if stale and (
|
||||
self._last_stale_fetch_at is None or (now - self._last_stale_fetch_at) >= self.min_refresh_interval
|
||||
):
|
||||
self._last_stale_fetch_at = now
|
||||
self._fetch_into_cache()
|
||||
fetched_this_call = True
|
||||
|
||||
key = self._lookup(kid)
|
||||
if key is not None:
|
||||
return key
|
||||
|
||||
# Unknown kid (key rotation): refresh at most once per cooldown window (DoS floor).
|
||||
# Skip if we already refetched above this call — it would hit the same endpoint.
|
||||
if not fetched_this_call and (
|
||||
self._last_unknown_refresh_at is None or (now - self._last_unknown_refresh_at) >= self.min_refresh_interval
|
||||
):
|
||||
self._last_unknown_refresh_at = now
|
||||
self._fetch_into_cache()
|
||||
key = self._lookup(kid)
|
||||
|
||||
if key is None:
|
||||
raise _SigningKeyUnavailableError(kid)
|
||||
return key
|
||||
|
||||
|
||||
class IdentityVerifier:
|
||||
"""Resolves a verified ``user_id`` from request headers per :class:`IdentityConfig`."""
|
||||
|
||||
_ALLOWED_ALGS: ClassVar[frozenset[str]] = frozenset(ALLOWED_ALGORITHMS)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: IdentityConfig,
|
||||
*,
|
||||
jwks_fetcher: Callable[[str], dict] | None = None,
|
||||
openid_fetcher: Callable[[str], dict] | None = None,
|
||||
) -> None:
|
||||
self._config = config
|
||||
self._json_fetch = jwks_fetcher or _fetch_json
|
||||
self._openid_fetch = openid_fetcher or jwks_fetcher or _fetch_json
|
||||
self._resolved_jwks_url: str | None = config.jwt_jwks_url
|
||||
self._jwks: _JwksCache | None = None
|
||||
if config.mode == "jwt":
|
||||
source_label = config.jwt_jwks_url or f"OIDC discovery from {config.jwt_issuer}"
|
||||
self._jwks = _JwksCache(fetch_jwks=self._fetch_signing_jwks, source_label=source_label)
|
||||
|
||||
@property
|
||||
def config(self) -> IdentityConfig:
|
||||
return self._config
|
||||
|
||||
# -- startup -----------------------------------------------------------
|
||||
|
||||
def prefetch(self) -> bool:
|
||||
"""In ``jwt`` mode, warm the JWKS cache; otherwise a no-op.
|
||||
|
||||
Returns ``True`` when the cache is warm (or there is nothing to warm, as in
|
||||
``header`` mode), ``False`` when a ``jwt``-mode prefetch left the cache empty.
|
||||
"""
|
||||
if self._jwks is not None:
|
||||
return self._jwks.prefetch()
|
||||
return True
|
||||
|
||||
def _resolve_jwks_url(self) -> str:
|
||||
if self._resolved_jwks_url:
|
||||
return self._resolved_jwks_url
|
||||
if not self._config.jwt_issuer:
|
||||
msg = "Cannot resolve JWKS URL without an issuer."
|
||||
raise IdentityConfigError(msg)
|
||||
discovery_url = self._config.jwt_issuer.rstrip("/") + "/.well-known/openid-configuration"
|
||||
document = self._openid_fetch(discovery_url)
|
||||
# A well-behaved discovery doc is a JSON object with a string jwks_uri.
|
||||
# Guard the shape so a malformed response (e.g. a JSON array) fails closed
|
||||
# as a config error instead of an uncaught AttributeError at request time.
|
||||
if not isinstance(document, dict):
|
||||
msg = f"OIDC discovery at {discovery_url} returned a non-object document."
|
||||
raise IdentityConfigError(msg)
|
||||
jwks_uri = document.get("jwks_uri")
|
||||
if not isinstance(jwks_uri, str) or not jwks_uri:
|
||||
msg = f"OIDC discovery at {discovery_url} did not return a string jwks_uri."
|
||||
raise IdentityConfigError(msg)
|
||||
self._config._require_fetch_scheme("OIDC discovery jwks_uri", jwks_uri) # noqa: SLF001
|
||||
self._resolved_jwks_url = jwks_uri
|
||||
return jwks_uri
|
||||
|
||||
def _fetch_signing_jwks(self) -> dict:
|
||||
return self._json_fetch(self._resolve_jwks_url())
|
||||
|
||||
# -- per-request -------------------------------------------------------
|
||||
|
||||
def authenticate(self, headers: Mapping[str, str]) -> str | None:
|
||||
"""Return the caller's identity, or raise ``HTTPException(401)``.
|
||||
|
||||
``off`` mode returns ``None`` (no identity); the caller treats that as
|
||||
"no per-user attribution", preserving the pre-identity behavior.
|
||||
"""
|
||||
if self._config.mode == "off":
|
||||
return None
|
||||
if self._config.mode == "header":
|
||||
return self._authenticate_header(headers)
|
||||
return self._authenticate_jwt(headers)
|
||||
|
||||
def _authenticate_header(self, headers: Mapping[str, str]) -> str:
|
||||
value = headers.get(self._config.trusted_header)
|
||||
if not value or not value.strip():
|
||||
raise HTTPException(status_code=401, detail="Missing identity header")
|
||||
return value.strip()
|
||||
|
||||
def _extract_token(self, headers: Mapping[str, str]) -> str | None:
|
||||
raw = headers.get(self._config.jwt_header)
|
||||
if not raw:
|
||||
return None
|
||||
raw = raw.strip()
|
||||
if raw.lower().startswith("bearer "):
|
||||
return raw[len("bearer ") :].strip()
|
||||
return raw
|
||||
|
||||
def _authenticate_jwt(self, headers: Mapping[str, str]) -> str:
|
||||
token = self._extract_token(headers)
|
||||
if not token:
|
||||
raise HTTPException(status_code=401, detail="Missing identity token")
|
||||
|
||||
# Inspect the unverified header FIRST so alg-confusion and ``alg: none``
|
||||
# are rejected before any key resolution or crypto runs.
|
||||
try:
|
||||
unverified_header = jwt.get_unverified_header(token)
|
||||
except jwt.PyJWTError as exc:
|
||||
raise HTTPException(status_code=401, detail="Malformed identity token") from exc
|
||||
|
||||
algorithm = unverified_header.get("alg")
|
||||
if algorithm not in self._ALLOWED_ALGS:
|
||||
raise HTTPException(status_code=401, detail="Unsupported token algorithm")
|
||||
kid = unverified_header.get("kid")
|
||||
if not kid:
|
||||
raise HTTPException(status_code=401, detail="Token missing key id")
|
||||
|
||||
if self._jwks is None: # pragma: no cover - constructed for jwt mode
|
||||
raise HTTPException(status_code=401, detail="Identity verification unavailable")
|
||||
try:
|
||||
signing_key = self._jwks.get_signing_key(kid)
|
||||
except _SigningKeyUnavailableError as exc:
|
||||
raise HTTPException(status_code=401, detail="Unknown token signing key") from exc
|
||||
|
||||
try:
|
||||
claims = jwt.decode(
|
||||
token,
|
||||
signing_key.key,
|
||||
algorithms=list(ALLOWED_ALGORITHMS),
|
||||
audience=self._config.jwt_audience,
|
||||
issuer=self._config.jwt_issuer,
|
||||
leeway=60,
|
||||
options={"require": ["exp"], "verify_aud": True, "verify_iss": True},
|
||||
)
|
||||
except jwt.PyJWTError as exc:
|
||||
# Static detail only — never echo the token or signature material.
|
||||
raise HTTPException(status_code=401, detail="Invalid identity token") from exc
|
||||
|
||||
identity = claims.get(self._config.claim)
|
||||
# Reject missing, empty, and whitespace-only string claims; a non-string claim
|
||||
# (e.g. a numeric sub) is left to str() below.
|
||||
if identity is None or (isinstance(identity, str) and not identity.strip()):
|
||||
raise HTTPException(status_code=401, detail=f"Token missing claim '{self._config.claim}'")
|
||||
return str(identity)
|
||||
|
||||
|
||||
def build_identity_verifier(config: IdentityConfig) -> IdentityVerifier | None:
|
||||
"""Construct a verifier for an enabled config, or ``None`` for ``off`` mode."""
|
||||
if not config.enabled:
|
||||
return None
|
||||
return IdentityVerifier(config)
|
||||
@ -1432,7 +1432,7 @@ class TestServeAppEndpoints:
|
||||
headers = {"x-api-key": "test-api-key"}
|
||||
|
||||
# Mock execute_graph_with_capture to raise an error
|
||||
async def mock_execute_error(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_error(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
msg = "Flow execution failed"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
@ -1457,7 +1457,7 @@ class TestServeAppEndpoints:
|
||||
headers = {"x-api-key": "test-api-key"}
|
||||
|
||||
# Mock execute_graph_with_capture to return empty results
|
||||
async def mock_execute_empty(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_empty(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
return [], "" # Empty results and logs
|
||||
|
||||
with (
|
||||
@ -1478,7 +1478,7 @@ class TestServeAppEndpoints:
|
||||
"""The /run endpoint must forward session_id from RunRequest to the executor."""
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["session_id"] = session_id
|
||||
return [], ""
|
||||
|
||||
@ -1497,7 +1497,7 @@ class TestServeAppEndpoints:
|
||||
"""The /stream endpoint must forward session_id from StreamRequest to the executor."""
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["session_id"] = session_id
|
||||
return [], ""
|
||||
|
||||
@ -1592,7 +1592,7 @@ class TestServeAppEndpoints:
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["request_variables"] = dict(graph.context.get("request_variables") or {})
|
||||
return [], ""
|
||||
|
||||
@ -1628,7 +1628,7 @@ class TestServeAppEndpoints:
|
||||
app = create_multi_serve_app(registry=registry)
|
||||
monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True)
|
||||
|
||||
async def mock_execute_noop(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_noop(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
return [], ""
|
||||
|
||||
headers = {"x-api-key": "test-api-key"}
|
||||
@ -1670,7 +1670,7 @@ class TestServeAppEndpoints:
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["no_env_fallback"] = graph.context.get("no_env_fallback")
|
||||
return [], ""
|
||||
|
||||
@ -1705,7 +1705,7 @@ class TestServeAppEndpoints:
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["no_env_fallback"] = graph.context.get("no_env_fallback")
|
||||
return [], ""
|
||||
|
||||
@ -1744,7 +1744,7 @@ class TestServeAppEndpoints:
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["request_variables"] = dict(graph.context.get("request_variables") or {})
|
||||
return [], ""
|
||||
|
||||
@ -1786,7 +1786,7 @@ class TestServeAppEndpoints:
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_capture(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["request_variables"] = graph.context.get("request_variables")
|
||||
return [], ""
|
||||
|
||||
@ -1823,7 +1823,7 @@ class TestServeAppEndpoints:
|
||||
app = create_multi_serve_app(registry=registry)
|
||||
monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True)
|
||||
|
||||
async def mock_execute_noop(graph, input_value, session_id=None): # noqa: ARG001
|
||||
async def mock_execute_noop(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
return [], ""
|
||||
|
||||
headers = {"x-api-key": "test-api-key"}
|
||||
|
||||
@ -275,7 +275,14 @@ class TestMultiServeStreaming:
|
||||
"""Test error handling in streaming endpoint."""
|
||||
with patch("lfx.cli.serve_app.run_flow_generator_for_serve") as mock_generator:
|
||||
# Mock an error in the generator that properly terminates the stream
|
||||
async def mock_error_generator(graph, input_request, flow_id, event_manager, client_consumed_queue): # noqa: ARG001
|
||||
async def mock_error_generator(
|
||||
graph, # noqa: ARG001
|
||||
input_request, # noqa: ARG001
|
||||
flow_id, # noqa: ARG001
|
||||
event_manager,
|
||||
client_consumed_queue, # noqa: ARG001
|
||||
user_id=None, # noqa: ARG001
|
||||
):
|
||||
try:
|
||||
msg = "Test error during streaming"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
939
src/lfx/tests/unit/cli/test_serve_identity.py
Normal file
939
src/lfx/tests/unit/cli/test_serve_identity.py
Normal file
@ -0,0 +1,939 @@
|
||||
"""Tests for per-user identity forwarding in ``lfx serve`` (verified JWT / header).
|
||||
|
||||
Crypto is exercised with locally generated RSA and EC P-256 keypairs and an
|
||||
in-memory fake JWKS endpoint — no network. See ``lfx.cli.serve_identity``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives.asymmetric import ec, rsa
|
||||
from fastapi import HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
from jwt.algorithms import ECAlgorithm, RSAAlgorithm
|
||||
from lfx.cli.common import execute_graph_with_capture
|
||||
from lfx.cli.serve_app import FlowMeta, FlowRegistry, create_multi_serve_app, create_serve_app
|
||||
from lfx.cli.serve_identity import (
|
||||
IdentityConfig,
|
||||
IdentityConfigError,
|
||||
IdentityVerifier,
|
||||
_fetch_json,
|
||||
build_identity_verifier,
|
||||
)
|
||||
from lfx.graph import Graph
|
||||
|
||||
ISSUER = "https://accounts.example.com"
|
||||
AUDIENCE = "my-oauth-client-id"
|
||||
KID = "test-key-1"
|
||||
EC_KID = "test-key-ec"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Crypto / JWKS test helpers (no network)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_keypair() -> rsa.RSAPrivateKey:
|
||||
return rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
|
||||
|
||||
def _jwks_for(public_key, kid: str = KID) -> dict:
|
||||
"""Build a JWKS document exposing an RSA ``public_key`` under ``kid``."""
|
||||
jwk = json.loads(RSAAlgorithm.to_jwk(public_key))
|
||||
jwk.update({"kid": kid, "use": "sig", "alg": "RS256"})
|
||||
return {"keys": [jwk]}
|
||||
|
||||
|
||||
def _make_ec_keypair() -> ec.EllipticCurvePrivateKey:
|
||||
return ec.generate_private_key(ec.SECP256R1())
|
||||
|
||||
|
||||
def _ec_jwks_for(public_key, kid: str = EC_KID) -> dict:
|
||||
"""Build a JWKS document exposing an EC P-256 ``public_key`` under ``kid``."""
|
||||
jwk = json.loads(ECAlgorithm.to_jwk(public_key))
|
||||
jwk.update({"kid": kid, "use": "sig", "alg": "ES256"})
|
||||
return {"keys": [jwk]}
|
||||
|
||||
|
||||
def _sign(private_key, claims: dict, *, kid: str = KID, alg: str = "RS256") -> str:
|
||||
payload = {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, **claims}
|
||||
return jwt.encode(payload, private_key, algorithm=alg, headers={"kid": kid})
|
||||
|
||||
|
||||
def _alg_none_token(claims: dict, *, kid: str = KID) -> str:
|
||||
"""Hand-craft an ``alg: none`` token (unsigned)."""
|
||||
|
||||
def b64(obj: dict) -> str:
|
||||
return base64.urlsafe_b64encode(json.dumps(obj).encode()).rstrip(b"=").decode()
|
||||
|
||||
payload = {"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, **claims}
|
||||
return f"{b64({'alg': 'none', 'typ': 'JWT', 'kid': kid})}.{b64(payload)}."
|
||||
|
||||
|
||||
class _CountingFetcher:
|
||||
"""A fake JWKS fetcher that counts calls and can be flipped to fail."""
|
||||
|
||||
def __init__(self, jwks: dict) -> None:
|
||||
self.jwks = jwks
|
||||
self.calls = 0
|
||||
self.fail = False
|
||||
|
||||
def __call__(self, _url: str) -> dict:
|
||||
self.calls += 1
|
||||
if self.fail:
|
||||
msg = "simulated JWKS endpoint down"
|
||||
raise OSError(msg)
|
||||
return self.jwks
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def keypair():
|
||||
return _make_keypair()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def jwks(keypair):
|
||||
return _jwks_for(keypair.public_key())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fetcher(jwks):
|
||||
return _CountingFetcher(jwks)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def jwt_config():
|
||||
# Explicit jwks_url avoids OIDC discovery so the fake fetcher serves the JWKS directly.
|
||||
return IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://accounts.example.com/jwks",
|
||||
claim="email",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def verifier(jwt_config, fetcher):
|
||||
v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher)
|
||||
v.prefetch()
|
||||
return v
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IdentityConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIdentityConfig:
|
||||
def test_off_is_default_and_disabled(self):
|
||||
cfg = IdentityConfig()
|
||||
assert cfg.mode == "off"
|
||||
assert cfg.enabled is False
|
||||
assert cfg.to_env() == {"LFX_SERVE_IDENTITY_MODE": "off"}
|
||||
|
||||
def test_jwt_requires_issuer_and_audience(self):
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(mode="jwt", jwt_audience=AUDIENCE)
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(mode="jwt", jwt_issuer=ISSUER)
|
||||
|
||||
def test_unknown_mode_rejected(self):
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(mode="bogus") # type: ignore[arg-type]
|
||||
|
||||
def test_non_http_jwks_url_rejected_offline(self):
|
||||
# A bad scheme must fail fast at config construction, not at first fetch.
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="file:///etc/passwd")
|
||||
|
||||
def test_http_jwks_url_rejected_by_default(self):
|
||||
# Plaintext http:// is MITM-able (forged JWKS), so HTTPS is required by default.
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(
|
||||
mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="http://issuer.test/jwks.json"
|
||||
)
|
||||
|
||||
def test_http_jwks_url_allowed_with_insecure_flag(self):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="http://issuer.test/jwks.json",
|
||||
allow_insecure_http=True,
|
||||
)
|
||||
assert cfg.jwt_jwks_url == "http://issuer.test/jwks.json"
|
||||
|
||||
def test_http_issuer_for_discovery_rejected_by_default(self):
|
||||
# No explicit jwks_url → the issuer is fetched for OIDC discovery, so its
|
||||
# scheme is subject to the same HTTPS policy.
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(mode="jwt", jwt_issuer="http://accounts.example.com", jwt_audience=AUDIENCE, claim="email")
|
||||
|
||||
def test_http_issuer_for_discovery_allowed_with_insecure_flag(self):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer="http://accounts.example.com",
|
||||
jwt_audience=AUDIENCE,
|
||||
claim="email",
|
||||
allow_insecure_http=True,
|
||||
)
|
||||
assert cfg.jwt_issuer == "http://accounts.example.com"
|
||||
|
||||
def test_insecure_http_flag_still_rejects_non_http_schemes(self):
|
||||
# The escape hatch widens to http:// only — it must not re-open file://.
|
||||
with pytest.raises(IdentityConfigError):
|
||||
IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="file:///etc/passwd",
|
||||
allow_insecure_http=True,
|
||||
)
|
||||
|
||||
def test_allow_insecure_http_env_round_trip(self):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="http://issuer.test/jwks.json",
|
||||
claim="email",
|
||||
allow_insecure_http=True,
|
||||
)
|
||||
assert IdentityConfig.from_env(cfg.to_env()) == cfg
|
||||
|
||||
def test_env_round_trip_preserves_all_fields(self):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://example.com/jwks",
|
||||
claim="email",
|
||||
jwt_header="X-Auth-Request-Access-Token",
|
||||
trusted_header="X-Consumer-Username",
|
||||
)
|
||||
assert IdentityConfig.from_env(cfg.to_env()) == cfg
|
||||
|
||||
def test_env_round_trip_header_mode(self):
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
assert IdentityConfig.from_env(cfg.to_env()) == cfg
|
||||
|
||||
def test_from_env_absent_is_off(self):
|
||||
assert IdentityConfig.from_env({}) == IdentityConfig(mode="off")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IdentityVerifier — JWT mode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIdentityVerifierJwt:
|
||||
def _headers(self, token: str) -> dict:
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
def test_valid_token_returns_identity_claim(self, verifier, keypair):
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
assert verifier.authenticate(self._headers(token)) == "alice@example.com"
|
||||
|
||||
def test_missing_token_rejected(self, verifier):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate({})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_bad_signature_rejected(self, verifier):
|
||||
other_key = _make_keypair()
|
||||
token = _sign(other_key, {"email": "mallory@example.com"}) # signed by wrong key, valid kid
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_wrong_issuer_rejected(self, verifier, keypair):
|
||||
token = jwt.encode(
|
||||
{"iss": "https://evil.example.com", "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "a@b.c"},
|
||||
keypair,
|
||||
algorithm="RS256",
|
||||
headers={"kid": KID},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_wrong_audience_rejected(self, verifier, keypair):
|
||||
token = jwt.encode(
|
||||
{"iss": ISSUER, "aud": "some-other-client", "exp": int(time.time()) + 3600, "email": "a@b.c"},
|
||||
keypair,
|
||||
algorithm="RS256",
|
||||
headers={"kid": KID},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_expired_token_rejected(self, verifier, keypair):
|
||||
token = jwt.encode(
|
||||
{"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) - 3600, "email": "a@b.c"},
|
||||
keypair,
|
||||
algorithm="RS256",
|
||||
headers={"kid": KID},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_alg_none_rejected(self, verifier):
|
||||
token = _alg_none_token({"email": "a@b.c"})
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_hs256_alg_confusion_rejected(self, verifier, keypair):
|
||||
# Attacker signs with HMAC using the PUBLIC key bytes as the shared secret.
|
||||
# pyjwt refuses to use an asymmetric key for HMAC, so hand-craft the token.
|
||||
import hashlib
|
||||
import hmac
|
||||
|
||||
from cryptography.hazmat.primitives import serialization
|
||||
|
||||
public_pem = keypair.public_key().public_bytes(
|
||||
encoding=serialization.Encoding.PEM,
|
||||
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
||||
)
|
||||
|
||||
def b64(raw: bytes) -> str:
|
||||
return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()
|
||||
|
||||
header = b64(json.dumps({"alg": "HS256", "typ": "JWT", "kid": KID}).encode())
|
||||
payload = b64(
|
||||
json.dumps(
|
||||
{"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "mallory@example.com"}
|
||||
).encode()
|
||||
)
|
||||
signing_input = f"{header}.{payload}".encode()
|
||||
signature = b64(hmac.new(public_pem, signing_input, hashlib.sha256).digest())
|
||||
forged = f"{header}.{payload}.{signature}"
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(forged))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_missing_claim_rejected(self, verifier, keypair):
|
||||
token = _sign(keypair, {"sub": "no-email-here"}) # claim is 'email'
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
@pytest.mark.parametrize("blank", ["", " ", "\t", "\n "])
|
||||
def test_blank_claim_rejected(self, verifier, keypair, blank):
|
||||
# An empty or whitespace-only claim must not become a (blank) user_id.
|
||||
token = _sign(keypair, {"email": blank})
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_missing_kid_rejected(self, verifier, keypair):
|
||||
# Token with no `kid` in its header — distinct from an unknown kid.
|
||||
token = jwt.encode(
|
||||
{"iss": ISSUER, "aud": AUDIENCE, "exp": int(time.time()) + 3600, "email": "a@b.c"},
|
||||
keypair,
|
||||
algorithm="RS256", # no headers={"kid": ...}
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
assert "key id" in exc.value.detail.lower()
|
||||
|
||||
def test_unknown_kid_triggers_one_rate_limited_refresh_then_401(self, verifier, fetcher, keypair):
|
||||
assert fetcher.calls == 1 # prefetch
|
||||
token = _sign(keypair, {"email": "a@b.c"}, kid="rotated-kid-not-in-jwks")
|
||||
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert exc.value.status_code == 401
|
||||
assert fetcher.calls == 2 # one refresh attempt on the unknown kid
|
||||
|
||||
# A second immediate unknown-kid request is rate-limited: no extra fetch.
|
||||
with pytest.raises(HTTPException):
|
||||
verifier.authenticate(self._headers(token))
|
||||
assert fetcher.calls == 2
|
||||
|
||||
def test_cold_cache_fetch_is_rate_limited_during_outage(self, jwt_config, fetcher, keypair):
|
||||
# Prefetch fails and the issuer stays down: a flood of requests must NOT each
|
||||
# trigger its own (blocking) fetch — the cooldown floors bound it to a few.
|
||||
fetcher.fail = True
|
||||
v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher)
|
||||
v.prefetch() # fails; cache stays empty
|
||||
headers = self._headers(_sign(keypair, {"email": "a@b.c"}))
|
||||
|
||||
for _ in range(10):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
v.authenticate(headers)
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
# 1 prefetch + at most one fetch per path before both cooldown floors engage.
|
||||
assert fetcher.calls <= 4
|
||||
|
||||
def test_jwks_down_with_warm_cache_still_verifies(self, verifier, fetcher, keypair):
|
||||
assert fetcher.calls == 1
|
||||
fetcher.fail = True # endpoint goes down
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
# Known kid is already cached → verifies without re-fetching.
|
||||
assert verifier.authenticate(self._headers(token)) == "alice@example.com"
|
||||
assert fetcher.calls == 1
|
||||
|
||||
def test_custom_jwt_header_with_bearer_prefix(self, fetcher, keypair):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://accounts.example.com/jwks",
|
||||
claim="email",
|
||||
jwt_header="X-Auth-Request-Access-Token",
|
||||
)
|
||||
v = IdentityVerifier(cfg, jwks_fetcher=fetcher)
|
||||
v.prefetch()
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
assert v.authenticate({"X-Auth-Request-Access-Token": f"Bearer {token}"}) == "alice@example.com"
|
||||
# Raw token without Bearer prefix also accepted.
|
||||
assert v.authenticate({"X-Auth-Request-Access-Token": token}) == "alice@example.com"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# IdentityVerifier — ES256 (the second accepted algorithm)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIdentityVerifierEs256:
|
||||
"""ES256 (ECDSA P-256) is accepted alongside RS256; keys are matched by ``kid``."""
|
||||
|
||||
def _verifier(self, jwks_doc: dict) -> IdentityVerifier:
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://accounts.example.com/jwks",
|
||||
claim="email",
|
||||
)
|
||||
v = IdentityVerifier(cfg, jwks_fetcher=lambda _u: jwks_doc)
|
||||
v.prefetch()
|
||||
return v
|
||||
|
||||
def test_valid_es256_token_returns_identity_claim(self):
|
||||
ec_key = _make_ec_keypair()
|
||||
v = self._verifier(_ec_jwks_for(ec_key.public_key()))
|
||||
token = _sign(ec_key, {"email": "ec-alice@example.com"}, kid=EC_KID, alg="ES256")
|
||||
assert v.authenticate({"Authorization": f"Bearer {token}"}) == "ec-alice@example.com"
|
||||
|
||||
def test_es256_bad_signature_rejected(self):
|
||||
ec_key = _make_ec_keypair()
|
||||
other = _make_ec_keypair() # advertised key differs from the signer
|
||||
v = self._verifier(_ec_jwks_for(ec_key.public_key()))
|
||||
token = _sign(other, {"email": "mallory@example.com"}, kid=EC_KID, alg="ES256")
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
v.authenticate({"Authorization": f"Bearer {token}"})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_es256_wrong_audience_rejected(self):
|
||||
ec_key = _make_ec_keypair()
|
||||
v = self._verifier(_ec_jwks_for(ec_key.public_key()))
|
||||
token = jwt.encode(
|
||||
{"iss": ISSUER, "aud": "some-other-client", "exp": int(time.time()) + 3600, "email": "a@b.c"},
|
||||
ec_key,
|
||||
algorithm="ES256",
|
||||
headers={"kid": EC_KID},
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
v.authenticate({"Authorization": f"Bearer {token}"})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_mixed_jwks_verifies_both_rs256_and_es256_by_kid(self):
|
||||
# A single JWKS carrying one RSA and one EC key; the verifier picks by kid.
|
||||
rsa_key = _make_keypair()
|
||||
ec_key = _make_ec_keypair()
|
||||
jwks_doc = {"keys": _jwks_for(rsa_key.public_key())["keys"] + _ec_jwks_for(ec_key.public_key())["keys"]}
|
||||
v = self._verifier(jwks_doc)
|
||||
rs_token = _sign(rsa_key, {"email": "rs@example.com"}) # RS256, kid=test-key-1
|
||||
es_token = _sign(ec_key, {"email": "es@example.com"}, kid=EC_KID, alg="ES256")
|
||||
assert v.authenticate({"Authorization": f"Bearer {rs_token}"}) == "rs@example.com"
|
||||
assert v.authenticate({"Authorization": f"Bearer {es_token}"}) == "es@example.com"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Startup prefetch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestJwksSourceResolution:
|
||||
"""OIDC discovery (when no explicit jwks_url) and the SSRF scheme guard."""
|
||||
|
||||
def _disco_config(self):
|
||||
# No jwt_jwks_url → the verifier must discover it from the issuer.
|
||||
return IdentityConfig(mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, claim="email")
|
||||
|
||||
def test_oidc_discovery_resolves_jwks_uri_then_verifies(self, jwks, keypair):
|
||||
seen = {}
|
||||
|
||||
def openid_fetch(url):
|
||||
seen["openid"] = url
|
||||
return {"jwks_uri": "https://accounts.example.com/discovered-jwks"}
|
||||
|
||||
def jwks_fetch(url):
|
||||
seen["jwks"] = url
|
||||
return jwks
|
||||
|
||||
v = IdentityVerifier(self._disco_config(), jwks_fetcher=jwks_fetch, openid_fetcher=openid_fetch)
|
||||
v.prefetch()
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com"
|
||||
# Discovery hit the well-known doc, then fetched the URL it advertised.
|
||||
assert seen["openid"] == "https://accounts.example.com/.well-known/openid-configuration"
|
||||
assert seen["jwks"] == "https://accounts.example.com/discovered-jwks"
|
||||
|
||||
def test_oidc_discovery_without_jwks_uri_degrades_to_401(self, keypair):
|
||||
# Discovery doc missing jwks_uri → IdentityConfigError, swallowed at fetch,
|
||||
# cache stays empty, prefetch does not raise, requests fail closed (401).
|
||||
v = IdentityVerifier(self._disco_config(), jwks_fetcher=lambda _u: {"keys": []}, openid_fetcher=lambda _u: {})
|
||||
v.prefetch()
|
||||
token = _sign(keypair, {"email": "a@b.c"})
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
v.authenticate({"Authorization": f"Bearer {token}"})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_fetch_json_rejects_non_http_scheme(self):
|
||||
# The SSRF guard: a file:// (or other non-http) URL must be refused.
|
||||
with pytest.raises(IdentityConfigError):
|
||||
_fetch_json("file:///etc/passwd")
|
||||
|
||||
def test_fetch_json_refuses_redirect(self):
|
||||
# A redirect (e.g. https -> http) must be refused, not followed: following it
|
||||
# would issue the plaintext request before any scheme check could run.
|
||||
redirect = httpx.Response(status_code=302, headers={"location": "http://evil.example/jwks"})
|
||||
with (
|
||||
patch.object(httpx.Client, "get", return_value=redirect),
|
||||
pytest.raises(IdentityConfigError, match="redirect"),
|
||||
):
|
||||
_fetch_json("https://accounts.example.com/jwks")
|
||||
|
||||
def _assert_fails_closed(self, verifier, keypair):
|
||||
"""A malformed external shape must not raise out of prefetch/auth — 401."""
|
||||
assert verifier.prefetch() is False # empty cache, no AttributeError/500
|
||||
token = _sign(keypair, {"email": "a@b.c"})
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
verifier.authenticate({"Authorization": f"Bearer {token}"})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_jwks_non_object_fails_closed(self, keypair):
|
||||
# A JWKS endpoint returning a JSON array (not an object) must fail closed,
|
||||
# not raise AttributeError out of PyJWKSet.from_dict.
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt", jwt_issuer=ISSUER, jwt_audience=AUDIENCE, jwt_jwks_url="https://x/jwks", claim="email"
|
||||
)
|
||||
v = IdentityVerifier(cfg, jwks_fetcher=lambda _u: [])
|
||||
self._assert_fails_closed(v, keypair)
|
||||
|
||||
def test_oidc_discovery_non_object_fails_closed(self, keypair):
|
||||
# Discovery doc that is a JSON array (no .get) must fail closed, not raise.
|
||||
v = IdentityVerifier(self._disco_config(), jwks_fetcher=lambda _u: {"keys": []}, openid_fetcher=lambda _u: [])
|
||||
self._assert_fails_closed(v, keypair)
|
||||
|
||||
def test_oidc_discovery_non_string_jwks_uri_fails_closed(self, keypair):
|
||||
# jwks_uri present but not a string must fail closed, not slip through to
|
||||
# a later AttributeError when the non-string is fetched.
|
||||
v = IdentityVerifier(
|
||||
self._disco_config(),
|
||||
jwks_fetcher=lambda _u: {"keys": []},
|
||||
openid_fetcher=lambda _u: {"jwks_uri": 123},
|
||||
)
|
||||
self._assert_fails_closed(v, keypair)
|
||||
|
||||
def test_oidc_discovery_http_jwks_uri_rejected_by_default(self, keypair):
|
||||
# A discovered jwks_uri over plaintext http:// is refused under the default
|
||||
# HTTPS policy (defense-in-depth even when the issuer itself was https).
|
||||
v = IdentityVerifier(
|
||||
self._disco_config(),
|
||||
jwks_fetcher=lambda _u: {"keys": []},
|
||||
openid_fetcher=lambda _u: {"jwks_uri": "http://accounts.example.com/jwks"},
|
||||
)
|
||||
self._assert_fails_closed(v, keypair)
|
||||
|
||||
|
||||
class TestStartupPrefetch:
|
||||
def test_prefetch_fetches_before_first_request(self, jwt_config, fetcher, keypair):
|
||||
v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher)
|
||||
assert fetcher.calls == 0
|
||||
assert v.prefetch() is True # warm cache reported so callers can stay quiet
|
||||
assert fetcher.calls == 1 # JWKS fetched at startup, not on first request
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com"
|
||||
assert fetcher.calls == 1 # first request paid no fetch round-trip
|
||||
|
||||
def test_failed_prefetch_does_not_raise_and_recovers(self, jwt_config, fetcher, keypair):
|
||||
fetcher.fail = True
|
||||
v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher)
|
||||
assert v.prefetch() is False # empty cache reported so the caller can warn loudly
|
||||
assert fetcher.calls == 1
|
||||
|
||||
fetcher.fail = False # issuer recovers
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
# First real request re-fetches (cache was empty) and now succeeds.
|
||||
assert v.authenticate({"Authorization": f"Bearer {token}"}) == "alice@example.com"
|
||||
assert fetcher.calls == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Header mode
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIdentityVerifierHeader:
|
||||
def test_reads_configured_header(self):
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
v = IdentityVerifier(cfg)
|
||||
assert v.authenticate({"X-Consumer-Username": "kong-user-7"}) == "kong-user-7"
|
||||
|
||||
def test_missing_header_rejected(self):
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
v = IdentityVerifier(cfg)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
v.authenticate({})
|
||||
assert exc.value.status_code == 401
|
||||
|
||||
def test_header_mode_emits_startup_warning(self):
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
registry = FlowRegistry()
|
||||
with patch("lfx.cli.serve_app.logger") as mock_logger:
|
||||
create_multi_serve_app(registry=registry, identity_config=cfg)
|
||||
assert mock_logger.warning.called
|
||||
warned = " ".join(str(c.args[0]) for c in mock_logger.warning.call_args_list)
|
||||
assert "X-Consumer-Username" in warned
|
||||
assert "topology" in warned.lower()
|
||||
|
||||
def test_header_mode_warns_on_guaranteed_console_channel(self):
|
||||
# The structlog logger is not reliably surfaced on the serve stdout path, so the
|
||||
# header trust warning must also reach the operator via the stderr console.
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
registry = FlowRegistry()
|
||||
with patch("lfx.cli.serve_app._startup_console") as mock_console:
|
||||
create_multi_serve_app(registry=registry, identity_config=cfg)
|
||||
assert mock_console.print.called
|
||||
printed = " ".join(str(c.args[0]) for c in mock_console.print.call_args_list)
|
||||
assert "X-Consumer-Username" in printed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Startup notices on the guaranteed (stderr console) channel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestStartupConsoleNotices:
|
||||
def _jwt_config(self):
|
||||
return IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://accounts.example.com/jwks.json",
|
||||
claim="email",
|
||||
)
|
||||
|
||||
def test_failed_jwt_prefetch_warns_on_console(self):
|
||||
# A failed prefetch must not abort startup, but it must warn loudly on the
|
||||
# guaranteed channel so the operator knows every request will 401.
|
||||
stub = MagicMock()
|
||||
stub.prefetch.return_value = False
|
||||
with (
|
||||
patch("lfx.cli.serve_app.build_identity_verifier", return_value=stub),
|
||||
patch("lfx.cli.serve_app._startup_console") as mock_console,
|
||||
):
|
||||
create_multi_serve_app(registry=FlowRegistry(), identity_config=self._jwt_config())
|
||||
assert stub.prefetch.called
|
||||
printed = " ".join(str(c.args[0]) for c in mock_console.print.call_args_list)
|
||||
assert "JWKS prefetch failed" in printed
|
||||
assert "401" in printed
|
||||
|
||||
def test_successful_jwt_prefetch_is_quiet_on_console(self):
|
||||
stub = MagicMock()
|
||||
stub.prefetch.return_value = True
|
||||
with (
|
||||
patch("lfx.cli.serve_app.build_identity_verifier", return_value=stub),
|
||||
patch("lfx.cli.serve_app._startup_console") as mock_console,
|
||||
):
|
||||
create_multi_serve_app(registry=FlowRegistry(), identity_config=self._jwt_config())
|
||||
assert stub.prefetch.called
|
||||
assert not mock_console.print.called
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Token never logged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTokenNeverLogged:
|
||||
def test_token_absent_from_all_log_calls(self, jwt_config, fetcher, keypair):
|
||||
v = IdentityVerifier(jwt_config, jwks_fetcher=fetcher)
|
||||
v.prefetch()
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
|
||||
with patch("lfx.cli.serve_identity.logger") as mock_logger:
|
||||
v.authenticate(headers) # success path
|
||||
with pytest.raises(HTTPException):
|
||||
v.authenticate({"Authorization": "Bearer not.a.jwt"}) # malformed
|
||||
fetcher.fail = True
|
||||
bad_kid = _sign(keypair, {"email": "a@b.c"}, kid="rotated")
|
||||
with pytest.raises(HTTPException):
|
||||
v.authenticate({"Authorization": f"Bearer {bad_kid}"}) # forces failing refresh + logging
|
||||
|
||||
logged = " ".join(str(a) for call in mock_logger.mock_calls if call.args for a in call.args)
|
||||
assert token not in logged
|
||||
assert "not.a.jwt" not in logged
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Endpoint integration (run/stream + serve-key floor)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
API_KEY = "test-api-key" # pragma: allowlist secret
|
||||
FLOW_ID = "00000000-0000-0000-0000-000000000001"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def real_graph():
|
||||
data_dir = Path(__file__).parent.parent.parent / "data"
|
||||
with (data_dir / "simple_chat_no_llm.json").open() as f:
|
||||
payload = json.load(f)
|
||||
return Graph.from_payload(payload, flow_id=FLOW_ID)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _allow_custom_components(monkeypatch):
|
||||
from lfx.services.deps import get_settings_service
|
||||
|
||||
monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True)
|
||||
|
||||
|
||||
def _client_with_verifier(real_graph, verifier_obj) -> TestClient:
|
||||
"""Build a serve app (off mode) and swap in a test verifier on app.state."""
|
||||
registry = FlowRegistry()
|
||||
registry.add(real_graph, FlowMeta(id=FLOW_ID, relative_path="t.json", title="t", description=None))
|
||||
app = create_multi_serve_app(registry=registry)
|
||||
app.state.identity_verifier = verifier_obj # None ⇒ off
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
class TestServeCli:
|
||||
"""The identity options are reachable through the real ``lfx serve`` CLI."""
|
||||
|
||||
def test_serve_help_lists_identity_options(self):
|
||||
from lfx.__main__ import app
|
||||
from typer.testing import CliRunner
|
||||
|
||||
# Force a wide terminal so rich does not wrap long option names across lines.
|
||||
with patch.dict(os.environ, {"COLUMNS": "200"}):
|
||||
result = CliRunner().invoke(app, ["serve", "--help"])
|
||||
assert result.exit_code == 0
|
||||
# Strip ANSI styling: rich colorizes each hyphen-segment of an option name
|
||||
# (``-``/``-identity``/``-mode``), so the raw output never holds the literal
|
||||
# ``--identity-mode`` substring when color is forced (as it is in CI).
|
||||
plain = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
|
||||
assert "--identity-mode" in plain
|
||||
assert "--identity-jwt-issuer" in plain
|
||||
|
||||
def test_jwt_mode_without_issuer_fails_fast(self):
|
||||
from lfx.__main__ import app
|
||||
from typer.testing import CliRunner
|
||||
|
||||
env = {k: v for k, v in os.environ.items() if not k.startswith("LFX_SERVE_")}
|
||||
env["LANGFLOW_API_KEY"] = API_KEY
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
result = CliRunner().invoke(app, ["serve", "--identity-mode", "jwt"])
|
||||
assert result.exit_code != 0
|
||||
assert "identity-jwt-issuer" in result.output
|
||||
|
||||
|
||||
class TestWorkerEnvRoundTrip:
|
||||
"""Identity config must survive the env-var round-trip into uvicorn workers."""
|
||||
|
||||
def _clean_serve_env(self) -> dict[str, str]:
|
||||
return {k: v for k, v in os.environ.items() if not k.startswith("LFX_SERVE_")}
|
||||
|
||||
def test_header_mode_round_trips_through_create_serve_app(self):
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
env = {**self._clean_serve_env(), "LANGFLOW_API_KEY": API_KEY, **cfg.to_env()}
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
app = create_serve_app()
|
||||
assert app.state.identity_config == cfg
|
||||
assert app.state.identity_verifier is not None
|
||||
|
||||
def test_jwt_mode_round_trips_through_create_serve_app(self, jwks):
|
||||
cfg = IdentityConfig(
|
||||
mode="jwt",
|
||||
jwt_issuer=ISSUER,
|
||||
jwt_audience=AUDIENCE,
|
||||
jwt_jwks_url="https://accounts.example.com/jwks",
|
||||
claim="email",
|
||||
jwt_header="X-Auth-Request-Access-Token",
|
||||
)
|
||||
env = {**self._clean_serve_env(), "LANGFLOW_API_KEY": API_KEY, **cfg.to_env()}
|
||||
# Patch the network fetcher so the worker's startup prefetch stays offline.
|
||||
with (
|
||||
patch.dict(os.environ, env, clear=True),
|
||||
patch("lfx.cli.serve_identity._fetch_json", return_value=jwks),
|
||||
):
|
||||
app = create_serve_app()
|
||||
# Every field reconstructed identically in the worker.
|
||||
assert app.state.identity_config == cfg
|
||||
assert app.state.identity_verifier is not None
|
||||
|
||||
|
||||
class TestServeAppIdentityEndpoints:
|
||||
def test_off_mode_requires_no_identity_and_threads_none(self, real_graph):
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["user_id"] = user_id
|
||||
return [], ""
|
||||
|
||||
client = _client_with_verifier(real_graph, None)
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute),
|
||||
):
|
||||
resp = client.post(f"/flows/{FLOW_ID}/run", json={"input_value": "hi"}, headers={"x-api-key": API_KEY})
|
||||
assert resp.status_code != 401
|
||||
assert captured["user_id"] is None # off mode ⇒ no identity threaded
|
||||
|
||||
def test_jwt_valid_token_threads_identity_into_run(self, real_graph, verifier, keypair):
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["user_id"] = user_id
|
||||
return [], ""
|
||||
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
client = _client_with_verifier(real_graph, verifier)
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute),
|
||||
):
|
||||
resp = client.post(
|
||||
f"/flows/{FLOW_ID}/run",
|
||||
json={"input_value": "hi"},
|
||||
headers={"x-api-key": API_KEY, "Authorization": f"Bearer {token}"},
|
||||
)
|
||||
assert resp.status_code != 401
|
||||
assert captured["user_id"] == "alice@example.com"
|
||||
|
||||
def test_jwt_valid_token_threads_identity_into_stream(self, real_graph, verifier, keypair):
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["user_id"] = user_id
|
||||
return [], ""
|
||||
|
||||
token = _sign(keypair, {"email": "bob@example.com"})
|
||||
client = _client_with_verifier(real_graph, verifier)
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute),
|
||||
client.stream(
|
||||
"POST",
|
||||
f"/flows/{FLOW_ID}/stream",
|
||||
json={"input_value": "hi"},
|
||||
headers={"x-api-key": API_KEY, "Authorization": f"Bearer {token}"},
|
||||
) as resp,
|
||||
):
|
||||
assert resp.status_code == 200
|
||||
for _ in resp.iter_bytes():
|
||||
pass
|
||||
assert captured["user_id"] == "bob@example.com"
|
||||
|
||||
def test_jwt_invalid_token_rejected_before_execution(self, real_graph, verifier):
|
||||
execute = MagicMock()
|
||||
client = _client_with_verifier(real_graph, verifier)
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", execute),
|
||||
):
|
||||
resp = client.post(
|
||||
f"/flows/{FLOW_ID}/run",
|
||||
json={"input_value": "hi"},
|
||||
headers={"x-api-key": API_KEY, "Authorization": "Bearer garbage.token.here"},
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
execute.assert_not_called()
|
||||
|
||||
def test_serve_key_floor_enforced_even_with_valid_token(self, real_graph, verifier, keypair):
|
||||
"""A valid identity token must NOT bypass the serve API key."""
|
||||
execute = MagicMock()
|
||||
token = _sign(keypair, {"email": "alice@example.com"})
|
||||
client = _client_with_verifier(real_graph, verifier)
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", execute),
|
||||
):
|
||||
resp = client.post(
|
||||
f"/flows/{FLOW_ID}/run",
|
||||
json={"input_value": "hi"},
|
||||
headers={"Authorization": f"Bearer {token}"}, # no x-api-key
|
||||
)
|
||||
assert resp.status_code == 401
|
||||
execute.assert_not_called()
|
||||
|
||||
def test_header_mode_threads_identity(self, real_graph):
|
||||
captured: dict = {}
|
||||
|
||||
async def mock_execute(graph, input_value, session_id=None, user_id=None): # noqa: ARG001
|
||||
captured["user_id"] = user_id
|
||||
return [], ""
|
||||
|
||||
cfg = IdentityConfig(mode="header", trusted_header="X-Consumer-Username")
|
||||
client = _client_with_verifier(real_graph, build_identity_verifier(cfg))
|
||||
with (
|
||||
patch.dict(os.environ, {"LANGFLOW_API_KEY": API_KEY}),
|
||||
patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute),
|
||||
):
|
||||
resp = client.post(
|
||||
f"/flows/{FLOW_ID}/run",
|
||||
json={"input_value": "hi"},
|
||||
headers={"x-api-key": API_KEY, "X-Consumer-Username": "kong-user-7"},
|
||||
)
|
||||
assert resp.status_code != 401
|
||||
assert captured["user_id"] == "kong-user-7"
|
||||
|
||||
|
||||
class TestExecuteGraphUserIdThreading:
|
||||
"""The user_id contract in execute_graph_with_capture.
|
||||
|
||||
Covers the off-mode preservation the docstrings describe: None keeps the
|
||||
graph's existing user_id, a verified value overwrites it. async_start is
|
||||
stubbed so only the apply_run_defaults stamping logic runs, not a full flow.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _stub_async_start(real_graph):
|
||||
async def _no_results(*_args, **_kwargs):
|
||||
return
|
||||
yield # make it an (empty) async generator
|
||||
|
||||
real_graph.async_start = _no_results
|
||||
|
||||
async def test_none_user_id_preserves_existing_graph_user_id(self, real_graph):
|
||||
self._stub_async_start(real_graph)
|
||||
real_graph.user_id = "preexisting-user"
|
||||
await execute_graph_with_capture(real_graph, "hi", user_id=None)
|
||||
assert real_graph.user_id == "preexisting-user"
|
||||
|
||||
async def test_verified_user_id_overwrites_graph_user_id(self, real_graph):
|
||||
self._stub_async_start(real_graph)
|
||||
real_graph.user_id = "preexisting-user"
|
||||
await execute_graph_with_capture(real_graph, "hi", user_id="alice@example.com")
|
||||
assert real_graph.user_id == "alice@example.com"
|
||||
2
uv.lock
generated
2
uv.lock
generated
@ -8951,6 +8951,7 @@ dependencies = [
|
||||
{ name = "platformdirs" },
|
||||
{ name = "pydantic" },
|
||||
{ name = "pydantic-settings" },
|
||||
{ name = "pyjwt" },
|
||||
{ name = "pypdf" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "pyyaml" },
|
||||
@ -9021,6 +9022,7 @@ requires-dist = [
|
||||
{ name = "platformdirs", specifier = ">=4.3.8,<5.0.0" },
|
||||
{ name = "pydantic", specifier = ">=2.0.0,<3.0.0" },
|
||||
{ name = "pydantic-settings", specifier = ">=2.10.1,<3.0.0" },
|
||||
{ name = "pyjwt", specifier = ">=2.10.0,<3.0.0" },
|
||||
{ name = "pypdf", specifier = ">=6.10.0,<7.0.0" },
|
||||
{ name = "python-dotenv", specifier = ">=1.0.0,<2.0.0" },
|
||||
{ name = "pyyaml", specifier = ">=6.0.0,<7.0.0" },
|
||||
|
||||
Reference in New Issue
Block a user