From 71cd5ce37632e37e9dfcd538a9b2fe8c0912a613 Mon Sep 17 00:00:00 2001 From: Jordan Frazier <122494242+jordanrfrazier@users.noreply.github.com> Date: Wed, 3 Jun 2026 11:33:39 -0400 Subject: [PATCH] feat: improve lfx serve production-readiness (#13121) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: add design spec for lfx serve multi-flow startup and dynamic upload * docs: add implementation plan for lfx serve multi-flow startup and dynamic upload * feat(serve): add FlowRegistry for mutable in-process flow management * feat(serve): add flow_id_from_content for content-based flow ID generation * fix(test): update streaming tests to use FlowRegistry API * test(serve): add upload endpoint tests * feat(serve): add build_registry_from_directory and build_registry_from_paths helpers * refactor(serve): remove dead _analyze_graph_structure and _generate_dynamic_run_description helpers * feat(serve): accept directory and multiple file paths at startup; wire registry to serve_command * [autofix.ci] apply automated fixes * fix(serve): move traceback import to module level; remove planning docs from branch * [autofix.ci] apply automated fixes * fix(serve): deepcopy graph in stream_flow; use json.dumps for error stream payload * test(serve): verify stream setup error payload is valid JSON with special characters * [autofix.ci] apply automated fixes * fix(serve): use load_flow_from_json directly in _load_graph_and_meta for clear error messages * docs(serve): document multi-worker and persistence limitations on FlowRegistry * feat(serve): allow starting server with no flows; accept full Langflow flow JSON on upload - Remove the requirement for a flow path/JSON/stdin at startup; lfx serve now starts with an empty FlowRegistry when no input is provided, ready to accept flows via POST /flows/upload/ - UploadFlowRequest now accepts the full Langflow export JSON directly (extra="allow", id field added) so callers can do -d @flow.json without wrapping; load_flow_from_json receives the full dict and the flow's own stable id is used when present * fix(serve): harden upload endpoint and remove unused verbose_print param - Pass body.model_dump(exclude={"replace"}) to load_flow_from_json instead of body.model_dump(); keeps the top-level "data" key the function expects while excluding the internal replace flag - Set graph.flow_id after upload so it matches the registry key - Return 409 Conflict on duplicate flow ID; add replace=true field to allow explicit overwrite - Remove unused verbose_print param from _load_graph_and_meta and update callers - Clarify directory scan is top-level only (non-recursive) in CLI help text - Add test_upload_full_export_as_body regression test: sends a complete Langflow export JSON as the request body (curl -d @flow.json pattern) where body.data has no nested "data" key — would have caught the KeyError('data') regression * fix(serve): address PR review comments - Restore .py script support: _load_graph_and_meta now dispatches to load_graph_from_script for .py paths; CLI validation accepts .json and .py; help text updated in both _running_commands.py and commands.py - Fix flow ID collision in build_registry_from_paths: compute a shared common root from all supplied paths so same-named files in different directories get distinct UUIDv5 IDs instead of silently overwriting each other - Harden FlowRegistry.add(): add overwrite=False parameter that raises ValueError on duplicate IDs; startup paths keep default (collision = error), upload endpoint passes overwrite=body.replace - Add regression test: test_same_filename_in_different_dirs_gets_distinct_ids - Update test_duplicate_add_replaces_graph to reflect new overwrite semantics * test(serve): add unit tests for .py script path in serve command - test_load_graph_and_meta_dispatches_to_script_loader_for_py: verifies _load_graph_and_meta calls load_graph_from_script (not load_flow_from_json) for .py paths, and that meta.title and relative_path are set correctly - test_serve_command_accepts_py_file: end-to-end check that lfx serve my.py reaches the server (exit_code == 0) - test_serve_command_rejects_unsupported_extension: .txt and other unknown extensions must exit non-zero with a helpful message * [autofix.ci] apply automated fixes * Add --no-env-fallback option to support scoped request_vars * [autofix.ci] apply automated fixes * feat(serve): filesystem-backed flow registry for multi-worker sharing Adds a FlowStore abstraction with two implementations: - NullFlowStore (default): in-memory only, existing behavior unchanged - FilesystemFlowStore: persists each flow as {id}.json; writes are atomic (tmp→rename); reads and deletes are TOCTOU-safe; path traversal guard rejects unsafe flow IDs FlowRegistry now delegates persistence to FlowStore: - add(raw_json=...) writes to store when raw JSON is available - get() falls through to store on cache miss, reconstructs + caches - list_metas() merges in-memory and store-side entries - remove() propagates deletion to store - warm_from_store() pre-warms the in-memory cache at startup New --flow-dir option on lfx serve activates the filesystem store. All uvicorn workers pointing at the same directory serve the same flows — use /tmp/lfx-flows for single-pod sharing or a PVC mount for cross-pod sharing. Raw JSON is plumbed through _load_graph_and_meta and both builder functions so startup-loaded .json flows are persisted automatically; .py flows remain in-memory only. * feat(serve): add --workers option for multi-worker uvicorn support Adds --workers N to lfx serve, enabling multiple uvicorn processes to share flows via --flow-dir. serve_command is restructured from @syncify async to plain sync: all async registry-building is extracted into _build_serve_registry() and called via asyncio.run(), then uvicorn.run() handles the server. For workers > 1, uvicorn requires an import string rather than an app object. The factory create_serve_app() in serve_app.py reads LFX_SERVE_FLOW_DIR and LFX_SERVE_NO_ENV_FALLBACK from env, pre-warms from the shared FilesystemFlowStore, and returns a ready FastAPI app. The parent sets those env vars unconditionally before dispatching. For workers == 1, the pre-built app object is used directly (no change to existing single-worker behavior). A warning is emitted when --workers > 1 is used without --flow-dir since each worker would have an isolated in-memory registry. The startup panel flow count is now accurate for single-worker + --flow-dir by calling warm_from_store() before printing. * [autofix.ci] apply automated fixes * [autofix.ci] apply automated fixes (attempt 2/3) * [autofix.ci] apply automated fixes (attempt 3/3) * fix(serve): remove no_args_is_help, drop verbose_print param, harden _safe_path - Remove no_args_is_help=True from serve command so lfx serve with no arguments starts with an empty registry instead of showing help - Remove unused verbose_print parameter from create_multi_serve_app and update all callers (14 call sites across tests) - Wrap _safe_path's relative_to() check in try/except for a clear error message; restructure delete() to use try/except/else (TRY300) - Add TYPE_CHECKING guard for FlowStore annotation in commands.py; wrap long --workers help string; restore FBT003 noqa on no_env_fallback - Fix test_env_fallback_skipped_when_flag_true to allow logger's getenv calls while still asserting the credential variable is never looked up - Fix ARG002/D205/RUF059 lint issues in test files * [autofix.ci] apply automated fixes * fix(serve): resolve merge conflicts and clean up test section comments * ruff - test updates * fix(serve): cross-worker delete propagation and HTTP 500 for execution failures Cross-worker stale cache (multi-worker DELETE bug): - Add FlowStore.is_persistent to distinguish real stores from NullFlowStore - Track _store_sourced in FlowRegistry for flows written to a persistent store - get() re-reads the store for store-sourced flows to detect cross-worker deletes; evicts and returns None if the file is gone - list_metas() skips store-sourced flows whose store key is no longer present - Add _evict() helper to clean up _flows, _store_meta_cache, _store_keys, _store_sourced HTTP status for execution failures: - POST /flows/{id}/run now returns HTTP 500 (not 200) when execution raises or when the flow returns success=false; body is still RunResponse for debug detail - Updated OpenAPI responses declaration: 500 now maps to RunResponse model * fix(serve): harden multi-worker registry correctness and safety - Add FlowAlreadyRegisteredError (ValueError subclass) so upload_flow catches only the specific race-condition exception rather than any ValueError from unrelated internals. - Add _SERVE_ENV_PREFIX constant; rebuild LFX_SERVE_* env var names from it and add try/finally cleanup so those vars are always deleted after uvicorn.run() returns in multi-worker mode. - Reject .py startup files when --workers > 1: Python graphs cannot be serialized to the filesystem store so workers would silently start with empty registries for those flows. - Remove warm_from_store() from build_registry_from_paths/directory; the builders no longer pre-warm to avoid double-loading when the startup dir and flow-dir overlap. serve_command still warms for single-worker mode. - Fix FlowRegistry.__len__ to delegate to list_metas() so store-only flows (not yet cache-loaded) are counted; this makes the startup panel and /health flow_count accurate for multi-worker + pre-placed files. - Fix FlowRegistry.remove() to delete both the primary store key file and any alternate-keyed file (UUID or stem), ensuring cross-worker DELETE propagation: after a delete, every worker's next request to that flow triggers the per-request stale check, reads None from the store, and returns 404 regardless of which key its cache entry uses. - Document the per-request stale check in FlowRegistry docstring with latency guidance for network volumes vs. local tmpfs. - Add tests for all of the above, including a cross-worker DELETE propagation scenario validated by a live 2-worker integration test. * fix(serve): load startup flows in every uvicorn worker via LFX_SERVE_STARTUP_PATHS Previously, lfx serve flow.json --workers 4 silently dropped startup flows: the parent loaded them into a NullFlowStore registry that workers never inherited. Each worker called create_serve_app() fresh and started empty, returning flow_count: 0 from /health and 404 for the startup flow. Fix: before calling uvicorn.run() in multi-worker mode, set LFX_SERVE_STARTUP_PATHS to the JSON-encoded list of resolved startup file paths (or the inline-JSON temp file path). create_serve_app() reads this env var and re-loads the startup flows independently in each worker process using a ThreadPoolExecutor to avoid RuntimeError from asyncio.run() being called inside uvicorn's already-running event loop. Startup flows loaded this way live in each worker's in-memory registry. Uploaded flows still require --flow-dir for cross-worker sharing (existing warning preserved). The LFX_SERVE_* env vars are cleaned up after uvicorn.run() returns via the existing prefix-sweep finally block. Tests added: - test_serve_command_sets_startup_paths_env_for_multi_worker: verifies LFX_SERVE_STARTUP_PATHS is set before uvicorn.run() and cleaned up after - TestCreateServeAppFactory: unit tests for create_serve_app() in all three modes (empty start, with startup paths, empty paths list) * fix(serve): fix 3 worker/store consistency issues Issue 1 — redundant store writes with --flow-dir + startup files + workers > 1: When flow_dir is set the parent already persisted startup flows to the store. create_serve_app() was re-loading and re-writing them in every worker. Fix: create_serve_app() now skips LFX_SERVE_STARTUP_PATHS when flow_dir is set and relies solely on warm_from_store(). Similarly, serve_command now only populates LFX_SERVE_STARTUP_PATHS when flow_dir is NOT set. Issue 2 — .py restriction was too broad: The error blocked .py + --workers > 1 even without --flow-dir, where each worker can safely re-execute the .py via LFX_SERVE_STARTUP_PATHS. The error now only fires when --flow-dir is also set (workers would skip startup paths and warm_from_store, missing the .py flow entirely since it can't be stored). Issue 3 — parent built FastAPI app for multi-worker then discarded it: create_multi_serve_app() was called before the workers > 1 / else branch, building a full FastAPI app that was never served in multi-worker mode. Moved into the else branch so it's only built for single-worker. Tests: - test_create_serve_app_with_flow_dir_skips_startup_paths_uses_store: asserts build_registry_from_paths NOT called when flow_dir is set; flow comes from store via warm_from_store instead. - test_create_serve_app_without_flow_dir_loads_startup_paths_from_files: asserts flow IS loaded when no flow_dir. - test_serve_command_does_not_set_startup_paths_when_flow_dir_set: asserts LFX_SERVE_STARTUP_PATHS is [] when flow_dir is set. - test_serve_command_allows_py_with_multiple_workers_no_flow_dir: asserts .py + workers>1 + no flow_dir is allowed and path appears in env. * test(serve): remove duplicate test and fix mislabeled test name - Remove test_create_serve_app_loads_startup_paths: functionally identical to test_create_serve_app_without_flow_dir_loads_startup_paths_from_files (same setup, same assertions, added in a prior session). - Rename test_verify_api_key_header_takes_precedence to test_verify_api_key_query_param_takes_precedence: the test passes the correct key as the query_param argument and the wrong key as header_param, then asserts the query param value is returned — query param has precedence, not header. The old name was the opposite of what the test actually checks. * [autofix.ci] apply automated fixes * fix(serve): fix concurrent request isolation and startup UX Three bug fixes discovered via live QA: 1. Component.__deepcopy__ produced shallow copies of _outputs_map, _edges, and _components, causing concurrent requests to share the same Output objects across graph copies. Output.cache=True by default: the first request's output.value was served to all subsequent concurrent requests, producing identical responses for different inputs. Fixed by deep-copying all three fields with memo. 2. POST /flows/upload/ without an explicit id used uuid5(content), which is deterministic — two uploads of any flow with the same nodes/edges structure always generated the same UUID and the second always returned 409. Fixed by using uuid4() (random). 3. Startup CLI emitted emoji and exclamation-point decoration in the serve banner and verbose output. Replaced with plain text. * [autofix.ci] apply automated fixes * fix(serve): address PR review comments - Move 409 conflict check before load_flow_from_json+prepare so callers with an explicit id don't pay the parse cost for a flow that will be rejected anyway - Remove unreachable FlowAlreadyRegisteredError catch: no await exists between registry.get() and registry.add(), so asyncio cannot interleave another request in that window - Update FlowMeta.id and UploadFlowResponse.id field descriptions from stale "Deterministic UUIDv5" to plain "Flow identifier (UUID)" * test(serve): fix upload duplicate/replace tests after uuid4 change Both tests relied on uuid5(content) determinism to generate the same flow ID from identical data. With uuid4, each upload gets a fresh ID, so duplicate detection was never triggered. Tests now supply an explicit id field to exercise the 409/replace paths directly. * fix(serve): address code review suggestions - Fix test_write_is_atomic to glob for any *.tmp files rather than checking a hardcoded filename that could never exist - Remove incorrect noqa: B904 from validate_id_is_uuid (not in an except block); use `from None` instead - Add TTL cache (1 s) for FlowStore.list_ids() in FlowRegistry to avoid one filesystem glob per /health request; mutations invalidate immediately, warm_from_store and delete-scan bypass it - Document the UUID/stem dual-key alias invariant on _flows - Add test verifying create_serve_app() propagates load errors from the ThreadPoolExecutor startup path (fail fast, not silent empty registry) * [autofix.ci] apply automated fixes * fix: lfx serve no_env_fallback, deepcopy cycle, cross-worker delete (#13297) fix: address concurrency and isolation issues in lfx serve - Re-apply no_env_fallback to the per-request graph copy in the run and stream endpoints; deepcopy() drops graph.context so the flag was lost. - Register the copy in memo before recursive deepcopy calls in Component.__deepcopy__ so reference cycles are preserved. - Check the store directly in list_metas() so cross-worker deletes are reflected immediately instead of after the TTL cache window. - Remove the redundant upload-conflict test (already covered by test_upload_duplicate_without_replace_returns_409). Also mark test secrets with allowlist pragmas and fix two docstring lint warnings in the touched test file. * [autofix.ci] apply automated fixes * feat(lfx): WXO/TRM request-scoped variables for lfx serve (child of #13121) (#13327) * feat(lfx): WXO/TRM request-scoped variables for lfx serve global_vars Align lfx serve with the TRM/WXO credential contract on top of #13121: - Parse LANGFLOW_REQUEST_VARIABLES from global_vars payloads - ContextVar-backed lookup in VariableService (no os.environ mutation) - Activate request scope during execute_graph_with_capture Child of add-workers-and-upload-lfx (PR #13121). * [autofix.ci] apply automated fixes * fix(lfx): harden request-scoped variable resolution + add isolation t… (#13347) fix(lfx): harden request-scoped variable resolution + add isolation tests Addresses PR review of WXO/TRM request-scoped variables for lfx serve: - service: honor no_env_fallback in VariableService so model/KB credential resolution no longer falls back to os.environ under --no-env-fallback; correct the get_variable precedence docstring (5-step interleaved order) - request_scope: add a per-request no_env_fallback ContextVar; use precise Token type; document the None-vs-empty scope contract - common: activate/reset the no_env_fallback scope from graph.context - runtime_variables: log (instead of silently swallowing) a malformed or non-object LANGFLOW_REQUEST_VARIABLES blob; anchor the override comment to the TRM contract Tests: - add request-scope concurrency isolation + no-env-fallback suite - add execute_graph_with_capture wiring test (graph.context -> ContextVars) - give TestGraphExecution mock graphs a real context, fixing 11 tests that broke once the helper began reading graph.context * [autofix.ci] apply automated fixes --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Jordan Frazier <122494242+jordanrfrazier@users.noreply.github.com> * Update use of stem alias tracking to ensure proper deletes * fix(lfx): enforce no_env_fallback across all credential resolution paths Closes credential-isolation gaps in lfx serve and hardens request-scoped variable handling: - VariableService: skip the process-wide LANGFLOW_REQUEST_VARIABLES env read when env fallback is disabled (previously leaked process credentials for an empty-global_vars request despite --no-env-fallback). - Unified models: gate os.environ reads in get_api_key_for_provider and the connection-config fallbacks in instantiation.py (WatsonX URL/project-id, Ollama base-url, OpenRouter headers) on the no-env-fallback flag. - run_until_complete: propagate the caller's contextvars into the worker thread so request-scoped vars and the flag survive the thread hop. - Request variables: serialize structured (dict/list) values as valid JSON and drop nulls instead of producing lossy Python reprs or a truthy "None". - FlowRegistry: extract the duplicated store-alias scan shared by add/remove. Adds tests for request-scope isolation, credential/env gating, and contextvar propagation. * fix(lfx): close no_env_fallback credential leaks and harden serve isolation - credentials: honor no_env_fallback in get_all_variables_for_provider so served flows no longer leak process env (was defeating the _env_if_allowed guards in instantiation.py via pre-filled provider_vars) - loading: apply the no_env_fallback contract to table load_from_db column resolution (sibling of load_from_env_vars, previously still read os.getenv) - variable service: resolve request-scoped vars (exact name + x-langflow-global-var alias) before env fallback so a caller's per-request credential is never shadowed by a same-named process env var - runtime_variables: flatten the LANGFLOW_REQUEST_VARIABLES blob into graph context for parity between the ContextVar and load_from_db resolution paths - flow_store: clean up the partial temp file when an atomic write fails - commands: on multi-worker shutdown remove only the serve env vars we set, not the whole LFX_SERVE_* namespace - components: add check_component_env_writes lint banning os.environ writes; stop the mem0 component from writing the OpenAI key into the process environment * fix(lfx): contain corrupt flow-store files so one bad file can't crash worker startup FilesystemFlowStore.read() only caught FileNotFoundError, so a corrupted/partial {id}.json in a shared store (PVC corruption, an external writer killed mid-write) raised JSONDecodeError. warm_from_store() calls get() for every store id at startup with no per-file guard, so a single bad file would crash every worker's startup and 500 GET /flows. - read(): catch JSONDecodeError/OSError/UnicodeDecodeError, log, return None (treat as absent) so warm_from_store/list_metas/get skip it and other flows keep serving. - warm_from_store(): skip+log a flow that fails to reconstruct (corrupt JSON, or valid JSON referencing a component not available in this build) instead of aborting startup. - tests: corrupt-JSON read returns None; warm_from_store skips unloadable flows. * fix test -- request wins over env now * fix test. mock graph context correctly --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida Co-authored-by: Himavarsha <40851462+HimavarshaVS@users.noreply.github.com> --- .pre-commit-config.yaml | 6 + .secrets.baseline | 2 +- pyproject.toml | 1 + scripts/lint/check_component_env_writes.py | 226 +++ .../bundles/mem0/test_mem0_component.py | 17 +- .../unit/interface/initialize/test_loading.py | 8 + src/lfx/src/lfx/_assets/component_index.json | 6 +- .../base/models/unified_models/credentials.py | 24 +- .../models/unified_models/instantiation.py | 27 +- src/lfx/src/lfx/cli/_extension_commands.py | 2 +- src/lfx/src/lfx/cli/_running_commands.py | 45 +- src/lfx/src/lfx/cli/commands.py | 641 +++--- src/lfx/src/lfx/cli/common.py | 13 + src/lfx/src/lfx/cli/flow_store.py | 149 ++ src/lfx/src/lfx/cli/run.py | 2 +- src/lfx/src/lfx/cli/runtime_variables.py | 70 + src/lfx/src/lfx/cli/serve_app.py | 919 ++++++--- .../lfx/components/mem0/mem0_chat_memory.py | 21 +- .../lfx/custom/custom_component/component.py | 19 +- .../src/lfx/interface/initialize/loading.py | 30 +- .../lfx/services/variable/request_scope.py | 75 + src/lfx/src/lfx/services/variable/service.py | 103 +- src/lfx/src/lfx/utils/async_helpers.py | 12 +- .../unit/base/models/test_credentials.py | 216 ++ src/lfx/tests/unit/cli/test_common.py | 47 + src/lfx/tests/unit/cli/test_flow_store.py | 143 ++ .../tests/unit/cli/test_runtime_variables.py | 98 + src/lfx/tests/unit/cli/test_serve.py | 887 ++++++++- src/lfx/tests/unit/cli/test_serve_app.py | 1734 ++++++++++++++++- .../unit/cli/test_serve_app_streaming.py | 85 +- .../tests/unit/cli/test_serve_components.py | 229 +-- .../test_concurrent_tool_invocation.py | 22 + .../interface/test_loading_no_env_fallback.py | 59 + .../unit/services/test_minimal_services.py | 49 + .../services/test_request_scope_isolation.py | 216 ++ .../tests/unit/utils/test_async_helpers.py | 44 + 36 files changed, 5242 insertions(+), 1005 deletions(-) create mode 100755 scripts/lint/check_component_env_writes.py create mode 100644 src/lfx/src/lfx/cli/flow_store.py create mode 100644 src/lfx/src/lfx/cli/runtime_variables.py create mode 100644 src/lfx/src/lfx/services/variable/request_scope.py create mode 100644 src/lfx/tests/unit/base/models/test_credentials.py create mode 100644 src/lfx/tests/unit/cli/test_flow_store.py create mode 100644 src/lfx/tests/unit/cli/test_runtime_variables.py create mode 100644 src/lfx/tests/unit/interface/test_loading_no_env_fallback.py create mode 100644 src/lfx/tests/unit/services/test_request_scope_isolation.py create mode 100644 src/lfx/tests/unit/utils/test_async_helpers.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cece539c6a..c69d7070ae 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -66,6 +66,12 @@ repos: language: system files: ^src/backend/base/langflow/api/.*\.py$ pass_filenames: false + - id: component-env-writes + name: Components must not write to os.environ (process-global; bleeds across requests) + entry: python scripts/lint/check_component_env_writes.py + language: system + files: ^(src/lfx/src/lfx/components/.*\.py|src/backend/base/langflow/components/.*\.py|src/bundles/.*\.py)$ + pass_filenames: false - repo: https://github.com/Yelp/detect-secrets rev: v1.5.0 hooks: diff --git a/.secrets.baseline b/.secrets.baseline index 36edf82053..b12441041c 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -8889,7 +8889,7 @@ "filename": "src/lfx/src/lfx/cli/serve_app.py", "hashed_secret": "b894b81be94cf8fa8d7536475aaec876addf05c8", "is_verified": false, - "line_number": 40, + "line_number": 49, "is_secret": false } ], diff --git a/pyproject.toml b/pyproject.toml index ecfd66a2c6..4ca5cd21b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -250,6 +250,7 @@ flake8-bugbear.extend-immutable-calls = [ "fastapi.Depends", "fastapi.File", "fastapi.Query", + "typer.Argument", "typer.Option", ] flake8-type-checking.runtime-evaluated-base-classes = [ diff --git a/scripts/lint/check_component_env_writes.py b/scripts/lint/check_component_env_writes.py new file mode 100755 index 0000000000..23b8575896 --- /dev/null +++ b/scripts/lint/check_component_env_writes.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""CI guard: components must not WRITE to ``os.environ``. + +``os.environ`` is process-global. A warm, long-lived serving process (``lfx serve``'s +worker pool, or the main Langflow server) handles many callers' requests in one +process, so a per-request write like ``os.environ["OPENAI_API_KEY"] = caller_key`` +bleeds into other requests on that worker and persists until overwritten — leaking +one caller's credential to another. (The bleed is per-worker: uvicorn ``--workers N`` +are separate processes, but every concurrent request inside one worker shares its +``os.environ``.) + +READS are fine — this guard only flags WRITES. Pass per-request values through the +component's config/params instead of mutating the environment. + +This walks every component source file (stdlib AST only — zero dependency on the lfx +package, since CI may run before it is importable) and flags: + + os.environ[k] = v os.environ = {...} del os.environ[k] + os.environ.update(...) os.environ.setdefault(...) .pop/.clear/.popitem + os.putenv(...) os.unsetenv(...) load_dotenv(...) + +Usage:: + + python scripts/lint/check_component_env_writes.py # scan default roots + python scripts/lint/check_component_env_writes.py FILE... # check specific files + python scripts/lint/check_component_env_writes.py --root path/to/components + +Exit codes: + 0 -- no disallowed env writes + 1 -- one or more disallowed env writes (or a file failed to parse) +""" + +from __future__ import annotations + +import argparse +import ast +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] + +# Component source roots (matches scripts/migrate/check_bare_names.py coverage). +DEFAULT_ROOTS = ( + "src/lfx/src/lfx/components", + "src/backend/base/langflow/components", + "src/bundles", +) + +# Deliberate, reviewed exceptions: repo-relative path -> set of allowed violation kinds. +# Keep this list tiny and justified; an entry here means "this component is knowingly +# allowed to mutate the process environment". Prefer fixing the component over adding one. +ALLOWLIST: dict[str, set[str]] = { + # The "Dotenv" component's entire purpose is to load a user-supplied .env blob into + # the process environment. It is a known multi-tenant hazard flagged for separate + # review; allow only the load_dotenv call, still flagging any direct os.environ write. + "src/lfx/src/lfx/components/datastax/dotenv.py": {"load_dotenv"}, +} + +_MUTATING_METHODS = {"setdefault", "update", "pop", "clear", "popitem", "__setitem__", "__delitem__"} + + +class _EnvWriteVisitor(ast.NodeVisitor): + """Flags writes to ``os.environ`` (and ``load_dotenv()``), import-alias-aware. + + A first pass over the tree collects local names that refer to ``os``, + ``os.environ``, ``os.putenv``, ``os.unsetenv`` and ``dotenv.load_dotenv`` + (including ``as``-aliases), so ``import os as system; system.environ[X]=Y`` + and ``from dotenv import load_dotenv as ld; ld()`` are caught. + """ + + def __init__(self) -> None: + self.violations: list[tuple[int, str, str]] = [] # (lineno, kind, snippet) + self.os_aliases: set[str] = {"os"} + self.environ_aliases: set[str] = {"environ"} + self.putenv_aliases: set[str] = set() + self.unsetenv_aliases: set[str] = set() + self.load_dotenv_aliases: set[str] = {"load_dotenv"} + + # First pass: collect import aliases module-wide (over-broad on scope is fine for a + # security lint — we'd rather flag a write that imports `os` locally than miss it). + def collect_aliases(self, tree: ast.AST) -> None: + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name == "os": + self.os_aliases.add(alias.asname or alias.name) + elif isinstance(node, ast.ImportFrom): + if node.module == "os": + for alias in node.names: + local = alias.asname or alias.name + if alias.name == "environ": + self.environ_aliases.add(local) + elif alias.name == "putenv": + self.putenv_aliases.add(local) + elif alias.name == "unsetenv": + self.unsetenv_aliases.add(local) + elif node.module == "dotenv": + for alias in node.names: + if alias.name == "load_dotenv": + self.load_dotenv_aliases.add(alias.asname or alias.name) + + def _is_environ(self, node: ast.AST) -> bool: + """True for ``.environ`` or a bare ````.""" + if isinstance(node, ast.Attribute): + return node.attr == "environ" and isinstance(node.value, ast.Name) and node.value.id in self.os_aliases + return isinstance(node, ast.Name) and node.id in self.environ_aliases + + def _flag(self, lineno: int, kind: str, snippet: str) -> None: + self.violations.append((lineno, kind, snippet)) + + def visit_Assign(self, node: ast.Assign) -> None: + for target in node.targets: + if isinstance(target, ast.Subscript) and self._is_environ(target.value): + self._flag(node.lineno, "assign", "os.environ[...] = ...") + elif self._is_environ(target): + self._flag(node.lineno, "rebind", "os.environ = ...") + self.generic_visit(node) + + def visit_AnnAssign(self, node: ast.AnnAssign) -> None: + if isinstance(node.target, ast.Subscript) and self._is_environ(node.target.value): + self._flag(node.lineno, "assign", "os.environ[...]: ... = ...") + self.generic_visit(node) + + def visit_AugAssign(self, node: ast.AugAssign) -> None: + target = node.target + if (isinstance(target, ast.Subscript) and self._is_environ(target.value)) or self._is_environ(target): + self._flag(node.lineno, "assign", "os.environ augmented assignment") + self.generic_visit(node) + + def visit_Delete(self, node: ast.Delete) -> None: + for target in node.targets: + if isinstance(target, ast.Subscript) and self._is_environ(target.value): + self._flag(node.lineno, "del", "del os.environ[...]") + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: + func = node.func + if isinstance(func, ast.Attribute): + if func.attr in _MUTATING_METHODS and self._is_environ(func.value): + self._flag(node.lineno, f"mutate:{func.attr}", f"os.environ.{func.attr}(...)") + elif ( + func.attr in {"putenv", "unsetenv"} + and isinstance(func.value, ast.Name) + and func.value.id in self.os_aliases + ): + self._flag(node.lineno, func.attr, f"os.{func.attr}(...)") + elif func.attr == "load_dotenv": + self._flag(node.lineno, "load_dotenv", "load_dotenv(...)") + elif isinstance(func, ast.Name): + if func.id in self.putenv_aliases: + self._flag(node.lineno, "putenv", f"{func.id}(...)") + elif func.id in self.unsetenv_aliases: + self._flag(node.lineno, "unsetenv", f"{func.id}(...)") + elif func.id in self.load_dotenv_aliases: + self._flag(node.lineno, "load_dotenv", f"{func.id}(...)") + self.generic_visit(node) + + +def _iter_py_files(roots: list[Path]) -> list[Path]: + files: list[Path] = [] + for root in roots: + if not root.exists(): + continue + for path in sorted(root.rglob("*.py")): + parts = path.parts + if "tests" in parts or path.name.startswith("test_"): + continue + files.append(path) + return files + + +def _check_file(path: Path) -> list[tuple[str, int, str, str]]: + try: + rel = path.resolve().relative_to(REPO_ROOT).as_posix() + except ValueError: + rel = path.as_posix() + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except (OSError, SyntaxError) as exc: + return [(rel, getattr(exc, "lineno", 0) or 0, "parse-error", str(exc))] + visitor = _EnvWriteVisitor() + visitor.collect_aliases(tree) + visitor.visit(tree) + allowed = ALLOWLIST.get(rel, set()) + out: list[tuple[str, int, str, str]] = [] + for lineno, kind, snippet in visitor.violations: + if kind in allowed: + continue + out.append((rel, lineno, kind, snippet)) + return out + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("files", nargs="*", help="Specific files to check (default: scan component roots).") + parser.add_argument("--root", action="append", default=None, help="Override component root(s) to scan.") + args = parser.parse_args(argv) + + if args.files: + targets = [Path(f).resolve() for f in args.files if f.endswith(".py")] + else: + roots = [Path(r) if Path(r).is_absolute() else REPO_ROOT / r for r in (args.root or DEFAULT_ROOTS)] + targets = _iter_py_files(roots) + + violations: list[tuple[str, int, str, str]] = [] + for path in targets: + violations.extend(_check_file(path)) + + if not violations: + return 0 + + print("ERROR: component(s) write to os.environ (process-global — bleeds across requests") + print(" in a shared serving process such as `lfx serve`):\n") + for rel, lineno, kind, snippet in sorted(violations): + print(f" {rel}:{lineno} [{kind}] {snippet}") + print( + "\nFix: pass per-request values through the component's config/params, not the\n" + "process environment (reads are fine; only writes bleed). If a write is a\n" + "deliberate, reviewed exception, add it to ALLOWLIST in\n" + "scripts/lint/check_component_env_writes.py with a justification." + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/backend/tests/unit/components/bundles/mem0/test_mem0_component.py b/src/backend/tests/unit/components/bundles/mem0/test_mem0_component.py index 3736831eac..48256750a3 100644 --- a/src/backend/tests/unit/components/bundles/mem0/test_mem0_component.py +++ b/src/backend/tests/unit/components/bundles/mem0/test_mem0_component.py @@ -24,11 +24,24 @@ class TestMem0CloudValidation: @patch("lfx.components.mem0.mem0_chat_memory.Memory") def test_build_mem0_works_when_not_in_cloud(self, mock_memory): - """Test that build_mem0 works when ASTRA_CLOUD_DISABLE_COMPONENT is false.""" + """build_mem0 works when not in cloud, passing the OpenAI key through mem0 config. + + The key must NOT be written to os.environ: that is process-global and persists + across requests, so in a shared serving process (lfx serve) one caller's key + would leak into other concurrent requests. + """ with patch.dict(os.environ, {"ASTRA_CLOUD_DISABLE_COMPONENT": "false"}): + os.environ.pop("OPENAI_API_KEY", None) component = Mem0MemoryComponent(openai_api_key="test-key") component.build_mem0() - mock_memory.assert_called_once() + + # Key flows through config, not the environment. + mock_memory.from_config.assert_called_once() + config = mock_memory.from_config.call_args.kwargs["config_dict"] + assert config["llm"]["config"]["api_key"] == "test-key" + assert config["embedder"]["config"]["api_key"] == "test-key" + # Regression guard for the credential-bleed bug: never pollute os.environ. + assert "OPENAI_API_KEY" not in os.environ def test_build_search_results_uses_mem0_2_filters_for_search(self): """Test that search uses filters for Mem0 2.x entity scoping.""" diff --git a/src/backend/tests/unit/interface/initialize/test_loading.py b/src/backend/tests/unit/interface/initialize/test_loading.py index 126e64ca34..e2fba2886f 100644 --- a/src/backend/tests/unit/interface/initialize/test_loading.py +++ b/src/backend/tests/unit/interface/initialize/test_loading.py @@ -369,6 +369,10 @@ async def test_update_table_params_with_fallback_to_env(): # Create mock custom component custom_component = MagicMock() + # graph.context is a dict in production; set it explicitly so no_env_fallback resolves + # to False (a bare MagicMock's .context.get(...) returns a truthy mock and would + # silently disable env fallback). + custom_component.graph.context = {} custom_component.get_variable = AsyncMock(side_effect=ValueError("variable not found.")) # Set up table params @@ -406,6 +410,10 @@ async def test_update_table_params_mixed_db_and_env(): # Create mock custom component custom_component = MagicMock() + # graph.context is a dict in production; set it explicitly so no_env_fallback resolves + # to False (a bare MagicMock's .context.get(...) returns a truthy mock and would + # silently disable env fallback). + custom_component.graph.context = {} async def mock_get_variable(name, **_kwargs): if name == "DB_KEY": diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index 09cee00bf5..4927891fa5 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -89796,7 +89796,7 @@ "icon": "Mem0", "legacy": false, "metadata": { - "code_hash": "f2832e1bc71f", + "code_hash": "036442d6c64c", "dependencies": { "dependencies": [ { @@ -89863,7 +89863,7 @@ "show": true, "title_case": false, "type": "code", - "value": "import os\n\nfrom mem0 import Memory, MemoryClient\n\nfrom lfx.base.memory.model import LCChatMemoryComponent\nfrom lfx.inputs.inputs import DictInput, HandleInput, MessageTextInput, NestedDictInput, SecretStrInput\nfrom lfx.io import Output\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.utils.validate_cloud import raise_error_if_astra_cloud_disable_component\n\ndisable_component_in_astra_cloud_msg = (\n \"Mem0 chat memory is not supported in Astra cloud environment. Please use local storage mode or mem0 cloud.\"\n)\n\n\nclass Mem0MemoryComponent(LCChatMemoryComponent):\n display_name = \"Mem0 Chat Memory\"\n description = \"Retrieves and stores chat messages using Mem0 memory storage.\"\n name = \"mem0_chat_memory\"\n icon: str = \"Mem0\"\n inputs = [\n NestedDictInput(\n name=\"mem0_config\",\n display_name=\"Mem0 Configuration\",\n info=\"\"\"Configuration dictionary for initializing Mem0 memory instance.\n Example:\n {\n \"graph_store\": {\n \"provider\": \"neo4j\",\n \"config\": {\n \"url\": \"neo4j+s://your-neo4j-url\",\n \"username\": \"neo4j\",\n \"password\": \"your-password\"\n }\n },\n \"version\": \"v1.1\"\n }\"\"\",\n input_types=[\"Data\", \"JSON\"],\n ),\n MessageTextInput(\n name=\"ingest_message\",\n display_name=\"Message to Ingest\",\n info=\"The message content to be ingested into Mem0 memory.\",\n ),\n HandleInput(\n name=\"existing_memory\",\n display_name=\"Existing Memory Instance\",\n input_types=[\"Memory\"],\n info=\"Optional existing Mem0 memory instance. If not provided, a new instance will be created.\",\n ),\n MessageTextInput(\n name=\"user_id\", display_name=\"User ID\", info=\"Identifier for the user associated with the messages.\"\n ),\n MessageTextInput(\n name=\"search_query\", display_name=\"Search Query\", info=\"Input text for searching related memories in Mem0.\"\n ),\n SecretStrInput(\n name=\"mem0_api_key\",\n display_name=\"Mem0 API Key\",\n info=\"API key for Mem0 platform. Leave empty to use the local version.\",\n ),\n DictInput(\n name=\"metadata\",\n display_name=\"Metadata\",\n info=\"Additional metadata to associate with the ingested message.\",\n advanced=True,\n ),\n SecretStrInput(\n name=\"openai_api_key\",\n display_name=\"OpenAI API Key\",\n required=False,\n info=\"API key for OpenAI. Required if using OpenAI Embeddings without a provided configuration.\",\n ),\n ]\n\n outputs = [\n Output(name=\"memory\", display_name=\"Mem0 Memory\", method=\"ingest_data\"),\n Output(\n name=\"search_results\",\n display_name=\"Search Results\",\n method=\"build_search_results\",\n ),\n ]\n\n def build_mem0(self) -> Memory:\n \"\"\"Initializes a Mem0 memory instance based on provided configuration and API keys.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n if self.openai_api_key:\n os.environ[\"OPENAI_API_KEY\"] = self.openai_api_key\n\n try:\n if not self.mem0_api_key:\n return Memory.from_config(config_dict=dict(self.mem0_config)) if self.mem0_config else Memory()\n if self.mem0_config:\n return MemoryClient.from_config(api_key=self.mem0_api_key, config_dict=dict(self.mem0_config))\n return MemoryClient(api_key=self.mem0_api_key)\n except ImportError as e:\n msg = \"Mem0 is not properly installed. Please install it with 'pip install -U mem0ai'.\"\n raise ImportError(msg) from e\n\n def ingest_data(self) -> Memory:\n \"\"\"Ingests a new message into Mem0 memory and returns the updated memory instance.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n mem0_memory = self.existing_memory or self.build_mem0()\n\n if not self.ingest_message or not self.user_id:\n logger.warning(\"Missing 'ingest_message' or 'user_id'; cannot ingest data.\")\n return mem0_memory\n\n metadata = self.metadata or {}\n\n logger.info(\"Ingesting message for user_id: %s\", self.user_id)\n\n try:\n mem0_memory.add(self.ingest_message, user_id=self.user_id, metadata=metadata)\n except Exception:\n logger.exception(\"Failed to add message to Mem0 memory.\")\n raise\n\n return mem0_memory\n\n def build_search_results(self) -> Data:\n \"\"\"Searches the Mem0 memory for related messages based on the search query and returns the results.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n mem0_memory = self.ingest_data()\n search_query = self.search_query\n user_id = self.user_id\n\n logger.info(\"Search query: %s\", search_query)\n\n try:\n if search_query:\n logger.info(\"Performing search with query.\")\n related_memories = mem0_memory.search(query=search_query, filters={\"user_id\": user_id})\n else:\n logger.info(\"Retrieving all memories for user_id: %s\", user_id)\n related_memories = mem0_memory.get_all(filters={\"user_id\": user_id})\n except Exception:\n logger.exception(\"Failed to retrieve related memories from Mem0.\")\n raise\n\n logger.info(\"Related memories retrieved: %s\", related_memories)\n return related_memories\n" + "value": "from mem0 import Memory, MemoryClient\n\nfrom lfx.base.memory.model import LCChatMemoryComponent\nfrom lfx.inputs.inputs import DictInput, HandleInput, MessageTextInput, NestedDictInput, SecretStrInput\nfrom lfx.io import Output\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.utils.validate_cloud import raise_error_if_astra_cloud_disable_component\n\ndisable_component_in_astra_cloud_msg = (\n \"Mem0 chat memory is not supported in Astra cloud environment. Please use local storage mode or mem0 cloud.\"\n)\n\n\nclass Mem0MemoryComponent(LCChatMemoryComponent):\n display_name = \"Mem0 Chat Memory\"\n description = \"Retrieves and stores chat messages using Mem0 memory storage.\"\n name = \"mem0_chat_memory\"\n icon: str = \"Mem0\"\n inputs = [\n NestedDictInput(\n name=\"mem0_config\",\n display_name=\"Mem0 Configuration\",\n info=\"\"\"Configuration dictionary for initializing Mem0 memory instance.\n Example:\n {\n \"graph_store\": {\n \"provider\": \"neo4j\",\n \"config\": {\n \"url\": \"neo4j+s://your-neo4j-url\",\n \"username\": \"neo4j\",\n \"password\": \"your-password\"\n }\n },\n \"version\": \"v1.1\"\n }\"\"\",\n input_types=[\"Data\", \"JSON\"],\n ),\n MessageTextInput(\n name=\"ingest_message\",\n display_name=\"Message to Ingest\",\n info=\"The message content to be ingested into Mem0 memory.\",\n ),\n HandleInput(\n name=\"existing_memory\",\n display_name=\"Existing Memory Instance\",\n input_types=[\"Memory\"],\n info=\"Optional existing Mem0 memory instance. If not provided, a new instance will be created.\",\n ),\n MessageTextInput(\n name=\"user_id\", display_name=\"User ID\", info=\"Identifier for the user associated with the messages.\"\n ),\n MessageTextInput(\n name=\"search_query\", display_name=\"Search Query\", info=\"Input text for searching related memories in Mem0.\"\n ),\n SecretStrInput(\n name=\"mem0_api_key\",\n display_name=\"Mem0 API Key\",\n info=\"API key for Mem0 platform. Leave empty to use the local version.\",\n ),\n DictInput(\n name=\"metadata\",\n display_name=\"Metadata\",\n info=\"Additional metadata to associate with the ingested message.\",\n advanced=True,\n ),\n SecretStrInput(\n name=\"openai_api_key\",\n display_name=\"OpenAI API Key\",\n required=False,\n info=\"API key for OpenAI. Required if using OpenAI Embeddings without a provided configuration.\",\n ),\n ]\n\n outputs = [\n Output(name=\"memory\", display_name=\"Mem0 Memory\", method=\"ingest_data\"),\n Output(\n name=\"search_results\",\n display_name=\"Search Results\",\n method=\"build_search_results\",\n ),\n ]\n\n def build_mem0(self) -> Memory:\n \"\"\"Initializes a Mem0 memory instance based on provided configuration and API keys.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n\n # Pass the OpenAI key through mem0's config instead of writing os.environ.\n # os.environ is process-global and persists across requests, so in a shared\n # serving process (e.g. lfx serve) one caller's key would leak into other\n # concurrent requests. mem0's OpenAI llm/embedder read config.api_key before\n # falling back to OPENAI_API_KEY, so config injection keeps it request-scoped.\n config = dict(self.mem0_config) if self.mem0_config else {}\n if self.openai_api_key:\n for section in (\"llm\", \"embedder\"):\n sec = dict(config.get(section) or {})\n if sec.get(\"provider\", \"openai\") != \"openai\":\n continue # the OpenAI key is irrelevant to non-OpenAI providers\n sec[\"provider\"] = \"openai\"\n sec_config = dict(sec.get(\"config\") or {})\n sec_config.setdefault(\"api_key\", self.openai_api_key)\n sec[\"config\"] = sec_config\n config[section] = sec\n\n try:\n if not self.mem0_api_key:\n return Memory.from_config(config_dict=config) if config else Memory()\n if self.mem0_config:\n return MemoryClient.from_config(api_key=self.mem0_api_key, config_dict=dict(self.mem0_config))\n return MemoryClient(api_key=self.mem0_api_key)\n except ImportError as e:\n msg = \"Mem0 is not properly installed. Please install it with 'pip install -U mem0ai'.\"\n raise ImportError(msg) from e\n\n def ingest_data(self) -> Memory:\n \"\"\"Ingests a new message into Mem0 memory and returns the updated memory instance.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n mem0_memory = self.existing_memory or self.build_mem0()\n\n if not self.ingest_message or not self.user_id:\n logger.warning(\"Missing 'ingest_message' or 'user_id'; cannot ingest data.\")\n return mem0_memory\n\n metadata = self.metadata or {}\n\n logger.info(\"Ingesting message for user_id: %s\", self.user_id)\n\n try:\n mem0_memory.add(self.ingest_message, user_id=self.user_id, metadata=metadata)\n except Exception:\n logger.exception(\"Failed to add message to Mem0 memory.\")\n raise\n\n return mem0_memory\n\n def build_search_results(self) -> Data:\n \"\"\"Searches the Mem0 memory for related messages based on the search query and returns the results.\"\"\"\n # Check if we're in Astra cloud environment and raise an error if we are.\n raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg)\n mem0_memory = self.ingest_data()\n search_query = self.search_query\n user_id = self.user_id\n\n logger.info(\"Search query: %s\", search_query)\n\n try:\n if search_query:\n logger.info(\"Performing search with query.\")\n related_memories = mem0_memory.search(query=search_query, filters={\"user_id\": user_id})\n else:\n logger.info(\"Retrieving all memories for user_id: %s\", user_id)\n related_memories = mem0_memory.get_all(filters={\"user_id\": user_id})\n except Exception:\n logger.exception(\"Failed to retrieve related memories from Mem0.\")\n raise\n\n logger.info(\"Related memories retrieved: %s\", related_memories)\n return related_memories\n" }, "existing_memory": { "_input_type": "HandleInput", @@ -118806,6 +118806,6 @@ "num_components": 356, "num_modules": 95 }, - "sha256": "1362819987d1b6a65b5fec2661b9a7667fe4c7159b458ee438e83d47f7bdf028", + "sha256": "893aff1f6652bf41616a20b725ab7cd7faae1dedd88237634f8b4058d3734923", "version": "0.5.0" } diff --git a/src/lfx/src/lfx/base/models/unified_models/credentials.py b/src/lfx/src/lfx/base/models/unified_models/credentials.py index 2f38346de0..ffceb42b33 100644 --- a/src/lfx/src/lfx/base/models/unified_models/credentials.py +++ b/src/lfx/src/lfx/base/models/unified_models/credentials.py @@ -11,6 +11,7 @@ from uuid import UUID from lfx.log.logger import logger from lfx.services.deps import get_variable_service, session_scope +from lfx.services.variable.request_scope import is_env_fallback_disabled from lfx.utils.async_helpers import run_until_complete from lfx.utils.secrets import secret_value_to_str @@ -60,9 +61,12 @@ def get_api_key_for_provider(user_id: UUID | str | None, provider: str, api_key: value = secret_value_to_str(value, strip=True) if value: return value - env_value = os.environ.get(var_name) - if env_value and env_value.strip(): - return env_value.strip() + # Honor the request's no-env-fallback contract: skip os.environ when disabled so a + # served flow stays isolated from process-wide credentials (matches VariableService). + if not is_env_fallback_disabled(): + env_value = os.environ.get(var_name) + if env_value and env_value.strip(): + return env_value.strip() return None if api_key and api_key.strip(): @@ -115,6 +119,8 @@ def get_api_key_for_provider(user_id: UUID | str | None, provider: str, api_key: if api_key: return api_key + if is_env_fallback_disabled(): + return None return os.getenv(variable_name) @@ -147,8 +153,13 @@ def get_all_variables_for_provider(user_id: UUID | str | None, provider: str) -> if not provider_vars: return result - # If no user_id, only check environment variables + # If no user_id, only check environment variables. Honor the request's no-env-fallback + # contract: a served flow under no_env_fallback stays isolated from process-wide + # credentials, so return nothing rather than leaking os.environ into provider_vars + # (which would defeat the _env_if_allowed guards in instantiation.py). if user_id is None or (isinstance(user_id, str) and user_id == "None"): + if is_env_fallback_disabled(): + return result for var_info in provider_vars: var_key = var_info.get("variable_key") if var_key: @@ -183,7 +194,10 @@ def get_all_variables_for_provider(user_id: UUID | str | None, provider: str) -> if value: values[var_key] = value except (ValueError, Exception): # noqa: BLE001 - # Variable not found - check environment + # Variable not found - check environment, unless the request disables + # env fallback (keeps served flows isolated from process-wide credentials). + if is_env_fallback_disabled(): + continue env_value = _env_value_for(var_key) if env_value: values[var_key] = env_value diff --git a/src/lfx/src/lfx/base/models/unified_models/instantiation.py b/src/lfx/src/lfx/base/models/unified_models/instantiation.py index 6aa33bd788..1333b4157a 100644 --- a/src/lfx/src/lfx/base/models/unified_models/instantiation.py +++ b/src/lfx/src/lfx/base/models/unified_models/instantiation.py @@ -6,6 +6,7 @@ import os from typing import TYPE_CHECKING, Any from lfx.base.models.model_utils import _to_str +from lfx.services.variable.request_scope import is_env_fallback_disabled from .class_registry import EMBEDDING_PARAM_MAPPINGS, EMBEDDING_PROVIDER_CLASS_MAPPING from .provider_queries import model_provider_metadata @@ -14,6 +15,18 @@ if TYPE_CHECKING: from uuid import UUID +def _env_if_allowed(key: str) -> str | None: + """Return ``os.environ.get(key)`` unless the active request disables env fallback. + + Connection config (provider URLs, project IDs, attribution headers) falls back to + process env only when env fallback is allowed, so a served flow under + ``no_env_fallback`` stays isolated from process-wide environment. + """ + if is_env_fallback_disabled(): + return None + return os.environ.get(key) + + def get_llm( model, user_id: UUID | str | None, @@ -180,12 +193,12 @@ def get_llm( # Priority: component value > database value > env var watsonx_url_value = ( - watsonx_url if watsonx_url else provider_vars.get("WATSONX_URL") or os.environ.get("WATSONX_URL") + watsonx_url if watsonx_url else provider_vars.get("WATSONX_URL") or _env_if_allowed("WATSONX_URL") ) watsonx_project_id_value = ( watsonx_project_id if watsonx_project_id - else provider_vars.get("WATSONX_PROJECT_ID") or os.environ.get("WATSONX_PROJECT_ID") + else provider_vars.get("WATSONX_PROJECT_ID") or _env_if_allowed("WATSONX_PROJECT_ID") ) has_url = bool(watsonx_url_value) @@ -217,7 +230,7 @@ def get_llm( ollama_base_url_value = ( ollama_base_url if ollama_base_url - else provider_vars.get("OLLAMA_BASE_URL") or os.environ.get("OLLAMA_BASE_URL") + else provider_vars.get("OLLAMA_BASE_URL") or _env_if_allowed("OLLAMA_BASE_URL") ) if ollama_base_url_value: kwargs[base_url_param] = ollama_base_url_value @@ -240,7 +253,7 @@ def get_llm( # KeyError on a misconfigured metadata entry beats silently # skipping a header the operator expects to be sent. variable_key = var["variable_key"] - value = provider_vars.get(variable_key) or os.environ.get(variable_key) + value = provider_vars.get(variable_key) or _env_if_allowed(variable_key) if header_name and value: default_headers[header_name] = value if default_headers: @@ -380,11 +393,11 @@ def get_embeddings( # Watson-specific parameters if provider in {"IBM WatsonX", "IBM watsonx.ai"}: watsonx_provider_vars = unified_models_module.get_all_variables_for_provider(user_id, provider) - url_value = watsonx_url or watsonx_provider_vars.get("WATSONX_URL") or os.environ.get("WATSONX_URL") + url_value = watsonx_url or watsonx_provider_vars.get("WATSONX_URL") or _env_if_allowed("WATSONX_URL") pid_value = ( watsonx_project_id or watsonx_provider_vars.get("WATSONX_PROJECT_ID") - or os.environ.get("WATSONX_PROJECT_ID") + or _env_if_allowed("WATSONX_PROJECT_ID") ) has_url = bool(url_value) @@ -430,7 +443,7 @@ def get_embeddings( base_url_value = ( ollama_base_url or provider_vars.get("OLLAMA_BASE_URL") - or os.environ.get("OLLAMA_BASE_URL") + or _env_if_allowed("OLLAMA_BASE_URL") or "http://localhost:11434" ) kwargs[param_mapping["base_url"]] = base_url_value diff --git a/src/lfx/src/lfx/cli/_extension_commands.py b/src/lfx/src/lfx/cli/_extension_commands.py index 4be2a1d3df..00d7391213 100644 --- a/src/lfx/src/lfx/cli/_extension_commands.py +++ b/src/lfx/src/lfx/cli/_extension_commands.py @@ -708,7 +708,7 @@ def dev_command( "external dev-server scripts." ), ), - extra_args: list[str] | None = typer.Argument( # noqa: B008 - Typer requires call-site default + extra_args: list[str] | None = typer.Argument( None, help="Extra arguments forwarded to ``langflow run`` (after a ``--`` separator).", ), diff --git a/src/lfx/src/lfx/cli/_running_commands.py b/src/lfx/src/lfx/cli/_running_commands.py index 2447fe48c3..6c64186913 100644 --- a/src/lfx/src/lfx/cli/_running_commands.py +++ b/src/lfx/src/lfx/cli/_running_commands.py @@ -92,17 +92,24 @@ def register(app: typer.Typer) -> None: session_id=session_id, ) - @app.command(name="serve", help="Serve a flow as an API", no_args_is_help=True, rich_help_panel="Running") + @app.command(name="serve", help="Serve a flow as an API", rich_help_panel="Running") def serve_command_wrapper( - script_path: str | None = typer.Argument( - None, + script_paths: list[str] | None = typer.Argument( + default=None, help=( - "Path to JSON flow (.json) or Python script (.py) file or stdin input. " + "Path(s) to JSON flow file(s) (.json), Python script(s) (.py), or a directory " + "containing .json files (top-level only, non-recursive). " "Optional when using --flow-json or --stdin." ), ), host: str = typer.Option("127.0.0.1", "--host", "-h", help="Host to bind the server to"), port: int = typer.Option(8000, "--port", "-p", help="Port to bind the server to"), + workers: int = typer.Option( + 1, + "--workers", + "-w", + help="Number of uvicorn worker processes. Use with --flow-dir for multi-worker flow sharing.", + ), verbose: bool = typer.Option(False, "--verbose", "-v", help="Show diagnostic output and execution details"), env_file: str | None = typer.Option( None, @@ -117,36 +124,58 @@ def register(app: typer.Typer) -> None: flow_json: str | None = typer.Option( None, "--flow-json", - help="Inline JSON flow content as a string (alternative to script_path)", + help="Inline JSON flow content as a string (alternative to script_paths)", + ), + flow_dir: str | None = typer.Option( + None, + "--flow-dir", + help=( + "Directory for filesystem-backed flow storage. " + "All uvicorn workers sharing this path will serve the same flows. " + "Use /tmp/lfx-flows for single-pod sharing or a PVC mount for cross-pod. " + "Defaults to in-memory only when omitted." + ), ), *, stdin: bool = typer.Option( False, "--stdin", - help="Read JSON flow content from stdin (alternative to script_path)", + help="Read JSON flow content from stdin (alternative to script_paths)", ), check_variables: bool = typer.Option( True, "--check-variables/--no-check-variables", help="Check global variables for environment compatibility", ), + no_env_fallback: bool = typer.Option( + False, + "--no-env-fallback/--env-fallback", + help=( + "Disable os.environ fallback for credential variables. " + "Variables not supplied via global_vars on each request resolve to None " + "instead of reading from the process environment." + ), + ), ) -> None: """Serve LFX flows as a web API (lazy-loaded).""" from pathlib import Path from lfx.cli.commands import serve_command - # Convert env_file string to Path if provided env_file_path = Path(env_file) if env_file else None + flow_dir_path = Path(flow_dir) if flow_dir else None serve_command( - script_path=script_path, + script_paths=script_paths, host=host, port=port, + workers=workers, verbose=verbose, env_file=env_file_path, log_level=log_level, flow_json=flow_json, + flow_dir=flow_dir_path, stdin=stdin, check_variables=check_variables, + no_env_fallback=no_env_fallback, ) diff --git a/src/lfx/src/lfx/cli/commands.py b/src/lfx/src/lfx/cli/commands.py index 71429badfa..bd88658ea8 100644 --- a/src/lfx/src/lfx/cli/commands.py +++ b/src/lfx/src/lfx/cli/commands.py @@ -2,16 +2,17 @@ from __future__ import annotations +import asyncio +import contextlib import json import os import sys import tempfile -from functools import partial from pathlib import Path +from typing import TYPE_CHECKING import typer import uvicorn -from asyncer import syncify from dotenv import load_dotenv from rich.console import Console from rich.panel import Panel @@ -23,9 +24,15 @@ from lfx.cli.common import ( get_best_access_host, get_free_port, is_port_in_use, - load_graph_from_path, ) -from lfx.cli.serve_app import FlowMeta, create_multi_serve_app +from lfx.cli.script_loader import find_graph_variable, load_graph_from_script +from lfx.cli.serve_app import FlowAlreadyRegisteredError, FlowMeta, FlowRegistry, create_multi_serve_app +from lfx.load import load_flow_from_json + +if TYPE_CHECKING: + from collections.abc import Callable + + from lfx.cli.flow_store import FlowStore # Initialize console console = Console() @@ -34,17 +41,119 @@ console = Console() API_KEY_MASK_LENGTH = 8 -@partial(syncify, raise_sync_error=False) -async def serve_command( - script_path: str | None = typer.Argument( - None, +async def _build_serve_registry( + *, + script_paths: list[str] | None, + flow_json: str | None, + stdin: bool, + check_variables: bool, + no_env_fallback: bool, + flow_store: FlowStore, + verbose_print: Callable[[str], None], +) -> tuple[FlowRegistry, str | None]: + """Build the FlowRegistry from startup inputs. + + Returns (registry, temp_file_path_or_None). Caller must unlink the temp + file if not None. + """ + temp_file_to_cleanup: str | None = None + + if flow_json is not None or stdin: + if flow_json is not None: + try: + json_data = json.loads(flow_json) + except json.JSONDecodeError as e: + typer.echo(f"Error: Invalid JSON content: {e}", err=True) + raise typer.Exit(1) from e + else: + stdin_content = sys.stdin.read().strip() + if not stdin_content: + typer.echo("Error: No content received from stdin", err=True) + raise typer.Exit(1) + try: + json_data = json.loads(stdin_content) + except json.JSONDecodeError as e: + typer.echo(f"Error: Invalid JSON content from stdin: {e}", err=True) + raise typer.Exit(1) from e + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp: + json.dump(json_data, tmp, indent=2) + temp_file_to_cleanup = tmp.name + try: + registry = await build_registry_from_paths( + [Path(temp_file_to_cleanup)], + verbose_print, + check_variables=check_variables, + no_env_fallback=no_env_fallback, + store=flow_store, + ) + except ValueError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from e + + elif script_paths: + resolved = [Path(p).resolve() for p in script_paths] + missing = [p for p in resolved if not p.exists()] + if missing: + for m in missing: + typer.echo(f"Error: Path '{m}' does not exist.", err=True) + raise typer.Exit(1) + + if len(resolved) == 1 and resolved[0].is_dir(): + dir_path = resolved[0] + try: + registry = await build_registry_from_directory( + dir_path, + verbose_print, + check_variables=check_variables, + no_env_fallback=no_env_fallback, + store=flow_store, + ) + except ValueError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from e + verbose_print(f"Loaded {len(registry)} flow(s) from directory {dir_path}") + else: + non_supported = [p for p in resolved if p.suffix not in {".json", ".py"}] + if non_supported: + for p in non_supported: + typer.echo(f"Error: '{p}' must be a .json or .py file.", err=True) + raise typer.Exit(1) + try: + registry = await build_registry_from_paths( + resolved, + verbose_print, + check_variables=check_variables, + no_env_fallback=no_env_fallback, + store=flow_store, + ) + except ValueError as e: + typer.echo(f"Error: {e}", err=True) + raise typer.Exit(1) from e + + else: + registry = FlowRegistry(no_env_fallback=no_env_fallback, store=flow_store) + verbose_print("Starting with empty registry — flows can be uploaded at runtime") + + return registry, temp_file_to_cleanup + + +def serve_command( + script_paths: list[str] | None = typer.Argument( + default=None, help=( - "Path to JSON flow (.json) or Python script (.py) file or stdin input. " + "Path(s) to JSON flow file(s) (.json), Python script(s) (.py), or a directory " + "containing .json files (top-level only, non-recursive). " "Optional when using --flow-json or --stdin." ), ), host: str = typer.Option("127.0.0.1", "--host", "-h", help="Host to bind the server to"), port: int = typer.Option(8000, "--port", "-p", help="Port to bind the server to"), + workers: int = typer.Option( + 1, + "--workers", + "-w", + help="Number of uvicorn worker processes. Use with --flow-dir for multi-worker flow sharing.", + ), verbose: bool = typer.Option(False, "--verbose", "-v", help="Show diagnostic output and execution details"), # noqa: FBT001, FBT003 env_file: Path | None = typer.Option( None, @@ -59,274 +168,356 @@ async def serve_command( flow_json: str | None = typer.Option( None, "--flow-json", - help="Inline JSON flow content as a string (alternative to script_path)", + help="Inline JSON flow content as a string (alternative to script_paths)", + ), + flow_dir: Path | None = typer.Option( + None, + "--flow-dir", + help=( + "Directory for filesystem-backed flow storage. " + "All uvicorn workers sharing this path will serve the same flows. " + "Use /tmp/lfx-flows for single-pod sharing or a PVC mount for cross-pod. " + "Defaults to in-memory only when omitted." + ), ), *, stdin: bool = typer.Option( False, # noqa: FBT003 "--stdin", - help="Read JSON flow content from stdin (alternative to script_path)", + help="Read JSON flow content from stdin (alternative to script_paths)", ), check_variables: bool = typer.Option( True, # noqa: FBT003 "--check-variables/--no-check-variables", help="Check global variables for environment compatibility", ), + no_env_fallback: bool = typer.Option( + False, # noqa: FBT003 + "--no-env-fallback/--env-fallback", + help=( + "Disable os.environ fallback for credential variables. " + "Variables not supplied via global_vars on each request resolve to None " + "instead of reading from the process environment." + ), + ), ) -> None: """Serve LFX flows as a web API. - Supports single files, inline JSON, and stdin input. - - Examples: - # Serve from file - lfx serve my_flow.json - - # Serve inline JSON - lfx serve --flow-json '{"nodes": [...], "edges": [...]}' - - # Serve from stdin - cat my_flow.json | lfx serve --stdin - echo '{"nodes": [...]}' | lfx serve --stdin + Supports single files, inline JSON, stdin, multiple explicit files, + and a directory of JSON flows. """ - # Configure logging with the specified level and import logger - from lfx.log.logger import configure, logger - - configure(log_level=log_level) - - verbose_print = create_verbose_printer(verbose=verbose) - - # Validate input sources - exactly one must be provided - input_sources = [script_path is not None, flow_json is not None, stdin] - if sum(input_sources) != 1: - if sum(input_sources) == 0: - verbose_print("Error: Must provide either script_path, --flow-json, or --stdin") - else: - verbose_print("Error: Cannot use script_path, --flow-json, and --stdin together. Choose exactly one.") - raise typer.Exit(1) - - # Load environment variables from .env file if provided - if env_file: - if not env_file.exists(): - verbose_print(f"Error: Environment file '{env_file}' does not exist.") - raise typer.Exit(1) - - verbose_print(f"Loading environment variables from: {env_file}") - load_dotenv(env_file) - - # Validate API key - try: - api_key = get_api_key() - verbose_print("✓ LANGFLOW_API_KEY is configured") - except ValueError as e: - typer.echo(f"✗ {e}", err=True) - typer.echo("Set the LANGFLOW_API_KEY environment variable before serving.", err=True) - raise typer.Exit(1) from e - - # Validate log level - valid_log_levels = {"debug", "info", "warning", "error", "critical"} - if log_level.lower() not in valid_log_levels: - verbose_print(f"Error: Invalid log level '{log_level}'. Must be one of: {', '.join(sorted(valid_log_levels))}") - raise typer.Exit(1) - - # Configure logging with the specified level - # Disable pretty logs for serve command to avoid ANSI codes in API responses - os.environ["LANGFLOW_PRETTY_LOGS"] = "false" - verbose_print(f"Configuring logging with level: {log_level}") from lfx.log.logger import configure configure(log_level=log_level) + verbose_print = create_verbose_printer(verbose=verbose) - # ------------------------------------------------------------------ - # Handle inline JSON content or stdin input - # ------------------------------------------------------------------ - temp_file_to_cleanup = None + # Validate at most one input source + has_paths = bool(script_paths) + input_sources = [has_paths, flow_json is not None, stdin] + if sum(input_sources) > 1: + typer.echo("Error: Cannot combine path(s), --flow-json, and --stdin. Choose exactly one.", err=True) + raise typer.Exit(1) - if flow_json is not None: - logger.info("Processing inline JSON content...") - try: - # Validate JSON syntax - json_data = json.loads(flow_json) - logger.info("JSON content is valid") - - # Create a temporary file with the JSON content - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as temp_file: - json.dump(json_data, temp_file, indent=2) - temp_file_to_cleanup = temp_file.name - - script_path = temp_file_to_cleanup - logger.info(f"Created temporary file: {script_path}") - - except json.JSONDecodeError as e: - typer.echo(f"Error: Invalid JSON content: {e}", err=True) - raise typer.Exit(1) from e - except Exception as e: - verbose_print(f"Error processing JSON content: {e}") - raise typer.Exit(1) from e - - elif stdin: - logger.info("Reading JSON content from stdin...") - try: - # Read all content from stdin - stdin_content = sys.stdin.read().strip() - if not stdin_content: - logger.error("No content received from stdin") - raise typer.Exit(1) - - # Validate JSON syntax - json_data = json.loads(stdin_content) - logger.info("JSON content from stdin is valid") - - # Create a temporary file with the JSON content - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as temp_file: - json.dump(json_data, temp_file, indent=2) - temp_file_to_cleanup = temp_file.name - - script_path = temp_file_to_cleanup - logger.info(f"Created temporary file from stdin: {script_path}") - - except json.JSONDecodeError as e: - verbose_print(f"Error: Invalid JSON content from stdin: {e}") - raise typer.Exit(1) from e - except Exception as e: - verbose_print(f"Error reading from stdin: {e}") - raise typer.Exit(1) from e + if env_file: + if not env_file.exists(): + typer.echo(f"Error: Environment file '{env_file}' does not exist.", err=True) + raise typer.Exit(1) + verbose_print(f"Loading environment variables from: {env_file}") + load_dotenv(env_file) try: - # Load the graph - if script_path is None: - verbose_print("Error: script_path is None after input validation") - raise typer.Exit(1) + api_key = get_api_key() + verbose_print("LANGFLOW_API_KEY is configured") + except ValueError as e: + typer.echo(str(e), err=True) + typer.echo("Set the LANGFLOW_API_KEY environment variable before serving.", err=True) + raise typer.Exit(1) from e - resolved_path = Path(script_path).resolve() + valid_log_levels = {"debug", "info", "warning", "error", "critical"} + if log_level.lower() not in valid_log_levels: + typer.echo( + f"Error: Invalid log level '{log_level}'. Must be one of: {', '.join(sorted(valid_log_levels))}", + err=True, + ) + raise typer.Exit(1) - if not resolved_path.exists(): - typer.echo(f"Error: File '{resolved_path}' does not exist.", err=True) - raise typer.Exit(1) + if workers < 1: + typer.echo("Error: --workers must be at least 1.", err=True) + raise typer.Exit(1) - if resolved_path.suffix == ".json": - graph = await load_graph_from_path(resolved_path, resolved_path.suffix, verbose_print, verbose=verbose) - elif resolved_path.suffix == ".py": - verbose_print("Loading graph from Python script...") - from lfx.cli.script_loader import load_graph_from_script + os.environ["LANGFLOW_PRETTY_LOGS"] = "false" + configure(log_level=log_level) - graph = await load_graph_from_script(resolved_path) - verbose_print("✓ Graph loaded from Python script") - else: - err_msg = "Error: Only JSON flow files (.json) or Python scripts (.py) are supported. " - err_msg += f"Got: {resolved_path.suffix}" - verbose_print(err_msg) - raise typer.Exit(1) + from lfx.cli.flow_store import FilesystemFlowStore, NullFlowStore - # Prepare the graph - logger.info("Preparing graph for serving...") - try: - graph.prepare() - logger.info("Graph prepared successfully") + flow_store = FilesystemFlowStore(flow_dir) if flow_dir else NullFlowStore() - # Validate global variables for environment compatibility - if check_variables: - from lfx.cli.validation import validate_global_variables_for_env - - validation_errors = validate_global_variables_for_env(graph) - if validation_errors: - logger.error("Global variable validation failed:") - for error in validation_errors: - logger.error(f" - {error}") - raise typer.Exit(1) - else: - logger.info("Global variable validation skipped") - except Exception as e: - verbose_print(f"✗ Failed to prepare graph: {e}") - raise typer.Exit(1) from e - - # Check if port is in use - if is_port_in_use(port, host): - available_port = get_free_port(port) - if verbose: - verbose_print(f"Port {port} is in use, using port {available_port} instead") - port = available_port - - # Create single-flow metadata - flow_id = flow_id_from_path(resolved_path, resolved_path.parent) - graph.flow_id = flow_id # annotate graph for reference - - title = resolved_path.stem - description = None - - metas = { - flow_id: FlowMeta( - id=flow_id, - relative_path=str(resolved_path.name), - title=title, - description=description, - ) - } - graphs = {flow_id: graph} - - source_display = "inline JSON" if flow_json else "stdin" if stdin else str(resolved_path) - verbose_print(f"✓ Prepared single flow '{title}' from {source_display} (id={flow_id})") - - # Create FastAPI app - serve_app = create_multi_serve_app( - root_dir=resolved_path.parent, - graphs=graphs, - metas=metas, - verbose_print=verbose_print, + if workers > 1 and flow_dir is None: + typer.echo( + "Warning: --workers > 1 without --flow-dir means each worker has an isolated " + "in-memory registry. Flows uploaded to one worker will not be visible to others. " + "Pass --flow-dir to enable shared flow storage across workers.", + err=True, ) - verbose_print("🚀 Starting single-flow server...") + if workers > 1 and flow_dir is not None and script_paths: + # With --flow-dir, workers skip LFX_SERVE_STARTUP_PATHS and rely on warm_from_store(). + # .py files produce no raw_json so they can't be written to the store — workers would + # silently start without those flows. Without --flow-dir workers re-execute the .py + # file directly from LFX_SERVE_STARTUP_PATHS, which works fine. + py_paths = [p for p in script_paths if Path(p).suffix == ".py"] + if py_paths: + for p in py_paths: + typer.echo( + f"Error: '{Path(p).name}' (.py) cannot be used with --workers > 1 and --flow-dir. " + "Python graphs cannot be serialized to the store, so workers would start without " + "those flows. Use a .json export, omit --flow-dir, or use --workers 1.", + err=True, + ) + raise typer.Exit(1) + + # Determine display name for startup panel + if flow_json is not None: + source_display = "inline JSON" + elif stdin: + source_display = "stdin" + elif script_paths: + resolved_display = [Path(p).resolve() for p in script_paths] + if len(resolved_display) == 1 and resolved_display[0].is_dir(): + source_display = str(resolved_display[0]) + else: + source_display = ", ".join(Path(p).name for p in script_paths) + else: + source_display = "none (upload flows via POST /flows/upload/)" + + temp_file_to_cleanup: str | None = None + try: + registry, temp_file_to_cleanup = asyncio.run( + _build_serve_registry( + script_paths=script_paths, + flow_json=flow_json, + stdin=stdin, + check_variables=check_variables, + no_env_fallback=no_env_fallback, + flow_store=flow_store, + verbose_print=verbose_print, + ) + ) + + if flow_dir: + if workers == 1: + # Single-worker: warm now so the startup panel shows the real count and the + # first request to any pre-existing flow doesn't pay a cold-load penalty. + # Multi-worker: each worker warms inside create_serve_app(); the parent can't + # safely warm and pass graphs across processes, so we skip it here. + registry.warm_from_store() + verbose_print(f"Flow store at {flow_dir} ({len(registry)} flows available)") + + if is_port_in_use(port, host): + port = get_free_port(port) + verbose_print(f"Port in use; using {port} instead") + + verbose_print("Starting server...") protocol = "http" access_host = get_best_access_host(host) - masked_key = f"{api_key[:API_KEY_MASK_LENGTH]}..." if len(api_key) > API_KEY_MASK_LENGTH else "***" + server_line = f"{protocol}://{access_host}:{port}" + if access_host != host: + server_line += f" [dim](bound to {host}:{port})[/dim]" console.print() console.print( Panel.fit( - f"[bold green]🎯 Single Flow Served Successfully![/bold green]\n\n" + f"[bold green]LFX Server Ready[/bold green]\n\n" f"[bold]Source:[/bold] {source_display}\n" - f"[bold]Server:[/bold] {protocol}://{access_host}:{port}\n" + f"[bold]Flows:[/bold] {len(registry)}\n" + f"[bold]Workers:[/bold] {workers}\n" + f"[bold]Server:[/bold] {server_line}\n" f"[bold]API Key:[/bold] {masked_key}\n\n" - f"[dim]Send POST requests to:[/dim]\n" - f"[blue]{protocol}://{access_host}:{port}/flows/{flow_id}/run[/blue]\n\n" - f"[dim]With headers:[/dim]\n" - f"[blue]x-api-key: {masked_key}[/blue]\n\n" - f"[dim]Or query parameter:[/dim]\n" - f"[blue]?x-api-key={masked_key}[/blue]\n\n" - f"[dim]Request body:[/dim]\n" - f"[blue]{{'input_value': 'Your input message'}}[/blue]", + f"[dim]List flows:[/dim]\n" + f"[blue]{protocol}://{access_host}:{port}/flows[/blue]\n\n" + f"[dim]Upload new flow:[/dim]\n" + f"[blue]POST {protocol}://{access_host}:{port}/flows/upload/[/blue]\n\n" + f"[dim]Run a flow:[/dim]\n" + f"[blue]POST {protocol}://{access_host}:{port}/flows/{{flow_id}}/run[/blue]", title="[bold blue]LFX Server[/bold blue]", border_style="blue", ) ) console.print() - # Start the server - # Use uvicorn.Server to properly handle async context - # uvicorn.run() uses asyncio.run() internally which fails when - # an event loop is already running (due to syncify decorator) try: - config = uvicorn.Config( - serve_app, - host=host, - port=port, - log_level=log_level, - ) - server = uvicorn.Server(config) - await server.serve() + if workers > 1: + # uvicorn requires an import string (not an app object) for multi-worker mode. + # Set env vars so each worker's create_serve_app() factory can reconstruct config. + # The parent's in-memory app is never passed to workers — each worker calls + # create_serve_app() fresh, so we skip building the app here. + from lfx.cli.serve_app import ( + _SERVE_FLOW_DIR_ENV, + _SERVE_NO_ENV_FALLBACK_ENV, + _SERVE_STARTUP_PATHS_ENV, + ) + + 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" + + # 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(). + # When flow_dir is NOT set, workers must re-read the original files. + startup_paths_for_workers: list[str] = [] + if not flow_dir: + if script_paths: + startup_paths_for_workers = [str(Path(p).resolve()) for p in script_paths] + elif temp_file_to_cleanup: + startup_paths_for_workers = [temp_file_to_cleanup] + os.environ[_SERVE_STARTUP_PATHS_ENV] = json.dumps(startup_paths_for_workers) + try: + uvicorn.run( + "lfx.cli.serve_app:create_serve_app", + host=host, + port=port, + workers=workers, + log_level=log_level, + factory=True, + ) + 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): + os.environ.pop(k, None) + else: + serve_app = create_multi_serve_app(registry=registry) + uvicorn.run(serve_app, host=host, port=port, workers=1, log_level=log_level) except KeyboardInterrupt: - verbose_print("\n👋 Server stopped") + verbose_print("\nServer stopped") raise typer.Exit(0) from None except Exception as e: - verbose_print(f"✗ Failed to start server: {e}") + verbose_print(f"Failed to start server: {e}") raise typer.Exit(1) from e finally: - # Clean up temporary file if created if temp_file_to_cleanup: - try: + with contextlib.suppress(OSError): Path(temp_file_to_cleanup).unlink() - verbose_print(f"✓ Cleaned up temporary file: {temp_file_to_cleanup}") - except OSError as e: - verbose_print(f"Warning: Failed to clean up temporary file {temp_file_to_cleanup}: {e}") + + +async def _load_graph_and_meta( + path: Path, + root_dir: Path, + *, + check_variables: bool, +) -> tuple: + """Load and prepare one graph, returning (graph, FlowMeta, raw_json | None). + + raw_json is the parsed flow dict for .json files; None for .py files + (which cannot be round-tripped to JSON for store persistence). + """ + raw_json: dict | None = None + try: + if path.suffix == ".py": + find_graph_variable(path) # validates a 'graph' variable exists + graph = await load_graph_from_script(path) + else: + raw_json = json.loads(path.read_text(encoding="utf-8")) + graph = load_flow_from_json(raw_json) + except Exception as exc: + msg = f"Failed to load {path.name}: {exc}" + raise ValueError(msg) from exc + graph.prepare() + if check_variables: + from lfx.cli.validation import validate_global_variables_for_env + + errors = validate_global_variables_for_env(graph) + if errors: + msg = f"Global variable validation failed for {path.name}: {'; '.join(errors)}" + raise ValueError(msg) + path_flow_id = flow_id_from_path(path, root_dir) + # Prefer the JSON's own id so the startup ID matches what workers reconstruct + # from the store (which also prefers raw_json["id"]). Fall back to the + # path-derived UUID when the JSON has no id (e.g. hand-written flows). + flow_id = (raw_json.get("id") if raw_json else None) or path_flow_id + graph.flow_id = flow_id + meta = FlowMeta( + id=flow_id, + relative_path=str(path.relative_to(root_dir)), + title=(raw_json.get("name") or path.stem) if raw_json else path.stem, + description=(raw_json.get("description")) if raw_json else None, + ) + return graph, meta, raw_json + + +async def _populate_registry( + paths: list[Path], + root_dir: Path, + registry: FlowRegistry, + verbose_print: Callable[[str], None], + *, + check_variables: bool, +) -> None: + """Load each path into *registry*, collecting errors and raising at the end.""" + errors: list[str] = [] + for path in paths: + try: + graph, meta, raw_json = await _load_graph_and_meta(path, root_dir, check_variables=check_variables) + registry.add(graph, meta, raw_json=raw_json) + verbose_print(f"Loaded flow '{meta.title}' (id={meta.id})") + except FlowAlreadyRegisteredError: + verbose_print(f"Skipping duplicate flow id={meta.id} from {path.name}") + except Exception as exc: # noqa: BLE001 + errors.append(f"{path.name}: {exc}") + if errors: + msg = "Failed to load flows:\n" + "\n".join(f" - {e}" for e in errors) + raise ValueError(msg) + + +async def build_registry_from_directory( + dir_path: Path, + verbose_print: Callable[[str], None], + *, + check_variables: bool, + no_env_fallback: bool = False, + store: FlowStore | None = None, +) -> FlowRegistry: + """Build a FlowRegistry by scanning *dir_path* for ``*.json`` files (non-recursive). + + Callers that want pre-existing store flows (e.g. from a prior run) to be + reachable are responsible for calling ``registry.warm_from_store()`` after + this returns. ``serve_command`` does this in the right place; calling it + here too would cause double-loading when ``store`` backs the same directory. + """ + from lfx.cli.flow_store import NullFlowStore + + json_files = sorted(dir_path.glob("*.json")) + if not json_files: + msg = f"No .json files found in directory: {dir_path}" + raise ValueError(msg) + + registry = FlowRegistry(no_env_fallback=no_env_fallback, store=store or NullFlowStore()) + await _populate_registry(json_files, dir_path, registry, verbose_print, check_variables=check_variables) + return registry + + +async def build_registry_from_paths( + paths: list[Path], + verbose_print: Callable[[str], None], + *, + check_variables: bool, + no_env_fallback: bool = False, + store: FlowStore | None = None, +) -> FlowRegistry: + """Build a FlowRegistry from an explicit list of ``.json`` or ``.py`` paths. + + Callers that want pre-existing store flows to be reachable are responsible + for calling ``registry.warm_from_store()`` after this returns. + ``serve_command`` does this in the right place. + """ + from lfx.cli.flow_store import NullFlowStore + + # Use a shared root so same-named files in different directories get distinct IDs. + common_root = ( + Path(os.path.commonpath([str(p) for p in paths])) if len(paths) > 1 else paths[0].parent if paths else Path() + ) + registry = FlowRegistry(no_env_fallback=no_env_fallback, store=store or NullFlowStore()) + await _populate_registry(paths, common_root, registry, verbose_print, check_variables=check_variables) + return registry diff --git a/src/lfx/src/lfx/cli/common.py b/src/lfx/src/lfx/cli/common.py index 25ee5ab84c..9f3cb67ac9 100644 --- a/src/lfx/src/lfx/cli/common.py +++ b/src/lfx/src/lfx/cli/common.py @@ -23,6 +23,7 @@ from urllib.parse import urlparse import httpx import typer +from lfx.cli.runtime_variables import build_request_variables_from_global_vars from lfx.cli.script_loader import ( extract_structured_result, find_graph_variable, @@ -31,6 +32,12 @@ from lfx.cli.script_loader import ( from lfx.load import load_flow_from_json from lfx.run._defaults import apply_run_defaults, resolve_fallback_to_env_vars from lfx.schema.schema import InputValueRequest +from lfx.services.variable.request_scope import ( + activate_no_env_fallback, + activate_request_variables, + reset_no_env_fallback, + reset_request_variables, +) if TYPE_CHECKING: from types import ModuleType @@ -335,6 +342,10 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id: fallback_to_env_vars = resolve_fallback_to_env_vars() + scope_vars = build_request_variables_from_global_vars(graph.context.get("request_variables")) + scope_token = activate_request_variables(scope_vars or None) + no_env_fallback_token = activate_no_env_fallback(disabled=bool(graph.context.get("no_env_fallback"))) + try: sys.stdout = captured_stdout sys.stderr = captured_stderr @@ -347,6 +358,8 @@ async def execute_graph_with_capture(graph, input_value: str | None, session_id: exc.args = (f"{exc.args[0] if exc.args else str(exc)}\n\nCaptured stderr:\n{error_output}",) raise finally: + reset_no_env_fallback(no_env_fallback_token) + reset_request_variables(scope_token) # Restore original stdout/stderr sys.stdout = original_stdout sys.stderr = original_stderr diff --git a/src/lfx/src/lfx/cli/flow_store.py b/src/lfx/src/lfx/cli/flow_store.py new file mode 100644 index 0000000000..c2a00a1577 --- /dev/null +++ b/src/lfx/src/lfx/cli/flow_store.py @@ -0,0 +1,149 @@ +"""Flow persistence backends for lfx serve. + +Three deployment modes — all use the same code path: + +- ``NullFlowStore`` (default): in-memory only, no disk writes. +- ``FilesystemFlowStore(directory)``: writes each flow as ``{id}.json``. + Point at ``/tmp/lfx-flows`` for single-pod sharing across uvicorn workers, + or at a PVC mount path for cross-pod sharing. +""" + +from __future__ import annotations + +import json +import uuid +from pathlib import Path +from typing import Protocol, runtime_checkable + +from lfx.log.logger import logger + + +@runtime_checkable +class FlowStore(Protocol): + """Persistence backend for flow JSON. + + Implementations must be safe to call from multiple OS processes sharing + the same filesystem (i.e. all writes must be atomic). + """ + + def write(self, flow_id: str, flow_json: dict) -> None: + """Persist *flow_json* under *flow_id*. Overwrites if already present.""" + ... + + def read(self, flow_id: str) -> dict | None: + """Return the stored JSON dict for *flow_id*, or ``None`` if absent.""" + ... + + def delete(self, flow_id: str) -> bool: + """Remove *flow_id* from the store. Returns True if it existed.""" + ... + + def list_ids(self) -> list[str]: + """Return all stored flow IDs.""" + ... + + def exists(self, flow_id: str) -> bool: + """Return True if the store has a record for *flow_id*.""" + ... + + @property + def is_persistent(self) -> bool: + """True if writes are visible to other processes (e.g. filesystem store).""" + ... + + +class NullFlowStore: + """No-op store — flows live in worker memory only (default).""" + + def write(self, _flow_id: str, _flow_json: dict) -> None: + pass + + def read(self, _flow_id: str) -> dict | None: + return None + + def delete(self, _flow_id: str) -> bool: + return False + + def list_ids(self) -> list[str]: + return [] + + def exists(self, _flow_id: str) -> bool: + return False + + @property + def is_persistent(self) -> bool: + return False + + +class FilesystemFlowStore: + """Filesystem-backed store. + + Each flow is persisted as ``{directory}/{flow_id}.json``. Writes are + atomic (write-to-tmp then replace) so concurrent readers in other workers + never see a partial file. + + Use ``/tmp/lfx-flows`` for single-pod worker sharing or a PVC mount path + for cross-pod sharing — the implementation is identical either way. + """ + + def __init__(self, directory: Path) -> None: + self._dir = Path(directory) + self._dir.mkdir(parents=True, exist_ok=True) + + def _safe_path(self, flow_id: str) -> Path: + """Return the path for *flow_id*, raising ValueError for unsafe IDs.""" + if "/" in flow_id or "\\" in flow_id or flow_id.startswith("."): + msg = f"Invalid flow_id: {flow_id!r}" + raise ValueError(msg) + path = self._dir / f"{flow_id}.json" + try: + path.resolve().relative_to(self._dir.resolve()) + except ValueError as exc: + msg = f"Invalid flow_id: {flow_id!r}" + raise ValueError(msg) from exc + return path + + def write(self, flow_id: str, flow_json: dict) -> None: + target = self._safe_path(flow_id) + tmp = self._dir / f"{flow_id}.{uuid.uuid4().hex}.tmp" + try: + tmp.write_text(json.dumps(flow_json), encoding="utf-8") + tmp.replace(target) + except OSError: + # Don't leak the partial temp file (e.g. disk full mid-write); list_ids() + # only globs *.json so an orphaned *.tmp would silently accumulate. + tmp.unlink(missing_ok=True) + raise + + def read(self, flow_id: str) -> dict | None: + path = self._safe_path(flow_id) + try: + return json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return None + except (json.JSONDecodeError, OSError, UnicodeDecodeError) as exc: + # A corrupted/unreadable file in a shared store (e.g. an external writer + # killed mid-write, or PVC corruption) must not take down the worker. + # Treat it as absent (logged) so warm_from_store/list_metas skip it and + # every other flow keeps serving, instead of raising into a 500/startup crash. + logger.warning("Skipping unreadable flow store file %s: %r", path.name, exc) + return None + + def delete(self, flow_id: str) -> bool: + path = self._safe_path(flow_id) + try: + path.unlink() + except FileNotFoundError: + return False + else: + return True + + def list_ids(self) -> list[str]: + return [p.stem for p in self._dir.glob("*.json")] + + def exists(self, flow_id: str) -> bool: + return self._safe_path(flow_id).exists() + + @property + def is_persistent(self) -> bool: + return True diff --git a/src/lfx/src/lfx/cli/run.py b/src/lfx/src/lfx/cli/run.py index 0098b54c3f..d5c2e91a43 100644 --- a/src/lfx/src/lfx/cli/run.py +++ b/src/lfx/src/lfx/cli/run.py @@ -49,7 +49,7 @@ def _check_langchain_version_compatibility(error_message: str) -> str | None: @partial(syncify, raise_sync_error=False) async def run( - script_path: Path | None = typer.Argument( # noqa: B008 + script_path: Path | None = typer.Argument( None, help="Path to the Python script (.py) or JSON flow (.json) containing a graph" ), input_value: str | None = typer.Argument(None, help="Input value to pass to the graph"), diff --git a/src/lfx/src/lfx/cli/runtime_variables.py b/src/lfx/src/lfx/cli/runtime_variables.py new file mode 100644 index 0000000000..518c5aad98 --- /dev/null +++ b/src/lfx/src/lfx/cli/runtime_variables.py @@ -0,0 +1,70 @@ +"""Normalize runtime ``global_vars`` payloads (WXO ADK / TRM contract).""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from lfx.log.logger import logger +from lfx.services.variable.request_scope import normalize_parsed_variables + +if TYPE_CHECKING: + from lfx.graph.graph.base import Graph + +_LANGFLOW_REQUEST_VARIABLES_KEY = "LANGFLOW_REQUEST_VARIABLES" + + +def build_request_variables_from_global_vars(global_vars: dict[str, str] | None) -> dict[str, str]: + """Flatten TRM-style ``global_vars`` into a single lookup map. + + TRM sends: + - ``LANGFLOW_REQUEST_VARIABLES``: JSON blob of merged context + connection_details + - raw credential keys (e.g. ``access_token``, ``wxo_github_access_token``) + - ``x-langflow-global-var-*`` aliases + + Parsed JSON is applied first; explicit keys in *global_vars* (except the JSON + blob key itself) override on collision. Under the TRM contract connection_details + is sent as raw keys and context inside the JSON blob, so connection_details wins. + """ + if not global_vars: + return {} + + merged: dict[str, str] = {} + raw_json = global_vars.get(_LANGFLOW_REQUEST_VARIABLES_KEY) + if raw_json: + try: + parsed = json.loads(raw_json) + except json.JSONDecodeError as exc: + logger.warning( + f"Failed to parse {_LANGFLOW_REQUEST_VARIABLES_KEY} JSON ({exc}); " + "credentials carried only in this blob will be unavailable." + ) + parsed = None + if isinstance(parsed, dict): + merged.update(normalize_parsed_variables(parsed)) + elif parsed is not None: + logger.warning( + f"{_LANGFLOW_REQUEST_VARIABLES_KEY} must be a JSON object, got {type(parsed).__name__}; ignoring blob." + ) + + for key, value in global_vars.items(): + if key == _LANGFLOW_REQUEST_VARIABLES_KEY: + continue + normalized = key.strip() if isinstance(key, str) else "" + if normalized: + merged[normalized] = str(value) + + return merged + + +def apply_global_vars_to_graph(graph: Graph, global_vars: dict[str, str] | None) -> None: + """Inject *global_vars* into ``graph.context['request_variables']``.""" + if not global_vars: + return + if "request_variables" not in graph.context: + graph.context["request_variables"] = {} + # Flatten the LANGFLOW_REQUEST_VARIABLES blob here so the no-DB load_from_db path + # (load_from_env_vars reads graph.context) resolves the same names as the + # VariableService/ContextVar path (common.py) — otherwise a caller sending a credential + # only inside the JSON blob would resolve on one path and return None on the other. + graph.context["request_variables"].update(build_request_variables_from_global_vars(global_vars)) diff --git a/src/lfx/src/lfx/cli/serve_app.py b/src/lfx/src/lfx/cli/serve_app.py index 6e2851e89f..8e7fc00cab 100644 --- a/src/lfx/src/lfx/cli/serve_app.py +++ b/src/lfx/src/lfx/cli/serve_app.py @@ -18,27 +18,42 @@ endpoints require the ``x-api-key`` header (or query parameter) validated by from __future__ import annotations import asyncio +import json import time +import traceback +import uuid from copy import deepcopy from typing import TYPE_CHECKING, Annotated, Any -from fastapi import APIRouter, Depends, FastAPI, HTTPException, Security -from fastapi.responses import StreamingResponse +from fastapi import Depends, FastAPI, HTTPException, Response, Security +from fastapi.responses import JSONResponse, StreamingResponse from fastapi.security import APIKeyHeader, APIKeyQuery -from pydantic import BaseModel, Field +from pydantic import BaseModel, ConfigDict, Field, field_validator -from lfx.cli.common import execute_graph_with_capture, extract_result_data, get_api_key +from lfx.cli.common import ( + execute_graph_with_capture, + extract_result_data, + get_api_key, +) +from lfx.cli.runtime_variables import apply_global_vars_to_graph +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 if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Callable - from pathlib import Path + from collections.abc import AsyncGenerator + from lfx.cli.flow_store import FlowStore from lfx.graph import Graph # Security - use the same pattern as Langflow main API API_KEY_NAME = "x-api-key" + +# Constants for app factory env vars (used by uvicorn worker processes) +_SERVE_ENV_PREFIX = "LFX_SERVE_" +_SERVE_FLOW_DIR_ENV = f"{_SERVE_ENV_PREFIX}FLOW_DIR" +_SERVE_NO_ENV_FALLBACK_ENV = f"{_SERVE_ENV_PREFIX}NO_ENV_FALLBACK" +_SERVE_STARTUP_PATHS_ENV = f"{_SERVE_ENV_PREFIX}STARTUP_PATHS" api_key_query = APIKeyQuery(name=API_KEY_NAME, scheme_name="API key query", auto_error=False) api_key_header = APIKeyHeader(name=API_KEY_NAME, scheme_name="API key header", auto_error=False) @@ -62,161 +77,14 @@ def verify_api_key( return provided_key -def _analyze_graph_structure(graph: Graph) -> dict[str, Any]: - """Analyze the graph structure to extract dynamic documentation information. - - Args: - graph: The LFX graph to analyze - - Returns: - dict: Graph analysis including components, input/output types, and flow details - """ - analysis: dict[str, Any] = { - "components": [], - "input_types": set(), - "output_types": set(), - "node_count": 0, - "edge_count": 0, - "entry_points": [], - "exit_points": [], - } - - try: - # Analyze nodes - for node_id, node in graph.nodes.items(): - analysis["node_count"] += 1 - component_info = { - "id": node_id, - "type": node.data.get("type", "Unknown"), - "name": node.data.get("display_name", node.data.get("type", "Unknown")), - "description": node.data.get("description", ""), - "template": node.data.get("template", {}), - } - analysis["components"].append(component_info) - - # Identify entry points (nodes with no incoming edges) - if not any(edge.source == node_id for edge in graph.edges): - analysis["entry_points"].append(component_info) - - # Identify exit points (nodes with no outgoing edges) - if not any(edge.target == node_id for edge in graph.edges): - analysis["exit_points"].append(component_info) - - # Analyze edges - analysis["edge_count"] = len(graph.edges) - - # Try to determine input/output types from entry/exit points - for entry in analysis["entry_points"]: - template = entry.get("template", {}) - for field_config in template.values(): - if field_config.get("type") in ["str", "text", "string"]: - analysis["input_types"].add("text") - elif field_config.get("type") in ["int", "float", "number"]: - analysis["input_types"].add("numeric") - elif field_config.get("type") in ["file", "path"]: - analysis["input_types"].add("file") - - for exit_point in analysis["exit_points"]: - template = exit_point.get("template", {}) - for field_config in template.values(): - if field_config.get("type") in ["str", "text", "string"]: - analysis["output_types"].add("text") - elif field_config.get("type") in ["int", "float", "number"]: - analysis["output_types"].add("numeric") - elif field_config.get("type") in ["file", "path"]: - analysis["output_types"].add("file") - - except (KeyError, AttributeError): - # If analysis fails, provide basic info - analysis["components"] = [{"type": "Unknown", "name": "Graph Component"}] - analysis["input_types"] = {"text"} - analysis["output_types"] = {"text"} - - # Convert sets to lists for JSON serialization - analysis["input_types"] = list(analysis["input_types"]) - analysis["output_types"] = list(analysis["output_types"]) - - return analysis - - -def _generate_dynamic_run_description(graph: Graph) -> str: - """Generate dynamic description for the /run endpoint based on graph analysis. - - Args: - graph: The LFX graph - - Returns: - str: Dynamic description for the /run endpoint - """ - analysis = _analyze_graph_structure(graph) - - # Determine input examples based on entry points - input_examples = [] - for entry in analysis["entry_points"]: - template = entry.get("template", {}) - for field_name, field_config in template.items(): - if field_config.get("type") in ["str", "text", "string"]: - input_examples.append(f'"{field_name}": "Your input text here"') - elif field_config.get("type") in ["int", "float", "number"]: - input_examples.append(f'"{field_name}": 42') - elif field_config.get("type") in ["file", "path"]: - input_examples.append(f'"{field_name}": "/path/to/file.txt"') - - if not input_examples: - input_examples = ['"input_value": "Your input text here"'] - - # Determine output examples based on exit points - output_examples = [] - for exit_point in analysis["exit_points"]: - template = exit_point.get("template", {}) - for field_name, field_config in template.items(): - if field_config.get("type") in ["str", "text", "string"]: - output_examples.append(f'"{field_name}": "Processed result"') - elif field_config.get("type") in ["int", "float", "number"]: - output_examples.append(f'"{field_name}": 123') - elif field_config.get("type") in ["file", "path"]: - output_examples.append(f'"{field_name}": "/path/to/output.txt"') - - if not output_examples: - output_examples = ['"result": "Processed result"'] - - description_parts = [ - f"Execute the deployed LFX graph with {analysis['node_count']} components.", - "", - "**Graph Analysis**:", - f"- Entry points: {len(analysis['entry_points'])}", - f"- Exit points: {len(analysis['exit_points'])}", - f"- Input types: {', '.join(analysis['input_types']) if analysis['input_types'] else 'text'}", - f"- Output types: {', '.join(analysis['output_types']) if analysis['output_types'] else 'text'}", - "", - "**Authentication Required**: Include your API key in the `x-api-key` header or as a query parameter.", - "", - "**Example Request**:", - "```json", - "{", - f" {', '.join(input_examples)}", - "}", - "```", - "", - "**Example Response**:", - "```json", - "{", - f" {', '.join(output_examples)},", - ' "success": true,', - ' "logs": "Graph execution completed successfully",', - ' "type": "message",', - ' "component": "FinalComponent"', - "}", - "```", - ] - - return "\n".join(description_parts) +class FlowAlreadyRegisteredError(ValueError): + """Raised by FlowRegistry.add() when a flow ID is already registered and overwrite=False.""" class FlowMeta(BaseModel): """Metadata returned by the ``/flows`` endpoint.""" - id: str = Field(..., description="Deterministic flow identifier (UUIDv5)") + id: str = Field(..., description="Flow identifier (UUID)") relative_path: str = Field(..., description="Path of the flow JSON relative to the deployed folder") title: str = Field(..., description="Human-readable title (filename stem if unknown)") description: str | None = Field(None, description="Optional flow description") @@ -227,6 +95,11 @@ class RunRequest(BaseModel): input_value: str = Field(..., description="Input value passed to the flow") session_id: str | None = Field(default=None, description="Session ID for maintaining conversation state") + global_vars: dict[str, str] | None = Field( + default=None, + description="Per-request variables injected into graph.context['request_variables'] on the deepcopy. " + "Use this to supply credentials or other scoped values without touching os.environ.", + ) class StreamRequest(BaseModel): @@ -238,6 +111,11 @@ class StreamRequest(BaseModel): output_component: str | None = Field(default=None, description="Specific output component to stream from") session_id: str | None = Field(default=None, description="Session ID for maintaining conversation state") tweaks: dict[str, Any] | None = Field(default=None, description="Optional tweaks to modify flow behavior") + global_vars: dict[str, str] | None = Field( + default=None, + description="Per-request variables injected into graph.context['request_variables'] on the deepcopy. " + "Use this to supply credentials or other scoped values without touching os.environ.", + ) class RunResponse(BaseModel): @@ -255,6 +133,317 @@ class ErrorResponse(BaseModel): success: bool = Field(default=False, description="Always false for errors") +class FlowRegistry: + """Mutable in-process registry of loaded flows. + + The in-memory dict is a per-worker cache. Backing persistence is provided + by a :class:`~lfx.cli.flow_store.FlowStore`; the default ``NullFlowStore`` + keeps everything in memory only. + + - ``NullFlowStore`` (default): single-worker, no disk I/O. + - ``FilesystemFlowStore("/tmp/lfx-flows")``: all uvicorn workers in the + same pod share flows via ``/tmp``. + - ``FilesystemFlowStore("/mnt/lfx-flows")``: same code, cross-pod if the + path is a PVC mount. + + **Per-request stale check (FilesystemFlowStore only)** + + Every call to :meth:`get` for a store-backed flow does one + ``FlowStore.read()`` to verify the file still exists. This is how DELETE + propagates across workers: the deleting worker removes the file; the next + inbound request on any other worker calls ``get()``, finds ``None``, evicts + the cached entry, and returns 404. + + On a local SSD this overhead is a single ``stat``-equivalent and is + negligible. On a **network volume** (NFS, CIFS, or a Kubernetes PVC backed + by a networked storage class) each read call crosses the network, which can + add measurable latency per request. If you are running on such a mount and + latency matters, prefer ``/tmp/lfx-flows`` (local tmpfs) for single-pod + worker sharing and accept that cross-pod DELETE propagation is eventual + (next request per pod) rather than immediate. + + When ``no_env_fallback=True``, every graph registered via ``add()`` has + ``graph.context['no_env_fallback']`` set to ``True`` at registration time, + preventing credential resolution from falling back to ``os.environ``. + """ + + # TTL for the cached result of store.list_ids() (a filesystem glob). + # Avoids one glob per /health call while still reflecting changes within ~1 s. + # Mutations (add / remove) invalidate the cache immediately. + _STORE_IDS_TTL: float = 1.0 + + def __init__(self, *, no_env_fallback: bool = False, store: FlowStore | None = None) -> None: + from lfx.cli.flow_store import NullFlowStore + + # Key invariant: a flow may be stored under TWO keys simultaneously — + # its JSON UUID *and* a filename stem (e.g. "prompt_one") when the + # pre-placed file's stem differs from the UUID in the JSON. Both keys + # map to the *same* (graph, meta) tuple. All deduplication logic + # (list_metas, __len__, remove) uses meta.id as the canonical identity. + self._flows: dict[str, tuple[Graph, FlowMeta]] = {} + self._no_env_fallback = no_env_fallback + self._store = store if store is not None else NullFlowStore() + # Maps meta.id → store key when they differ (pre-placed files with human-readable names). + self._store_keys: dict[str, str] = {} + # Lightweight metadata cache for store flows not yet fully loaded into _flows. + # Avoids re-reading JSON on every list_metas() call for multi-worker uploads. + self._store_meta_cache: dict[str, FlowMeta] = {} + # Flow IDs whose source of truth is the store. These are re-verified against + # the store in get() / list_metas() to catch cross-worker deletes. + self._store_sourced: set[str] = set() + # TTL cache for store.list_ids() to avoid one filesystem glob per /health call. + self._store_ids_cache: list[str] | None = None + self._store_ids_cache_ts: float = 0.0 + + def stamp(self, graph: Graph) -> None: + """Apply the registry's env-fallback policy to ``graph.context``. + + Called again after ``deepcopy`` in the run/stream endpoints, since + ``Graph.__deepcopy__`` does not carry ``context`` over. + """ + if self._no_env_fallback: + graph.context["no_env_fallback"] = True + + def _get_cached_store_ids(self) -> list[str]: + """Return store.list_ids(), refreshing at most once per _STORE_IDS_TTL seconds.""" + now = time.monotonic() + if self._store_ids_cache is None or now - self._store_ids_cache_ts > self._STORE_IDS_TTL: + self._store_ids_cache = self._store.list_ids() + self._store_ids_cache_ts = now + return self._store_ids_cache + + def _invalidate_store_ids_cache(self) -> None: + self._store_ids_cache = None + + def add(self, graph: Graph, meta: FlowMeta, *, overwrite: bool = False, raw_json: dict | None = None) -> None: + if not overwrite and meta.id in self._flows: + msg = f"Flow '{meta.id}' is already registered. Pass overwrite=True to replace it." + raise FlowAlreadyRegisteredError(msg) + # If overwriting a flow that was loaded from a differently-named store file + # (e.g. prompt_one.json whose JSON id is a UUID), delete the old file and + # clear the alias so the new file becomes the single source of truth. + if overwrite: + old_store_key = self._store_keys.pop(meta.id, None) + if old_store_key is not None: + self._store.delete(old_store_key) + self._flows.pop(old_store_key, None) + else: + # No alias recorded — the replace path skips the get() that would have + # learned it. A pre-placed file (e.g. my-flow.json) may still carry meta.id + # in its JSON "id" under a differently-named key; delete it so the new + # {uuid}.json becomes the single source of truth. + self._delete_aliased_store_files(meta.id) + if raw_json is not None: + self._store.write(meta.id, raw_json) + if getattr(self._store, "is_persistent", False): + self._store_sourced.add(meta.id) + self.stamp(graph) + self._flows[meta.id] = (graph, meta) + self._store_meta_cache.pop(meta.id, None) + self._invalidate_store_ids_cache() + + def _evict(self, meta_id: str) -> None: + """Remove all in-memory traces of a store-sourced flow (e.g. deleted by another worker).""" + store_key = self._store_keys.pop(meta_id, meta_id) + for k in {meta_id, store_key}: + self._flows.pop(k, None) + self._store_meta_cache.pop(k, None) + self._store_sourced.discard(meta_id) + + def _delete_aliased_store_files(self, canonical_id: str) -> None: + """Delete stem-keyed store files that are aliases of *canonical_id*. + + Targets persistent-store files whose JSON ``id`` equals *canonical_id* but are + keyed under a different stem (e.g. a pre-placed ``my-flow.json`` alongside + ``{uuid}.json``), so the canonical key stays the single source of truth. Also + drops any in-memory cache entries for those stems. No-op for non-persistent + stores (NullFlowStore has no other files). Uses a live ``list_ids()`` rather than + the TTL cache so a listing taken right after a delete is current. + """ + if not getattr(self._store, "is_persistent", False): + return + for stem_id in self._store.list_ids(): + if stem_id == canonical_id: + continue + stem_raw = self._store.read(stem_id) + if stem_raw and stem_raw.get("id") == canonical_id: + self._store.delete(stem_id) + self._flows.pop(stem_id, None) + self._store_meta_cache.pop(stem_id, None) + + def get(self, flow_id: str) -> tuple[Graph, FlowMeta] | None: + if flow_id in self._flows: + _, meta = self._flows[flow_id] + # Cross-worker stale check: if this flow came from the store, verify it + # hasn't been deleted by another worker since we cached it. + # Use _store_keys to find the real store key (which may be a filename stem + # rather than the JSON UUID for pre-placed flows). + if meta.id in self._store_sourced: + store_key = self._store_keys.get(meta.id, meta.id) + if self._store.read(store_key) is None: + self._evict(meta.id) + return None + return self._flows[flow_id] + raw_json = self._store.read(flow_id) + if raw_json is None: + return None + graph, meta = self._reconstruct(flow_id, raw_json) + # Cache under the authoritative JSON id so requests by UUID find it. + self._flows[meta.id] = (graph, meta) + self._store_meta_cache.pop(meta.id, None) + if flow_id != meta.id: + # Also keep the filename-stem alias so warm_from_store lookups hit. + self._flows[flow_id] = (graph, meta) + # Remember the store key so remove() can delete the right file. + self._store_keys[meta.id] = flow_id + return graph, meta + + @staticmethod + def _meta_from_raw_json(flow_id: str, raw_json: dict) -> FlowMeta: + actual_id = raw_json.get("id") or flow_id + return FlowMeta( + id=actual_id, + relative_path="", + title=raw_json.get("name", actual_id), + description=raw_json.get("description"), + ) + + def _reconstruct(self, flow_id: str, raw_json: dict) -> tuple[Graph, FlowMeta]: + meta = self._meta_from_raw_json(flow_id, raw_json) + graph = load_flow_from_json(raw_json) + graph.prepare() + graph.flow_id = meta.id + self.stamp(graph) + if getattr(self._store, "is_persistent", False): + self._store_sourced.add(meta.id) + return graph, meta + + def warm_from_store(self) -> None: + """Load all flows from the backing store into the in-memory cache. + + Called once at startup so every worker can serve any flow that was + previously uploaded (by another worker or a prior run) without a + cache-miss penalty on the first request. + """ + # Bypass the TTL cache — startup must always see the current store state. + for flow_id in self._store.list_ids(): + try: + self.get(flow_id) # no-op if already cached; loads from store otherwise + except Exception as exc: # noqa: BLE001 + # One unloadable flow (corrupt JSON, or a flow referencing a component + # not available in this build) must not abort the whole worker's startup — + # skip it so every other flow still serves. + logger.warning("Skipping flow %r during store warm-up: %r", flow_id, exc) + + def list_metas(self) -> list[FlowMeta]: + seen: set[str] = set() + result: list[FlowMeta] = [] + store_ids = set(self._get_cached_store_ids()) + for _, meta in self._flows.values(): + if meta.id in seen: + continue + # Skip flows deleted by another worker (still in our cache but gone from store). + # Use the real store key (may be a filename stem for pre-placed flows). + # Check the store directly rather than the TTL-cached id list, so a delete + # by another worker is reflected immediately (same as the get() stale check). + if meta.id in self._store_sourced: + store_key = self._store_keys.get(meta.id, meta.id) + if not self._store.exists(store_key): + continue + result.append(meta) + seen.add(meta.id) + for flow_id in store_ids: + if flow_id not in self._flows: + # Check the lightweight metadata cache before reading from disk. + if flow_id in self._store_meta_cache: + cached = self._store_meta_cache[flow_id] + if cached.id not in seen: + result.append(cached) + seen.add(cached.id) + continue + raw_json = self._store.read(flow_id) + if raw_json: + meta = self._meta_from_raw_json(flow_id, raw_json) + if meta.id not in seen: + self._store_meta_cache[flow_id] = meta + result.append(meta) + seen.add(meta.id) + return result + + def remove(self, flow_id: str) -> bool: + # Resolve the canonical meta id and store key. + entry = self._flows.get(flow_id) + if entry is not None: + meta_id = entry[1].id + store_key = self._store_keys.get(meta_id, meta_id) + # Drop all cache keys that point to this flow (UUID and any filename alias). + for k in [meta_id, store_key]: + self._flows.pop(k, None) + self._store_meta_cache.pop(k, None) + self._store_keys.pop(meta_id, None) + self._store_sourced.discard(meta_id) + mem_had_it = True + else: + # Flow not in memory; derive both the stem key (flow_id) and the UUID from the file. + store_key = flow_id + self._store_meta_cache.pop(flow_id, None) + self._store_sourced.discard(flow_id) + mem_had_it = False + meta_id = flow_id # best-effort default; overridden below if we can read the file + raw = self._store.read(flow_id) + if raw: + meta_id = raw.get("id") or flow_id + + store_had_it = self._store.delete(store_key) + # Ensure cross-worker DELETE propagation: delete BOTH the primary store key + # and any alternate-keyed file so that workers whose in-memory stale-check + # uses a different key also see the deletion on their next request. + if store_key != meta_id: + # Alias known: primary key was a stem, UUID-keyed file may also exist. + self._store.delete(meta_id) + elif store_had_it: + # Primary key IS the UUID: delete any stem-keyed file carrying the same + # "id" field (e.g. a pre-placed my-flow.json alongside {uuid}.json). + self._delete_aliased_store_files(meta_id) + self._invalidate_store_ids_cache() + return mem_had_it or store_had_it + + def __len__(self) -> int: + # Delegates to list_metas() so store-only flows (not yet cache-loaded) are + # counted. list_metas() caches metadata after the first disk read, so + # repeated calls (e.g. /health) are cheap after the first. + return len(self.list_metas()) + + +class UploadFlowRequest(BaseModel): + model_config = ConfigDict(extra="allow") + + name: str = Field(..., description="Human-readable name for the flow") + data: dict = Field(..., description="Flow graph data — nodes and edges") + description: str | None = Field(default=None, description="Optional flow description") + id: str | None = Field(default=None, description="Stable flow ID from Langflow export (must be a valid UUID)") + replace: bool = Field(default=False, description="Overwrite the existing flow if the ID already exists") + + @field_validator("id", mode="before") + @classmethod + def validate_id_is_uuid(cls, v: object) -> object: + if v is None: + return v + try: + uuid.UUID(str(v)) + except ValueError: + msg = f"id must be a valid UUID, got {v!r}" + raise ValueError(msg) from None + return str(v) + + +class UploadFlowResponse(BaseModel): + id: str = Field(..., description="Flow identifier (UUID)") + name: str + description: str | None + run_url: str = Field(..., description="Endpoint to POST run requests, e.g. /flows/{id}/run") + + # ----------------------------------------------------------------------------- # Streaming helper functions # ----------------------------------------------------------------------------- @@ -352,204 +541,310 @@ async def run_flow_generator_for_serve( def create_multi_serve_app( *, - root_dir: Path, # noqa: ARG001 - graphs: dict[str, Graph], - metas: dict[str, FlowMeta], - verbose_print: Callable[[str], None], # noqa: ARG001 + registry: FlowRegistry, ) -> FastAPI: - """Create a FastAPI app exposing multiple LFX flows. + """Create a FastAPI app exposing LFX flows via a mutable registry. - Parameters - ---------- - root_dir - Folder originally supplied to the serve command. All *relative_path* - values are relative to this directory. - graphs - Mapping ``flow_id -> Graph`` containing prepared graph objects. - metas - Mapping ``flow_id -> FlowMeta`` containing metadata for each flow. - verbose_print - Diagnostic printer inherited from the CLI (unused, kept for backward compatibility). + Routes dispatch to ``registry`` at request time, so flows added after + startup (via ``POST /flows/upload/``) are immediately reachable. """ - if set(graphs) != set(metas): # pragma: no cover - sanity check - msg = "graphs and metas must contain the same keys" - raise ValueError(msg) - app = FastAPI( - title=f"LFX Multi-Flow Server ({len(graphs)})", + title=f"LFX Multi-Flow Server ({len(registry)})", description=( - "This server hosts multiple LFX graphs under the `/flows/{id}` prefix. " - "Use `/flows` to list available IDs then POST your input to `/flows/{id}/run`." + "Hosts LFX graphs under the `/flows/{{id}}` prefix. " + "Use `/flows` to list available IDs then POST your input to `/flows/{{id}}/run`. " + "Use `POST /flows/upload/` to register new flows at runtime." ), version="1.0.0", ) + app.state.registry = registry # ------------------------------------------------------------------ # Global endpoints # ------------------------------------------------------------------ - @app.get("/flows", response_model=list[FlowMeta], tags=["info"], summary="List available flows") + @app.get( + "/flows", + response_model=list[FlowMeta], + tags=["info"], + summary="List available flows", + dependencies=[Depends(verify_api_key)], + ) async def list_flows(): - """Return metadata for all flows hosted in this server.""" - return list(metas.values()) + return registry.list_metas() @app.get("/health", tags=["info"], summary="Global health check") async def global_health(): - return {"status": "healthy", "flow_count": len(graphs)} + return {"status": "healthy", "flow_count": len(registry)} # ------------------------------------------------------------------ - # Per-flow routers + # Upload endpoint — registered BEFORE /{flow_id} to avoid shadowing # ------------------------------------------------------------------ - def create_flow_router(flow_id: str, graph: Graph, meta: FlowMeta) -> APIRouter: - """Create a router for a specific flow to avoid loop variable binding issues.""" - analysis = _analyze_graph_structure(graph) - run_description = _generate_dynamic_run_description(graph) + @app.post( + "/flows/upload/", + response_model=UploadFlowResponse, + status_code=201, + tags=["upload"], + summary="Upload and register a new flow", + dependencies=[Depends(verify_api_key)], + ) + async def upload_flow(body: UploadFlowRequest) -> UploadFlowResponse: + # Conflict check before the expensive load+prepare — only possible when the + # caller supplies an explicit id. uuid4()-generated ids are always unique so + # there is no point checking before we have one. + flow_id = body.id or str(uuid.uuid4()) - router = APIRouter( - prefix=f"/flows/{flow_id}", - tags=[meta.title or flow_id], - dependencies=[Depends(verify_api_key)], # Auth for all routes inside + if not body.replace and registry.get(flow_id) is not None: + raise HTTPException( + status_code=409, + detail=f"Flow '{flow_id}' already exists. Pass replace=true to overwrite.", + ) + + try: + graph = load_flow_from_json(body.model_dump(exclude={"replace"})) + except Exception as exc: + raise HTTPException(status_code=422, detail=f"Invalid flow data: {exc}") from exc + + try: + graph.prepare() + except Exception as exc: + raise HTTPException(status_code=422, detail=f"Flow preparation failed: {exc}") from exc + + graph.flow_id = flow_id + meta = FlowMeta( + id=flow_id, + relative_path="", + title=body.name, + description=body.description, + ) + # graph.prepare() must run before registry.add() — add() stamps + # graph.context with no_env_fallback, and prepare() must not overwrite it. + registry.add(graph, meta, overwrite=body.replace, raw_json=body.model_dump(exclude={"replace"})) + return UploadFlowResponse( + id=flow_id, + name=body.name, + description=body.description, + run_url=f"/flows/{flow_id}/run", ) - @router.post( - "/run", - response_model=RunResponse, - responses={500: {"model": ErrorResponse}}, - summary="Execute flow", - description=run_description, - ) - async def run_flow( - request: RunRequest, - ) -> RunResponse: - try: - validate_flow_for_current_settings(graph) - graph_copy = deepcopy(graph) - results, logs = await execute_graph_with_capture( - graph_copy, request.input_value, session_id=request.session_id - ) - result_data = extract_result_data(results, logs) + # ------------------------------------------------------------------ + # Per-flow dispatch routes + # ------------------------------------------------------------------ - # Debug logging - logger.debug(f"Flow {flow_id} execution completed: {len(results)} results, {len(logs)} log chars") - logger.debug(f"Flow {flow_id} result data: {result_data}") + def _get_flow_or_404(flow_id: str) -> tuple[Graph, FlowMeta]: + result = registry.get(flow_id) + if result is None: + raise HTTPException( + status_code=404, + detail={"error": "flow not found", "flow_id": flow_id}, + ) + return result - # Check if the execution was successful - if not result_data.get("success", True): - # If the flow execution failed, return error details in the response - error_message = result_data.get("result", result_data.get("text", "No response generated")) + @app.delete( + "/flows/{flow_id}", + status_code=204, + tags=["flows"], + summary="Remove a registered flow", + description=( + "Remove a flow from this worker's registry and from the backing store (if any). " + "**Multi-worker note:** the store file is deleted immediately, but other workers " + "continue to serve the cached copy until their next stale check " + "(triggered by an incoming request for that flow). " + "This is the same eventual-consistency window that applies to uploads." + ), + dependencies=[Depends(verify_api_key)], + ) + async def delete_flow(flow_id: str) -> Response: + removed = registry.remove(flow_id) + if not removed: + raise HTTPException( + status_code=404, + detail={"error": "flow not found", "flow_id": flow_id}, + ) + return Response(status_code=204) - # Add more context to the logs when there's an error - error_logs = logs - if not error_logs.strip(): - error_logs = ( - f"Flow execution completed but no valid result was produced.\nResult data: {result_data}" - ) + @app.get( + "/flows/{flow_id}/info", + response_model=FlowMeta, + tags=["flows"], + summary="Flow metadata", + dependencies=[Depends(verify_api_key)], + ) + async def flow_info(flow_id: str) -> FlowMeta: + _, meta = _get_flow_or_404(flow_id) + return meta - return RunResponse( + @app.post( + "/flows/{flow_id}/run", + response_model=RunResponse, + responses={404: {"model": ErrorResponse}, 500: {"model": RunResponse}}, + tags=["flows"], + summary="Execute flow", + dependencies=[Depends(verify_api_key)], + ) + async def run_flow(flow_id: str, request: RunRequest) -> RunResponse: + graph, _ = _get_flow_or_404(flow_id) + try: + validate_flow_for_current_settings(graph) + graph_copy = deepcopy(graph) + # deepcopy() drops graph.context; re-apply the registry's env policy. + 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 + ) + result_data = extract_result_data(results, logs) + + if not result_data.get("success", True): + error_message = result_data.get("result", result_data.get("text", "No response generated")) + return JSONResponse( + status_code=500, + content=RunResponse( result=error_message, success=False, - logs=error_logs, + logs=logs + or f"Flow execution completed but no valid result was produced.\nResult data: {result_data}", type="error", component=result_data.get("component", ""), - ) - - return RunResponse( - result=result_data.get("result", result_data.get("text", "")), - success=result_data.get("success", True), - logs=logs, - type=result_data.get("type", "message"), - component=result_data.get("component", ""), + ).model_dump(), ) - except Exception as exc: # noqa: BLE001 - import traceback - - # Capture the full traceback for debugging - error_traceback = traceback.format_exc() - error_message = f"Flow execution failed: {exc!s}" - - # Log to server console for debugging - logger.error(f"Error running flow {flow_id}: {exc}") - logger.debug(f"Full traceback for flow {flow_id}:\n{error_traceback}") - - # Return error details in the API response instead of raising HTTPException - return RunResponse( + return RunResponse( + result=result_data.get("result", result_data.get("text", "")), + success=result_data.get("success", True), + logs=logs, + type=result_data.get("type", "message"), + component=result_data.get("component", ""), + ) + except Exception as exc: # noqa: BLE001 + error_traceback = traceback.format_exc() + error_message = f"Flow execution failed: {exc!s}" + logger.error(f"Error running flow {flow_id}: {exc}") + logger.debug(f"Full traceback for flow {flow_id}:\n{error_traceback}") + return JSONResponse( + status_code=500, + content=RunResponse( result=error_message, success=False, logs=f"ERROR: {error_message}\n\nFull traceback:\n{error_traceback}", type="error", component="", + ).model_dump(), + ) + + @app.post( + "/flows/{flow_id}/stream", + response_model=None, + tags=["flows"], + summary="Stream flow execution", + dependencies=[Depends(verify_api_key)], + ) + async def stream_flow(flow_id: str, request: StreamRequest) -> StreamingResponse: + graph, _ = _get_flow_or_404(flow_id) + try: + validate_flow_for_current_settings(graph) + from lfx.events.event_manager import create_stream_tokens_event_manager + + asyncio_queue: asyncio.Queue = asyncio.Queue() + asyncio_queue_client_consumed: asyncio.Queue = asyncio.Queue() + event_manager = create_stream_tokens_event_manager(queue=asyncio_queue) + + graph_copy = deepcopy(graph) + # deepcopy() drops graph.context; re-apply the registry's env policy. + registry.stamp(graph_copy) + apply_global_vars_to_graph(graph_copy, request.global_vars) + main_task = asyncio.create_task( + run_flow_generator_for_serve( + graph=graph_copy, + input_request=request, + flow_id=flow_id, + event_manager=event_manager, + client_consumed_queue=asyncio_queue_client_consumed, ) + ) - @router.post( - "/stream", - response_model=None, - summary="Stream flow execution", - description=f"Stream the execution of {meta.title or flow_id} with real-time events and token streaming.", - ) - async def stream_flow( - request: StreamRequest, - ) -> StreamingResponse: - """Stream the execution of the flow with real-time events.""" - try: - validate_flow_for_current_settings(graph) + async def on_disconnect() -> None: + logger.debug(f"Client disconnected from flow {flow_id}, closing tasks") + main_task.cancel() - # Import here to avoid potential circular imports - from lfx.events.event_manager import create_stream_tokens_event_manager + return StreamingResponse( + consume_and_yield(asyncio_queue, asyncio_queue_client_consumed), + background=on_disconnect, + media_type="text/event-stream", + ) + except Exception as exc: # noqa: BLE001 + logger.error(f"Error setting up streaming for flow {flow_id}: {exc}") + error_payload = json.dumps({"error": str(exc), "success": False}) - asyncio_queue: asyncio.Queue = asyncio.Queue() - asyncio_queue_client_consumed: asyncio.Queue = asyncio.Queue() - event_manager = create_stream_tokens_event_manager(queue=asyncio_queue) + async def error_stream(): + yield f"data: {error_payload}\n\n" - main_task = asyncio.create_task( - run_flow_generator_for_serve( - graph=graph, - input_request=request, - flow_id=flow_id, - event_manager=event_manager, - client_consumed_queue=asyncio_queue_client_consumed, - ) - ) - - async def on_disconnect() -> None: - logger.debug(f"Client disconnected from flow {flow_id}, closing tasks") - main_task.cancel() - - return StreamingResponse( - consume_and_yield(asyncio_queue, asyncio_queue_client_consumed), - background=on_disconnect, - media_type="text/event-stream", - ) - except Exception as exc: # noqa: BLE001 - logger.error(f"Error setting up streaming for flow {flow_id}: {exc}") - # Return a simple error stream - error_message = f"Failed to start streaming: {exc!s}" - - async def error_stream(): - yield f'data: {{"error": "{error_message}", "success": false}}\n\n' - - return StreamingResponse( - error_stream(), - media_type="text/event-stream", - ) - - @router.get("/info", summary="Flow metadata", response_model=FlowMeta) - async def flow_info(): - """Return metadata and basic analysis for this flow.""" - # Enrich meta with analysis data for convenience - return { - **meta.model_dump(), - "components": analysis["node_count"], - "connections": analysis["edge_count"], - "input_types": analysis["input_types"], - "output_types": analysis["output_types"], - } - - return router - - for flow_id, graph in graphs.items(): - meta = metas[flow_id] - router = create_flow_router(flow_id, graph, meta) - app.include_router(router) + return StreamingResponse(error_stream(), media_type="text/event-stream") return app + + +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. + + The parent process must set those env vars **before** calling + ``uvicorn.run("lfx.cli.serve_app:create_serve_app", workers=N, ...)``. + """ + import asyncio + import os + from pathlib import Path + + from lfx.cli.flow_store import FilesystemFlowStore, NullFlowStore + + 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, "") + + flow_dir = Path(flow_dir_str) if flow_dir_str else None + flow_store = FilesystemFlowStore(flow_dir) if flow_dir else NullFlowStore() + + startup_paths = [Path(p) for p in json.loads(startup_paths_json)] if startup_paths_json else [] + + if startup_paths and not flow_dir: + # No shared store: each worker reloads startup flows from the original file paths. + # When flow_dir IS set the parent already persisted startup JSON flows to the store; + # workers pick them up via warm_from_store() below — no need to re-read files. + # + # ``create_serve_app`` is called by uvicorn as an ASGI app factory while an + # event loop is already running in the worker process. ``asyncio.run()`` + # raises RuntimeError in that situation. Running the coroutine in a fresh + # thread gives it a clean event loop with no interference. + import concurrent.futures + + from lfx.cli.commands import build_registry_from_directory, build_registry_from_paths + + if len(startup_paths) == 1 and startup_paths[0].is_dir(): + coro = build_registry_from_directory( + startup_paths[0], + lambda _: None, + check_variables=False, + no_env_fallback=no_env_fallback, + store=flow_store, + ) + else: + coro = build_registry_from_paths( + startup_paths, + lambda _: None, + check_variables=False, + no_env_fallback=no_env_fallback, + store=flow_store, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool: + registry = pool.submit(asyncio.run, coro).result() + else: + registry = FlowRegistry(no_env_fallback=no_env_fallback, store=flow_store) + + registry.warm_from_store() + return create_multi_serve_app(registry=registry) diff --git a/src/lfx/src/lfx/components/mem0/mem0_chat_memory.py b/src/lfx/src/lfx/components/mem0/mem0_chat_memory.py index 7dcde99648..25c0231a9a 100644 --- a/src/lfx/src/lfx/components/mem0/mem0_chat_memory.py +++ b/src/lfx/src/lfx/components/mem0/mem0_chat_memory.py @@ -1,5 +1,3 @@ -import os - from mem0 import Memory, MemoryClient from lfx.base.memory.model import LCChatMemoryComponent @@ -87,12 +85,27 @@ class Mem0MemoryComponent(LCChatMemoryComponent): """Initializes a Mem0 memory instance based on provided configuration and API keys.""" # Check if we're in Astra cloud environment and raise an error if we are. raise_error_if_astra_cloud_disable_component(disable_component_in_astra_cloud_msg) + + # Pass the OpenAI key through mem0's config instead of writing os.environ. + # os.environ is process-global and persists across requests, so in a shared + # serving process (e.g. lfx serve) one caller's key would leak into other + # concurrent requests. mem0's OpenAI llm/embedder read config.api_key before + # falling back to OPENAI_API_KEY, so config injection keeps it request-scoped. + config = dict(self.mem0_config) if self.mem0_config else {} if self.openai_api_key: - os.environ["OPENAI_API_KEY"] = self.openai_api_key + for section in ("llm", "embedder"): + sec = dict(config.get(section) or {}) + if sec.get("provider", "openai") != "openai": + continue # the OpenAI key is irrelevant to non-OpenAI providers + sec["provider"] = "openai" + sec_config = dict(sec.get("config") or {}) + sec_config.setdefault("api_key", self.openai_api_key) + sec["config"] = sec_config + config[section] = sec try: if not self.mem0_api_key: - return Memory.from_config(config_dict=dict(self.mem0_config)) if self.mem0_config else Memory() + return Memory.from_config(config_dict=config) if config else Memory() if self.mem0_config: return MemoryClient.from_config(api_key=self.mem0_api_key, config_dict=dict(self.mem0_config)) return MemoryClient(api_key=self.mem0_api_key) diff --git a/src/lfx/src/lfx/custom/custom_component/component.py b/src/lfx/src/lfx/custom/custom_component/component.py index 61c65062ff..d9cd99a027 100644 --- a/src/lfx/src/lfx/custom/custom_component/component.py +++ b/src/lfx/src/lfx/custom/custom_component/component.py @@ -437,8 +437,16 @@ class Component(CustomComponent): kwargs = dict(config) kwargs.update(inputs_raw) new_component = type(self)(**kwargs) + # Register in memo before the recursive deepcopy calls below so reference + # cycles (e.g. components linked through _components) resolve to this same + # copy instead of producing a duplicate. + memo[id(self)] = new_component new_component._code = self._code - new_component._outputs_map = self._outputs_map + # Must deep-copy so each graph_copy has independent Output objects. + # Output.cache=True by default, and output.value is set during execution. + # A shallow copy causes all concurrent requests to share the same Output objects, + # so the first request's cached output.value is returned to all subsequent requests. + new_component._outputs_map = deepcopy(self._outputs_map, memo) # Safe deepcopy of inputs new_inputs = {} @@ -471,13 +479,16 @@ class Component(CustomComponent): new_inputs[k] = input_copy new_component._inputs = new_inputs - new_component._edges = self._edges - new_component._components = self._components + # Must deep-copy so each graph_copy has independent component instances. + # Shallow copies here caused all concurrent requests to share the same + # intermediate component objects (e.g. the LLM node), producing identical + # responses for different inputs under concurrent load. + new_component._edges = deepcopy(self._edges, memo) + new_component._components = deepcopy(self._components, memo) new_component._parameters = dict(self._parameters) new_component._attributes = dict(self._attributes) new_component._output_logs = self._output_logs new_component._logs = self._logs # type: ignore[attr-defined] - memo[id(self)] = new_component return new_component def set_class_code(self) -> None: diff --git a/src/lfx/src/lfx/interface/initialize/loading.py b/src/lfx/src/lfx/interface/initialize/loading.py index 9956ac7c1b..40828b7ca0 100644 --- a/src/lfx/src/lfx/interface/initialize/loading.py +++ b/src/lfx/src/lfx/interface/initialize/loading.py @@ -114,13 +114,14 @@ def convert_kwargs(params): def load_from_env_vars(params, load_from_db_fields, context=None): + no_env_fallback = bool(context and context.get("no_env_fallback")) for field in load_from_db_fields: if field not in params or not params[field]: continue variable_name = params[field] key = None - # Check request_variables in context + # Check request_variables in context first if context and "request_variables" in context: request_variables = context["request_variables"] if variable_name in request_variables: @@ -128,14 +129,20 @@ def load_from_env_vars(params, load_from_db_fields, context=None): logger.debug(f"Found context override for variable '{variable_name}'") if key is None: - key = os.getenv(variable_name) - if key: - logger.info(f"Using environment variable {variable_name} for {field}") + if no_env_fallback: + logger.warning( + f"Variable '{variable_name}' not found in request_variables and " + f"env fallback is disabled. Setting to None." + ) else: - logger.error(f"Environment variable {variable_name} is not set.") - params[field] = key if key is not None else None - if key is None: - logger.warning(f"Could not get value for {field}. Setting it to None.") + key = os.getenv(variable_name) + if key: + logger.info(f"Using environment variable {variable_name} for {field}") + else: + logger.error(f"Environment variable {variable_name} is not set.") + logger.warning(f"Could not get value for {field}. Setting it to None.") + + params[field] = key return params @@ -166,6 +173,9 @@ async def update_table_params_with_load_from_db_fields( context = None if hasattr(custom_component, "graph") and hasattr(custom_component.graph, "context"): context = custom_component.graph.context + # Honor the same no-env-fallback contract as load_from_env_vars so a served flow under + # no_env_fallback never resolves table columns from process-wide os.environ. + no_env_fallback = bool(context and context.get("no_env_fallback")) async with session_scope() as session: settings_service = get_settings_service() @@ -208,7 +218,7 @@ async def update_table_params_with_load_from_db_fields( key = request_variables[variable_name] logger.debug(f"Found context override for variable '{variable_name}'") - if key is None: + if key is None and not no_env_fallback: key = os.getenv(variable_name) if key: logger.info( @@ -229,7 +239,7 @@ async def update_table_params_with_load_from_db_fields( key = None # If we couldn't get from database and fallback is enabled, try environment - if fallback_to_env_vars and key is None: + if fallback_to_env_vars and key is None and not no_env_fallback: key = os.getenv(variable_name) if key: logger.info(f"Using environment variable {variable_name} for table column {column_name}") diff --git a/src/lfx/src/lfx/services/variable/request_scope.py b/src/lfx/src/lfx/services/variable/request_scope.py new file mode 100644 index 0000000000..f2cdf81715 --- /dev/null +++ b/src/lfx/src/lfx/services/variable/request_scope.py @@ -0,0 +1,75 @@ +"""Request-scoped variables for lfx serve (and other in-process runners). + +TRM/WXO inject credentials via ``global_vars`` on ``POST /flows/{id}/run`` without +mutating process-wide ``os.environ``. This module holds the active request's flat +lookup table in a ContextVar so :class:`~lfx.services.variable.service.VariableService` +can resolve the same names as the ``lfx run`` subprocess path. + +A second ContextVar mirrors ``graph.context['no_env_fallback']`` so the service can +honor the same "do not read ``os.environ``" contract that ``load_from_env_vars`` +enforces for ``load_from_db`` fields — keeping the two resolution paths consistent. +""" + +from __future__ import annotations + +import contextvars +import json + +_request_variables: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar( + "lfx_request_variables", + default=None, +) +_no_env_fallback: contextvars.ContextVar[bool] = contextvars.ContextVar( + "lfx_no_env_fallback", + default=False, +) + + +def normalize_parsed_variables(parsed: dict) -> dict[str, str]: + """Flatten a parsed JSON variable map to ``str`` values. + + Drops ``None`` (so a JSON ``null`` never becomes the truthy string "None" and + masquerades as a credential) and serializes dict/list values as valid JSON + (round-trippable via ``json.loads``) rather than a lossy Python repr. Scalars use ``str``. + """ + return { + str(key): json.dumps(value) if isinstance(value, dict | list) else str(value) + for key, value in parsed.items() + if value is not None + } + + +def get_active_request_variables() -> dict[str, str] | None: + """Return the current request's variable map, if any.""" + return _request_variables.get() + + +def is_env_fallback_disabled() -> bool: + """Return whether the current request forbids falling back to ``os.environ``.""" + return _no_env_fallback.get() + + +def activate_no_env_fallback(*, disabled: bool) -> contextvars.Token[bool]: + """Bind the no-env-fallback flag for the current async task / thread.""" + return _no_env_fallback.set(disabled) + + +def reset_no_env_fallback(token: contextvars.Token[bool]) -> None: + """Restore the previous no-env-fallback flag.""" + _no_env_fallback.reset(token) + + +def activate_request_variables(variables: dict[str, str] | None) -> contextvars.Token[dict[str, str] | None]: + """Bind *variables* for the current async task / thread. + + Pass ``None`` to mean "no active scope" — lookups then fall back to the + ``LANGFLOW_REQUEST_VARIABLES`` env var. An empty ``{}`` is an *active* empty + scope that suppresses that env fallback, so callers normalizing an empty + result should pass ``... or None`` to preserve it. + """ + return _request_variables.set(variables) + + +def reset_request_variables(token: contextvars.Token[dict[str, str] | None]) -> None: + """Restore the previous request scope.""" + _request_variables.reset(token) diff --git a/src/lfx/src/lfx/services/variable/service.py b/src/lfx/src/lfx/services/variable/service.py index 6121f2999d..6ff4b94f32 100644 --- a/src/lfx/src/lfx/services/variable/service.py +++ b/src/lfx/src/lfx/services/variable/service.py @@ -5,6 +5,11 @@ import os from lfx.log.logger import logger from lfx.services.base import Service +from lfx.services.variable.request_scope import ( + get_active_request_variables, + is_env_fallback_disabled, + normalize_parsed_variables, +) class VariableService(Service): @@ -30,7 +35,17 @@ class VariableService(Service): @staticmethod def _get_request_variables() -> dict[str, str]: - """Parse request-scoped variables from LANGFLOW_REQUEST_VARIABLES when available.""" + """Request-scoped variables from serve ContextVar or LANGFLOW_REQUEST_VARIABLES env.""" + active = get_active_request_variables() + if active is not None: + return active + + # LANGFLOW_REQUEST_VARIABLES is a process-wide env var, so reading it is an + # os.environ access. Honor the no-env-fallback contract and skip it when the + # request disables env fallback, keeping the "never reads os.environ" guarantee. + if is_env_fallback_disabled(): + return {} + raw = os.getenv("LANGFLOW_REQUEST_VARIABLES") if not raw: return {} @@ -42,86 +57,68 @@ class VariableService(Service): if not isinstance(parsed, dict): logger.debug("LANGFLOW_REQUEST_VARIABLES must be a JSON object; skipping request-scoped lookup") return {} - return {str(key): str(value) for key, value in parsed.items()} + return normalize_parsed_variables(parsed) async def get_variable(self, name: str, **kwargs) -> str | None: # noqa: ARG002 """Get a variable value. - First checks in-memory cache, then LANGFLOW_REQUEST_VARIABLES, then - environment variables, then ``x-langflow-global-var-*`` normalized keys. + Resolution order (first match wins): + 1. In-memory cache (``set_variable``). + 2. Request-scoped exact name (serve ContextVar or ``LANGFLOW_REQUEST_VARIABLES``). + 3. Request-scoped ``x-langflow-global-var-*`` alias. + 4. Environment variable (exact name). + 5. Environment variable (``x-langflow-global-var-*`` alias). - Async to match the call signature in custom_component.get_variable - (`await variable_service.get_variable(...)`), which is the path used - by component variable resolution. The lookup itself is sync — no I/O — - but the coroutine wrapper is required so callers can `await` it - regardless of which variable service implementation is registered - (lfx env-fallback vs langflow DB-backed). + Both request-scoped lookups (2, 3) run *before* any process-env fallback (4, 5), + so a credential the caller supplies for this request — in either the exact-name + or the ``x-langflow-global-var-*`` form — always beats a same-named variable left + in the worker's environment, preventing one caller's request from running on an + ambient process credential. - Args: - name: Variable name - **kwargs: Additional arguments (ignored; user_id/field/session - from langflow's call signature are absorbed and not used, - since this implementation has no per-user scope). - - Returns: - Variable value or None if not found + When the active request disables env fallback (``graph.context['no_env_fallback']``, + propagated via the request scope), steps 4 and 5 are skipped so resolution + never reads ``os.environ`` — matching ``load_from_env_vars`` for ``load_from_db`` + fields and keeping served flows isolated from process-wide credentials. """ - # Check in-memory first if name in self._variables: return self._variables[name] - # Contract-first: prefer request-scoped variables injected by runtime. request_variables = self._get_request_variables() if name in request_variables: - logger.debug(f"Variable '{name}' loaded from LANGFLOW_REQUEST_VARIABLES") + logger.debug(f"Variable '{name}' loaded from request-scoped variables") return request_variables[name] - # Fall back to environment variable - value = os.getenv(name) - if value: - logger.debug(f"Variable '{name}' loaded from environment") - return value - - # Fall back to x-langflow-global-var-* aliases global_alias = self._normalize_global_var_key(name) - value = os.getenv(global_alias) - if value: - logger.debug(f"Variable '{name}' loaded from global alias '{global_alias}'") - return value + if global_alias in request_variables: + logger.debug(f"Variable '{name}' loaded from request-scoped alias '{global_alias}'") + return request_variables[global_alias] + + if not is_env_fallback_disabled(): + value = os.getenv(name) + if value: + logger.debug(f"Variable '{name}' loaded from environment") + return value + + value = os.getenv(global_alias) + if value: + logger.debug(f"Variable '{name}' loaded from global alias '{global_alias}'") + return value return None def set_variable(self, name: str, value: str, **kwargs) -> None: # noqa: ARG002 - """Set a variable value (in-memory only). - - Args: - name: Variable name - value: Variable value - **kwargs: Additional arguments (ignored in minimal implementation) - """ + """Set a variable value (in-memory only).""" self._variables[name] = value logger.debug(f"Variable '{name}' set (in-memory only)") def delete_variable(self, name: str, **kwargs) -> None: # noqa: ARG002 - """Delete a variable (from in-memory cache only). - - Args: - name: Variable name - **kwargs: Additional arguments (ignored in minimal implementation) - """ + """Delete a variable (from in-memory cache only).""" if name in self._variables: del self._variables[name] logger.debug(f"Variable '{name}' deleted (from in-memory cache)") def list_variables(self, **kwargs) -> list[str]: # noqa: ARG002 - """List all variables (in-memory only). - - Args: - **kwargs: Additional arguments (ignored in minimal implementation) - - Returns: - List of variable names - """ + """List all variables (in-memory only).""" return list(self._variables.keys()) async def teardown(self) -> None: diff --git a/src/lfx/src/lfx/utils/async_helpers.py b/src/lfx/src/lfx/utils/async_helpers.py index 5a6513a743..bb4f8631f6 100644 --- a/src/lfx/src/lfx/utils/async_helpers.py +++ b/src/lfx/src/lfx/utils/async_helpers.py @@ -25,9 +25,15 @@ def run_until_complete(coro): except RuntimeError: # If there's no event loop, create a new one and run the coroutine return asyncio.run(coro) - # If there's already a running event loop, we can't call run_until_complete on it - # Instead, we need to run the coroutine in a new thread with a new event loop + # If there's already a running event loop, we can't call run_until_complete on it. + # Instead, run the coroutine in a new thread with a new event loop. Propagate the + # caller's contextvars context so request-scoped state (e.g. lfx serve request + # variables and the no-env-fallback flag) stays visible inside the worker thread; + # a bare ThreadPoolExecutor would otherwise reset every ContextVar to its default. import concurrent.futures + import contextvars + + ctx = contextvars.copy_context() def run_in_new_loop(): new_loop = asyncio.new_event_loop() @@ -38,5 +44,5 @@ def run_until_complete(coro): new_loop.close() with concurrent.futures.ThreadPoolExecutor() as executor: - future = executor.submit(run_in_new_loop) + future = executor.submit(ctx.run, run_in_new_loop) return future.result() diff --git a/src/lfx/tests/unit/base/models/test_credentials.py b/src/lfx/tests/unit/base/models/test_credentials.py new file mode 100644 index 0000000000..31aac4e3e3 --- /dev/null +++ b/src/lfx/tests/unit/base/models/test_credentials.py @@ -0,0 +1,216 @@ +"""Behavioral tests for the unified-model credential resolver. + +Covers ``get_api_key_for_provider`` and ``get_all_variables_for_provider`` with a +focus on the no-env-fallback contract: when a served flow disables env fallback, +the resolver must not read process-wide ``os.environ`` credentials, while +request-scoped variables (injected via lfx serve ``global_vars``) must still +resolve through VariableService. + +No behavior is mocked. Tests that exercise the VariableService path register the +real minimal ``VariableService`` on the global service manager; environment is +managed with pytest's ``monkeypatch``. The service-path tests are ``async`` so the +call runs inside a live event loop, exercising the ``run_until_complete`` +worker-thread hop (and its contextvars propagation) end to end. +""" + +from __future__ import annotations + +from contextlib import contextmanager + +import pytest +from lfx.base.models.unified_models import ( + get_all_variables_for_provider, + get_api_key_for_provider, + get_model_provider_variable_mapping, +) +from lfx.services.variable.request_scope import ( + activate_no_env_fallback, + activate_request_variables, + reset_no_env_fallback, + reset_request_variables, +) + +_PROVIDER = "OpenAI" +_USER_ID = "11111111-1111-1111-1111-111111111111" +_MISSING = object() + + +def _primary_var_name() -> str: + mapping = get_model_provider_variable_mapping() + assert _PROVIDER in mapping, f"expected {_PROVIDER} in provider mapping: {sorted(mapping)}" + return mapping[_PROVIDER] + + +@contextmanager +def _no_env_fallback(): + token = activate_no_env_fallback(disabled=True) + try: + yield + finally: + reset_no_env_fallback(token) + + +@contextmanager +def _request_scope(variables: dict[str, str] | None): + token = activate_request_variables(variables) + try: + yield + finally: + reset_request_variables(token) + + +@pytest.fixture +def variable_service_registered(): + """Register the real minimal VariableService on the global manager, then restore. + + Standalone lfx has no VariableService factory, so ``get_variable_service()`` + returns None by default. Registering the genuine service (not a mock) lets the + credential resolver exercise the request-scoped lookup path for real. + """ + from lfx.services.manager import get_service_manager + from lfx.services.schema import ServiceType + from lfx.services.variable.service import VariableService + + manager = get_service_manager() + prev_class = manager.service_classes.get(ServiceType.VARIABLE_SERVICE, _MISSING) + prev_instance = manager.services.get(ServiceType.VARIABLE_SERVICE, _MISSING) + + manager.register_service_class(ServiceType.VARIABLE_SERVICE, VariableService, override=True) + manager.services.pop(ServiceType.VARIABLE_SERVICE, None) # force a fresh instance + try: + yield + finally: + manager.services.pop(ServiceType.VARIABLE_SERVICE, None) + manager.service_classes.pop(ServiceType.VARIABLE_SERVICE, None) + if prev_class is not _MISSING: + manager.service_classes[ServiceType.VARIABLE_SERVICE] = prev_class + if prev_instance is not _MISSING: + manager.services[ServiceType.VARIABLE_SERVICE] = prev_instance + + +# --------------------------------------------------------------------------- +# get_api_key_for_provider — env gating (no VariableService, == standalone lfx) +# --------------------------------------------------------------------------- + + +def test_literal_key_returned_as_is_regardless_of_flag(): + """A literal key (not an env-var name) is returned verbatim, flag on or off.""" + literal = "sk-literal-abc1234567890" + assert get_api_key_for_provider(None, _PROVIDER, literal) == literal + with _no_env_fallback(): + assert get_api_key_for_provider(None, _PROVIDER, literal) == literal + + +def test_primary_key_resolves_from_env_by_default(monkeypatch): + """No api_key + env fallback enabled: the canonical provider var resolves from os.environ.""" + monkeypatch.setenv(_primary_var_name(), "sk-env") + assert get_api_key_for_provider(None, _PROVIDER, None) == "sk-env" + + +def test_primary_key_env_suppressed_under_no_env_fallback(monkeypatch): + """No api_key + env fallback disabled: the os.getenv read is skipped -> None.""" + monkeypatch.setenv(_primary_var_name(), "sk-env") + with _no_env_fallback(): + assert get_api_key_for_provider(None, _PROVIDER, None) is None + + +def test_named_variable_resolves_from_env_by_default(monkeypatch): + """api_key given as an env-var NAME resolves that name from os.environ.""" + monkeypatch.setenv("MY_OPENAI_KEY", "sk-resolved") + assert get_api_key_for_provider(None, _PROVIDER, "MY_OPENAI_KEY") == "sk-resolved" + + +def test_named_variable_env_suppressed_under_no_env_fallback(monkeypatch): + """Env-var-name api_key + fallback disabled + no user: env skipped, unresolved -> None.""" + monkeypatch.setenv("MY_OPENAI_KEY", "sk-resolved") + with _no_env_fallback(): + # No user_id -> service not consulted; all-caps name resolves to None when unresolved. + assert get_api_key_for_provider(None, _PROVIDER, "MY_OPENAI_KEY") is None + + +def test_unknown_provider_returns_none(): + """A provider with no variable mapping yields None (no env read attempted).""" + assert get_api_key_for_provider(None, "NoSuchProvider", None) is None + + +# --------------------------------------------------------------------------- +# get_all_variables_for_provider — env gating (no user_id branch) +# --------------------------------------------------------------------------- + + +def test_get_all_variables_resolves_env_by_default(monkeypatch): + """No user_id + env fallback enabled: provider vars are read from os.environ.""" + monkeypatch.setenv(_primary_var_name(), "sk-env") + result = get_all_variables_for_provider(None, _PROVIDER) + assert result.get(_primary_var_name()) == "sk-env" + + +def test_get_all_variables_unknown_provider_returns_empty(): + """A provider with no variables yields an empty dict.""" + assert get_all_variables_for_provider(None, "NoSuchProvider") == {} + + +def test_env_if_allowed_gates_connection_config(monkeypatch): + """instantiation._env_if_allowed (used for provider URLs/IDs/headers) honors the flag.""" + from lfx.base.models.unified_models.instantiation import _env_if_allowed + + monkeypatch.setenv("WATSONX_URL", "https://env.example") + assert _env_if_allowed("WATSONX_URL") == "https://env.example" + with _no_env_fallback(): + assert _env_if_allowed("WATSONX_URL") is None + + +# --------------------------------------------------------------------------- +# VariableService path — request-scope resolution + precedence (real service) +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("variable_service_registered") +async def test_request_scope_resolves_under_no_env_fallback(monkeypatch): + """Flag on + request scope set: the resolver returns the request credential, never env.""" + var = _primary_var_name() + monkeypatch.setenv(var, "sk-env-must-not-win") + with _no_env_fallback(), _request_scope({var: "sk-from-request"}): + # Async context -> run_until_complete takes the thread-hop branch; the request + # scope must survive into the worker thread (contextvars propagation). + assert get_api_key_for_provider(_USER_ID, _PROVIDER, None) == "sk-from-request" + + +@pytest.mark.usefixtures("variable_service_registered") +async def test_request_scope_wins_over_env_on_primary_path(monkeypatch): + """Primary path checks VariableService before os.environ, so request scope wins (flag off).""" + var = _primary_var_name() + monkeypatch.setenv(var, "sk-env") + with _request_scope({var: "sk-from-request"}): + assert get_api_key_for_provider(_USER_ID, _PROVIDER, None) == "sk-from-request" + + +@pytest.mark.usefixtures("variable_service_registered") +async def test_named_variable_resolves_from_request_scope_under_no_env_fallback(monkeypatch): + """Named-variable path + flag on: env skipped, request scope resolves the name.""" + monkeypatch.setenv("OPENAI_API_KEY", "sk-env-must-not-win") + with _no_env_fallback(), _request_scope({"OPENAI_API_KEY": "sk-from-request"}): + assert get_api_key_for_provider(_USER_ID, _PROVIDER, "OPENAI_API_KEY") == "sk-from-request" + + +@pytest.mark.usefixtures("variable_service_registered") +async def test_named_variable_request_scope_precedes_env_when_fallback_enabled(monkeypatch): + """Documents current precedence: for an explicitly-named var, VariableService is read first. + + Mirrors the primary path — ``_resolve_var_name`` consults VariableService (where the + user's encrypted global variable, and the request scope it rides on, live) before + ``os.environ``, so the request-scoped credential wins even with env fallback enabled. + """ + monkeypatch.setenv("OPENAI_API_KEY", "sk-env") + with _request_scope({"OPENAI_API_KEY": "sk-from-request"}): + assert get_api_key_for_provider(_USER_ID, _PROVIDER, "OPENAI_API_KEY") == "sk-from-request" + + +@pytest.mark.usefixtures("variable_service_registered") +async def test_get_all_variables_resolves_request_scope_under_no_env_fallback(monkeypatch): + """Flag on + user_id + request scope: provider vars come from the request, not env.""" + var = _primary_var_name() + monkeypatch.setenv(var, "sk-env-must-not-win") + with _no_env_fallback(), _request_scope({var: "sk-from-request"}): + result = get_all_variables_for_provider(_USER_ID, _PROVIDER) + assert result == {var: "sk-from-request"} diff --git a/src/lfx/tests/unit/cli/test_common.py b/src/lfx/tests/unit/cli/test_common.py index abf516b14b..bb6f6dd4a1 100644 --- a/src/lfx/tests/unit/cli/test_common.py +++ b/src/lfx/tests/unit/cli/test_common.py @@ -212,6 +212,7 @@ class TestGraphExecution: yield mock_result mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start results, logs = await execute_graph_with_capture(mock_graph, "test input") @@ -233,6 +234,7 @@ class TestGraphExecution: yield mock_result mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start results, _ = await execute_graph_with_capture(mock_graph, "test input") @@ -240,6 +242,42 @@ class TestGraphExecution: assert len(results) == 1 assert results[0].message.text == "Message text" + @pytest.mark.asyncio + async def test_execute_graph_activates_request_scope_from_context(self): + """Activate the request scope + no_env_fallback from graph.context, then reset them. + + The flag and variables must be live during execution and cleared afterward so + nothing leaks to the next request. Covers the wiring (graph.context -> ContextVars) + that the serve_app tests skip by mocking out execute_graph_with_capture. + """ + from lfx.services.variable.request_scope import ( + get_active_request_variables, + is_env_fallback_disabled, + ) + + observed: dict[str, object] = {} + + async def mock_async_start(inputs, **kwargs): # noqa: ARG001 + observed["scope"] = get_active_request_variables() + observed["no_env_fallback"] = is_env_fallback_disabled() + yield MagicMock(results={"text": "ok"}) + + mock_graph = MagicMock() + mock_graph.context = { + "request_variables": {"access_token": "secret"}, + "no_env_fallback": True, + } + mock_graph.async_start = mock_async_start + + await execute_graph_with_capture(mock_graph, "test input") + + # During execution the scope reflected this request's global_vars + flag. + assert observed["scope"] == {"access_token": "secret"} + assert observed["no_env_fallback"] is True + # After execution both ContextVars are reset (no bleed into the next request). + assert get_active_request_variables() is None + assert is_env_fallback_disabled() is False + @pytest.mark.asyncio async def test_execute_graph_with_capture_error(self): """Test graph execution with error.""" @@ -250,6 +288,7 @@ class TestGraphExecution: yield # This line never executes but makes it an async generator mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start_error with pytest.raises(RuntimeError, match="Execution failed"): @@ -267,6 +306,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start await execute_graph_with_capture(mock_graph, "test input") @@ -282,6 +322,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start await execute_graph_with_capture(mock_graph, "test input", session_id="fixed-session") @@ -301,6 +342,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start memory_vertex = MagicMock() memory_vertex.raw_params = {} @@ -323,6 +365,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start pinned_vertex = MagicMock() pinned_vertex.raw_params = {"session_id": "hardcoded-in-flow"} @@ -347,6 +390,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start mock_graph.user_id = None @@ -363,6 +407,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start mock_graph.user_id = "preset-user-uuid" @@ -384,6 +429,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start await execute_graph_with_capture(mock_graph, "test input") @@ -400,6 +446,7 @@ class TestGraphExecution: yield MagicMock(results={"text": "ok"}) mock_graph = MagicMock() + mock_graph.context = {} mock_graph.async_start = mock_async_start mock_settings = MagicMock() mock_settings.settings.fallback_to_env_var = False diff --git a/src/lfx/tests/unit/cli/test_flow_store.py b/src/lfx/tests/unit/cli/test_flow_store.py new file mode 100644 index 0000000000..bb5c78da06 --- /dev/null +++ b/src/lfx/tests/unit/cli/test_flow_store.py @@ -0,0 +1,143 @@ +"""Tests for FlowStore implementations.""" + +from __future__ import annotations + +import pytest +from lfx.cli.flow_store import FilesystemFlowStore, NullFlowStore + + +class TestNullFlowStore: + def test_write_is_noop(self): + store = NullFlowStore() + store.write("abc", {"name": "test"}) # should not raise + + def test_read_always_returns_none(self): + store = NullFlowStore() + store.write("abc", {"name": "test"}) + assert store.read("abc") is None + + def test_delete_returns_false(self): + assert NullFlowStore().delete("abc") is False + + def test_list_ids_empty(self): + assert NullFlowStore().list_ids() == [] + + +class TestFilesystemFlowStore: + def test_write_and_read_roundtrip(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + data = {"name": "My Flow", "data": {"nodes": [], "edges": []}} + store.write("flow-1", data) + result = store.read("flow-1") + assert result == data + + def test_read_missing_returns_none(self, tmp_path): + assert FilesystemFlowStore(tmp_path).read("nonexistent") is None + + def test_read_corrupt_json_returns_none(self, tmp_path): + """A corrupted store file must be treated as absent, not raise. + + On a shared/PVC store a partially-written or corrupted ``{id}.json`` would + otherwise raise JSONDecodeError into a 500 (or crash warm_from_store at startup). + """ + store = FilesystemFlowStore(tmp_path) + (tmp_path / "broken.json").write_text("{ this is not valid json ", encoding="utf-8") + # list_ids still surfaces it (globs *.json), but read() must contain the damage. + assert "broken" in store.list_ids() + assert store.read("broken") is None + + def test_write_leaves_no_temp_files(self, tmp_path): + """Write must clean up its temp file after the rename.""" + store = FilesystemFlowStore(tmp_path) + store.write("flow-1", {"name": "test"}) + assert list(tmp_path.glob("*.tmp")) == [], "no temp files should remain after write" + assert (tmp_path / "flow-1.json").exists() + + def test_write_is_atomic_under_concurrent_reads(self, tmp_path): + """Concurrent readers must only ever observe a complete old or new payload. + + A reader looping read() while a writer overwrites must never see a partial + file. A non-atomic write (writing into the target in place) would let a + reader see truncated JSON, which read() does not guard against, so + json.loads would raise and surface in the reader thread. + """ + import threading + + store = FilesystemFlowStore(tmp_path) + # Large, distinct payloads so a torn write is detectable and a partial + # read would fail json.loads. + payload_a = {"name": "A", "data": {"nodes": [{"i": i} for i in range(2000)]}} + payload_b = {"name": "B", "data": {"edges": [{"j": j} for j in range(2000)]}} + valid = (payload_a, payload_b) + + store.write("flow-1", payload_a) + + errors: list[Exception] = [] + observed: list[dict | None] = [] + stop = threading.Event() + + def reader() -> None: + try: + while not stop.is_set(): + observed.append(store.read("flow-1")) + except Exception as exc: + errors.append(exc) + + readers = [threading.Thread(target=reader) for _ in range(4)] + for t in readers: + t.start() + try: + for _ in range(200): + store.write("flow-1", payload_b) + store.write("flow-1", payload_a) + finally: + stop.set() + for t in readers: + t.join() + + assert not errors, f"reader saw a corrupt/partial file: {errors[:1]}" + assert observed, "reader never managed to read" + assert all(r in valid for r in observed), "reader observed a non-atomic intermediate state" + + def test_delete_existing_returns_true(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + store.write("flow-1", {"name": "test"}) + assert store.delete("flow-1") is True + assert store.read("flow-1") is None + + def test_delete_missing_returns_false(self, tmp_path): + assert FilesystemFlowStore(tmp_path).delete("nonexistent") is False + + def test_list_ids_empty(self, tmp_path): + assert FilesystemFlowStore(tmp_path).list_ids() == [] + + def test_list_ids_after_write(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + store.write("flow-a", {"name": "A"}) + store.write("flow-b", {"name": "B"}) + assert sorted(store.list_ids()) == ["flow-a", "flow-b"] + + def test_list_ids_excludes_tmp_files(self, tmp_path): + (tmp_path / "partial.json.tmp").write_text("{}") + assert FilesystemFlowStore(tmp_path).list_ids() == [] + + def test_overwrite_existing(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + store.write("flow-1", {"name": "original"}) + store.write("flow-1", {"name": "updated"}) + assert store.read("flow-1")["name"] == "updated" + + def test_creates_directory_if_missing(self, tmp_path): + nested = tmp_path / "deep" / "nested" + FilesystemFlowStore(nested) + assert nested.is_dir() + + def test_rejects_path_traversal_flow_id(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + with pytest.raises(ValueError, match="Invalid flow_id"): + store.write("../escape", {"name": "bad"}) + + def test_rejects_slash_in_flow_id(self, tmp_path): + store = FilesystemFlowStore(tmp_path) + with pytest.raises(ValueError, match="Invalid flow_id"): + store.write("a/b", {"name": "bad"}) diff --git a/src/lfx/tests/unit/cli/test_runtime_variables.py b/src/lfx/tests/unit/cli/test_runtime_variables.py new file mode 100644 index 0000000000..9ecd696d08 --- /dev/null +++ b/src/lfx/tests/unit/cli/test_runtime_variables.py @@ -0,0 +1,98 @@ +import json + +from lfx.cli.runtime_variables import build_request_variables_from_global_vars +from lfx.services.variable.service import VariableService + + +def test_build_request_variables_parses_langflow_request_variables(): + merged = {"access_token": "from-json", "USER_ID": "u1"} + global_vars = { + "LANGFLOW_REQUEST_VARIABLES": json.dumps(merged), + "access_token": "from-raw-override", + "x-langflow-global-var-access-token": "alias-value", + } + flat = build_request_variables_from_global_vars(global_vars) + assert flat["access_token"] == "from-raw-override" # noqa: S105 # test value, not a credential + assert flat["USER_ID"] == "u1" + assert flat["x-langflow-global-var-access-token"] == "alias-value" + + +def test_build_request_variables_keeps_raw_keys_when_blob_is_malformed(): + """A malformed LANGFLOW_REQUEST_VARIABLES blob is skipped; raw keys still merge.""" + global_vars = { + "LANGFLOW_REQUEST_VARIABLES": "{not valid json", + "access_token": "raw-tok", + } + flat = build_request_variables_from_global_vars(global_vars) + assert flat == {"access_token": "raw-tok"} + + +def test_build_request_variables_ignores_non_dict_blob(): + """A JSON blob that is not an object is ignored; raw keys still merge.""" + global_vars = { + "LANGFLOW_REQUEST_VARIABLES": json.dumps(["a", "b"]), + "access_token": "raw-tok", + } + flat = build_request_variables_from_global_vars(global_vars) + assert flat == {"access_token": "raw-tok"} + + +def test_build_request_variables_coerces_structured_blob_values(): + """JSON null is dropped (never the truthy "None"); dict/list become valid JSON strings.""" + global_vars = { + "LANGFLOW_REQUEST_VARIABLES": json.dumps( + {"good": "tok", "null_cred": None, "nested": {"a": 1}, "listy": [1, 2]} + ), + } + flat = build_request_variables_from_global_vars(global_vars) + assert flat == {"good": "tok", "nested": '{"a": 1}', "listy": "[1, 2]"} + # Structured values round-trip back through json.loads (a Python repr would not). + assert json.loads(flat["nested"]) == {"a": 1} + + +def test_normalize_parsed_variables_coercion_rules(): + """Scalars -> str, None dropped, dict/list -> valid JSON (round-trippable).""" + from lfx.services.variable.request_scope import normalize_parsed_variables + + out = normalize_parsed_variables({"s": "v", "n": 5, "f": 1.5, "b": True, "drop": None, "d": {"a": 1}, "l": [1, 2]}) + assert out == {"s": "v", "n": "5", "f": "1.5", "b": "True", "d": '{"a": 1}', "l": "[1, 2]"} + assert "drop" not in out + assert json.loads(out["d"]) == {"a": 1} + + +def test_reset_request_variables_restores_previous_scope(): + """Nested activate/reset restores the prior scope, not None.""" + from lfx.services.variable.request_scope import ( + activate_request_variables, + get_active_request_variables, + reset_request_variables, + ) + + outer = activate_request_variables({"A": "1"}) + try: + inner = activate_request_variables({"A": "2"}) + try: + assert get_active_request_variables() == {"A": "2"} + finally: + reset_request_variables(inner) + assert get_active_request_variables() == {"A": "1"} + finally: + reset_request_variables(outer) + assert get_active_request_variables() is None + + +async def test_variable_service_reads_active_request_scope(): + service = VariableService() + from lfx.services.variable.request_scope import activate_request_variables, reset_request_variables + + token = activate_request_variables( + { + "wxo_github_access_token": "tok-123", + "x-langflow-global-var-wxo-github-access-token": "tok-123", + } + ) + try: + assert await service.get_variable("wxo_github_access_token") == "tok-123" + assert await service.get_variable("access_token") is None + finally: + reset_request_variables(token) diff --git a/src/lfx/tests/unit/cli/test_serve.py b/src/lfx/tests/unit/cli/test_serve.py index 2668598247..e0d7b34351 100644 --- a/src/lfx/tests/unit/cli/test_serve.py +++ b/src/lfx/tests/unit/cli/test_serve.py @@ -97,78 +97,25 @@ def test_flow_meta(): ) -def test_create_multi_serve_app_single_flow(mock_graph, test_flow_meta): - """Test creating app for single flow.""" - with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret - app = create_multi_serve_app( - root_dir=Path("/tmp"), - graphs={"test-flow-id": mock_graph}, - metas={"test-flow-id": test_flow_meta}, - verbose_print=lambda x: None, # noqa: ARG005 - ) +def test_create_multi_serve_app_unknown_flow_id_returns_404(mock_graph, test_flow_meta): + from lfx.cli.serve_app import FlowRegistry + registry = FlowRegistry() + registry.add(mock_graph, test_flow_meta) + + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret + app = create_multi_serve_app(registry=registry) client = TestClient(app) - # Test health endpoint - response = client.get("/health") - assert response.status_code == 200 - assert response.json() == {"status": "healthy", "flow_count": 1} + response = client.get("/flows/does-not-exist/info", headers={"x-api-key": "test-key"}) + assert response.status_code == 404 - # Test run endpoint without auth - response = client.post("/flows/test-flow-id/run", json={"input_value": "test"}) - assert response.status_code == 401 - - # Test run endpoint with auth response = client.post( - "/flows/test-flow-id/run", + "/flows/does-not-exist/run", json={"input_value": "test"}, headers={"x-api-key": "test-key"}, ) - assert response.status_code == 200 - - -def test_create_multi_serve_app_multiple_flows(mock_graph, test_flow_meta): - """Test creating app for multiple flows.""" - meta2 = FlowMeta( - id="flow-2", - relative_path="flow2.json", - title="Flow 2", - description="Second flow", - ) - - with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret - app = create_multi_serve_app( - root_dir=Path("/tmp"), - graphs={"test-flow-id": mock_graph, "flow-2": mock_graph}, - metas={"test-flow-id": test_flow_meta, "flow-2": meta2}, - verbose_print=lambda x: None, # noqa: ARG005 - ) - - client = TestClient(app) - - # Test flows listing - response = client.get("/flows") - assert response.status_code == 200 - flows = response.json() - assert len(flows) == 2 - assert any(f["id"] == "test-flow-id" for f in flows) - assert any(f["id"] == "flow-2" for f in flows) - - # Test individual flow run - response = client.post( - "/flows/test-flow-id/run", - json={"input_value": "test"}, - headers={"x-api-key": "test-key"}, - ) - assert response.status_code == 200 - - # Test flow info - response = client.get( - "/flows/test-flow-id/info", - headers={"x-api-key": "test-key"}, - ) - assert response.status_code == 200 - assert response.json()["id"] == "test-flow-id" + assert response.status_code == 404 def test_serve_command_json_file(): @@ -186,8 +133,8 @@ def test_serve_command_json_file(): try: # Mock the necessary dependencies with ( - patch("lfx.cli.commands.load_graph_from_path") as mock_load, - patch("lfx.cli.commands.uvicorn.Server.serve", new=AsyncMock(return_value=None)) as mock_uvicorn, + patch("lfx.cli.commands.load_flow_from_json") as mock_load, + patch("lfx.cli.commands.uvicorn.run") as mock_uvicorn, patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret ): import typer @@ -227,13 +174,9 @@ def test_serve_command_json_file(): assert mock_uvicorn.called mock_load.assert_called_once() - # Check that the mock was called with the correct arguments - args, kwargs = mock_load.call_args - assert args[0] == Path(temp_path).resolve() - assert args[1] == ".json" - # args[2] is the verbose_print function, which is harder to assert - assert "verbose" in kwargs - assert kwargs["verbose"] is True + # Check that the mock was called with the parsed JSON dict (not a path) + args, _kwargs = mock_load.call_args + assert isinstance(args[0], dict) finally: Path(temp_path).unlink() @@ -244,8 +187,8 @@ def test_serve_command_inline_json(): flow_json = '{"nodes": [], "edges": []}' with ( - patch("lfx.cli.commands.load_graph_from_path") as mock_load, - patch("lfx.cli.commands.uvicorn.Server.serve", new=AsyncMock(return_value=None)) as mock_uvicorn, + patch("lfx.cli.commands.load_flow_from_json") as mock_load, + patch("lfx.cli.commands.uvicorn.run") as mock_uvicorn, patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret ): import typer @@ -271,9 +214,791 @@ def test_serve_command_inline_json(): assert mock_uvicorn.called mock_load.assert_called_once() - # Check that the mock was called with the correct arguments - args, kwargs = mock_load.call_args - assert args[0].suffix == ".json" - assert args[1] == ".json" - assert "verbose" in kwargs - assert kwargs["verbose"] is True + # Check that the mock was called with the parsed JSON dict (not a temp file path) + args, _kwargs = mock_load.call_args + assert isinstance(args[0], dict) + + +class TestBuildRegistryFromDirectory: + def test_loads_all_json_files(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + flow_data = {"nodes": [], "edges": []} + (tmp_path / "a.json").write_text(json.dumps(flow_data)) + (tmp_path / "b.json").write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_directory(tmp_path, lambda _: None, check_variables=False)) + + assert len(registry) == 2 + + def test_empty_directory_raises(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + with pytest.raises(ValueError, match=r"No \.json files found"): + asyncio.run(build_registry_from_directory(tmp_path, lambda _: None, check_variables=False)) + + def test_non_json_files_ignored(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + (tmp_path / "notes.txt").write_text("ignore me") + flow_data = {"nodes": [], "edges": []} + (tmp_path / "flow.json").write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_directory(tmp_path, lambda _: None, check_variables=False)) + + assert len(registry) == 1 + + def test_failed_file_raises_with_filename(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + (tmp_path / "bad.json").write_text('{"nodes": [], "edges": []}') + + with ( + patch("lfx.cli.commands.load_flow_from_json", side_effect=ValueError("corrupt")), + pytest.raises(ValueError, match=r"bad\.json"), + ): + asyncio.run(build_registry_from_directory(tmp_path, lambda _: None, check_variables=False)) + + +class TestBuildRegistryFromPaths: + def test_loads_explicit_paths(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_paths + + flow_data = {"nodes": [], "edges": []} + p1 = tmp_path / "flow1.json" + p2 = tmp_path / "flow2.json" + p1.write_text(json.dumps(flow_data)) + p2.write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_paths([p1, p2], lambda _: None, check_variables=False)) + + assert len(registry) == 2 + + def test_failed_path_raises_with_filename(self, tmp_path): + import asyncio + + from lfx.cli.commands import build_registry_from_paths + + p = tmp_path / "bad.json" + p.write_text('{"nodes": [], "edges": []}') + + with ( + patch("lfx.cli.commands.load_flow_from_json", side_effect=ValueError("oops")), + pytest.raises(ValueError, match=r"bad\.json"), + ): + asyncio.run(build_registry_from_paths([p], lambda _: None, check_variables=False)) + + def test_same_filename_in_different_dirs_gets_distinct_ids(self, tmp_path): + """Regression: lfx serve a/flow.json b/flow.json must not collide on ID.""" + import asyncio + + from lfx.cli.commands import build_registry_from_paths + + dir_a = tmp_path / "a" + dir_b = tmp_path / "b" + dir_a.mkdir() + dir_b.mkdir() + flow_data = {"nodes": [], "edges": []} + p1 = dir_a / "flow.json" + p2 = dir_b / "flow.json" + p1.write_text(json.dumps(flow_data)) + p2.write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_paths([p1, p2], lambda _: None, check_variables=False)) + + assert len(registry) == 2, "both flows must be registered with distinct IDs" + + +class TestServeCommandMultiFlow: + def test_serve_command_with_directory(self, tmp_path): + from lfx.cli.commands import serve_command + + flow_data = {"nodes": [], "edges": []} + (tmp_path / "flow1.json").write_text(json.dumps(flow_data)) + (tmp_path / "flow2.json").write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.nodes = {} + mock_graph.edges = [] + + with ( + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run"), + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + ): + import typer + from typer.testing import CliRunner + + app = typer.Typer() + app.command()(serve_command) + runner = CliRunner() + result = runner.invoke(app, [str(tmp_path)]) + + assert result.exit_code == 0, result.output + + def test_serve_command_with_multiple_files(self, tmp_path): + from lfx.cli.commands import serve_command + + flow_data = {"nodes": [], "edges": []} + p1 = tmp_path / "flow1.json" + p2 = tmp_path / "flow2.json" + p1.write_text(json.dumps(flow_data)) + p2.write_text(json.dumps(flow_data)) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.nodes = {} + mock_graph.edges = [] + + with ( + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run"), + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + ): + import typer + from typer.testing import CliRunner + + app = typer.Typer() + app.command()(serve_command) + runner = CliRunner() + result = runner.invoke(app, [str(p1), str(p2)]) + + assert result.exit_code == 0, result.output + + def test_serve_command_empty_directory_exits_nonzero(self, tmp_path): + from lfx.cli.commands import serve_command + + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret + import typer + from typer.testing import CliRunner + + app = typer.Typer() + app.command()(serve_command) + runner = CliRunner() + result = runner.invoke(app, [str(tmp_path)]) + + assert result.exit_code != 0 + + +class TestPythonScriptServe: + def test_load_graph_and_meta_dispatches_to_script_loader_for_py(self, tmp_path): + """_load_graph_and_meta must call load_graph_from_script for .py files, not load_flow_from_json.""" + import asyncio + + from lfx.cli.commands import _load_graph_and_meta + + script = tmp_path / "my_flow.py" + script.write_text("graph = None") + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + + # load_graph_from_script is lazily imported inside _load_graph_and_meta, + # so patch its module-level name directly. + with ( + patch("lfx.cli.commands.load_graph_from_script", new=AsyncMock(return_value=mock_graph)) as mock_script, + patch("lfx.cli.commands.find_graph_variable", return_value={"source": "x", "type": "Graph", "line": 1}), + patch("lfx.cli.commands.load_flow_from_json") as mock_json, + ): + _graph, meta, raw_json = asyncio.run(_load_graph_and_meta(script, tmp_path, check_variables=False)) + + mock_json.assert_not_called() + mock_script.assert_called_once_with(script) + assert raw_json is None + assert meta.title == "my_flow" + assert meta.relative_path == "my_flow.py" + + def test_serve_command_accepts_py_file(self, tmp_path): + """Lfx serve my_script.py must not be rejected as an unsupported file type.""" + from lfx.cli.commands import serve_command + + script = tmp_path / "my_flow.py" + script.write_text("graph = None") + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.nodes = {} + mock_graph.edges = [] + + with ( + patch("lfx.cli.commands.load_graph_from_script", new=AsyncMock(return_value=mock_graph)), + patch("lfx.cli.commands.find_graph_variable", return_value={"type": "assignment", "line": 1}), + patch("lfx.cli.commands.uvicorn.run"), + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + ): + import typer + from typer.testing import CliRunner + + app = typer.Typer() + app.command()(serve_command) + result = CliRunner().invoke(app, [str(script)]) + + assert result.exit_code == 0, result.output + + def test_serve_command_rejects_unsupported_extension(self, tmp_path): + """Non-.json/.py files must exit with an error.""" + from lfx.cli.commands import serve_command + + bad = tmp_path / "flow.txt" + bad.write_text("not a flow") + + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret + import typer + from typer.testing import CliRunner + + app = typer.Typer() + app.command()(serve_command) + result = CliRunner().invoke(app, [str(bad)]) + + assert result.exit_code != 0 + assert ".json or .py" in result.output + + +def test_build_registry_from_paths_no_env_fallback_stamps_graphs(tmp_path): + """build_registry_from_paths with no_env_fallback=True must stamp each graph's context.""" + import asyncio + + from lfx.cli.commands import build_registry_from_paths + + p = tmp_path / "flow.json" + p.write_text(json.dumps({"nodes": [], "edges": []})) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.context = {} + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run( + build_registry_from_paths([p], lambda _: None, check_variables=False, no_env_fallback=True) + ) + flow_id = registry.list_metas()[0].id + graph, _ = registry.get(flow_id) + assert graph.context.get("no_env_fallback") is True + + +def test_build_registry_from_paths_default_does_not_stamp(tmp_path): + """build_registry_from_paths without no_env_fallback must not stamp the context.""" + import asyncio + + from lfx.cli.commands import build_registry_from_paths + + p = tmp_path / "flow.json" + p.write_text(json.dumps({"nodes": [], "edges": []})) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.context = {} + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_paths([p], lambda _: None, check_variables=False)) + flow_id = registry.list_metas()[0].id + graph, _ = registry.get(flow_id) + assert not graph.context.get("no_env_fallback") + + +def test_build_registry_from_directory_no_env_fallback_stamps_graphs(tmp_path): + """build_registry_from_directory with no_env_fallback=True must stamp each graph's context.""" + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + p = tmp_path / "flow.json" + p.write_text(json.dumps({"nodes": [], "edges": []})) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.context = {} + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run( + build_registry_from_directory(tmp_path, lambda _: None, check_variables=False, no_env_fallback=True) + ) + flow_id = registry.list_metas()[0].id + graph, _ = registry.get(flow_id) + assert graph.context.get("no_env_fallback") is True + + +def test_build_registry_from_directory_default_does_not_stamp(tmp_path): + """build_registry_from_directory without no_env_fallback must not stamp the context.""" + import asyncio + + from lfx.cli.commands import build_registry_from_directory + + p = tmp_path / "flow.json" + p.write_text(json.dumps({"nodes": [], "edges": []})) + + mock_graph = MagicMock() + mock_graph.prepare = MagicMock() + mock_graph.flow_id = None + mock_graph.context = {} + + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_directory(tmp_path, lambda _: None, check_variables=False)) + flow_id = registry.list_metas()[0].id + graph, _ = registry.get(flow_id) + assert not graph.context.get("no_env_fallback") + + +def test_build_registry_from_paths_passes_raw_json_to_store(): + """build_registry_from_paths must pass raw_json so JSON flows are written to the store.""" + import asyncio + import json + import tempfile + from pathlib import Path + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import build_registry_from_paths + from lfx.cli.flow_store import NullFlowStore + + written = {} + + class SpyStore(NullFlowStore): + def write(self, flow_id, flow_json): + written[flow_id] = flow_json + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "flow.json" + p.write_text(json.dumps(flow_data)) + with patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run( + build_registry_from_paths([p], lambda _: None, check_variables=False, store=SpyStore()) + ) + + assert len(written) == 1 + flow_id = registry.list_metas()[0].id + assert written[flow_id]["name"] == "Test" + + +def test_build_registry_from_paths_py_file_skips_store(): + """.py flows must not write to the store (no raw JSON round-trip).""" + import asyncio + import tempfile + from pathlib import Path + from unittest.mock import AsyncMock, MagicMock, patch + + from lfx.cli.commands import build_registry_from_paths + from lfx.cli.flow_store import NullFlowStore + + written = {} + + class SpyStore(NullFlowStore): + def write(self, flow_id, flow_json): + written[flow_id] = flow_json + + mock_graph = MagicMock() + mock_graph.context = {} + + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "flow.py" + p.write_text("graph = None\n") + with ( + patch("lfx.cli.commands.load_graph_from_script", new=AsyncMock(return_value=mock_graph)), + patch("lfx.cli.commands.find_graph_variable", return_value={"source": "x", "type": "Graph", "line": 1}), + ): + asyncio.run(build_registry_from_paths([p], lambda _: None, check_variables=False, store=SpyStore())) + + assert written == {} + + +def test_startup_scan_store_flows_accessible_lazily(tmp_path): + """build_registry_from_paths() does NOT call warm_from_store() — that is serve_command's job. + + Pre-existing store flows are NOT eagerly loaded by the builder but are counted + via list_metas() and accessible on first get(). serve_command adds the + warm_from_store() call for single-worker mode after build_registry_from_paths returns. + """ + import asyncio + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import build_registry_from_paths + from lfx.cli.flow_store import FilesystemFlowStore + + store = FilesystemFlowStore(tmp_path) + raw = {"name": "Pre-existing", "description": None, "data": {"nodes": [], "edges": []}, "id": "pre-existing-id"} + store.write("pre-existing-id", raw) + + mock_graph = MagicMock() + mock_graph.context = {} + + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry = asyncio.run(build_registry_from_paths([], lambda _: None, check_variables=False, store=store)) + + # Builder does NOT eagerly pre-warm — that's serve_command's responsibility + assert "pre-existing-id" not in registry._flows + + # Counted by len() / list_metas() even before warm_from_store() is called + assert len(registry) == 1 + + # Accessible via get() which triggers lazy load (simulating serve_command's warm_from_store) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + result = registry.get("pre-existing-id") + assert result is not None + assert result[1].title == "Pre-existing" + + +def test_serve_command_passes_workers_to_uvicorn(): + """--workers N must be forwarded to uvicorn.run as the workers argument.""" + import json + import os + import tempfile + from pathlib import Path + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import serve_command + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "flow.json" + p.write_text(json.dumps(flow_data)) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run") as mock_run, + ): + serve_command( + script_paths=[str(p)], + host="127.0.0.1", + port=9999, + workers=4, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=None, + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + mock_run.assert_called_once() + call_args = mock_run.call_args[0] + call_kwargs = mock_run.call_args[1] if mock_run.call_args[1] else {} + assert call_kwargs.get("workers") == 4 + # For workers > 1, the app must be the factory import string, not an object + assert call_args[0] == "lfx.cli.serve_app:create_serve_app" + # The factory=True flag is required so uvicorn calls create_serve_app() at + # worker startup, not at request time. + assert call_kwargs.get("factory") is True + + +def test_serve_command_sets_startup_paths_env_for_multi_worker(tmp_path): + """serve_command must set LFX_SERVE_STARTUP_PATHS so each worker can reload flows.""" + import json + import os + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import serve_command + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + captured_env: dict = {} + + def capture_env(*_a, **_kw): + captured_env.update(os.environ) + + p = tmp_path / "flow.json" + p.write_text(json.dumps(flow_data)) + + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run", side_effect=capture_env), + ): + serve_command( + script_paths=[str(p)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=None, + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + assert _SERVE_STARTUP_PATHS_ENV in captured_env, "LFX_SERVE_STARTUP_PATHS must be set before uvicorn.run()" + paths = json.loads(captured_env[_SERVE_STARTUP_PATHS_ENV]) + assert len(paths) == 1 + assert paths[0].endswith("flow.json"), f"expected flow.json in paths, got {paths}" + # Must be cleaned up after uvicorn exits + assert _SERVE_STARTUP_PATHS_ENV not in os.environ, "LFX_SERVE_STARTUP_PATHS must be cleaned up after server exits" + + +def test_serve_command_does_not_set_startup_paths_when_flow_dir_set(tmp_path): + """When --flow-dir is set, LFX_SERVE_STARTUP_PATHS must be empty. + + Workers load startup flows via warm_from_store() (parent persisted them to + the store). Sending file paths would cause each worker to re-load and re-write + them redundantly, and would break the .py restriction logic. + """ + import json + import os + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import serve_command + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + captured_env: dict = {} + + def capture_env(*_a, **_kw): + captured_env.update(os.environ) + + p = tmp_path / "flow.json" + p.write_text(json.dumps(flow_data)) + + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run", side_effect=capture_env), + ): + serve_command( + script_paths=[str(p)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=tmp_path / "store", # flow_dir is set + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + assert _SERVE_STARTUP_PATHS_ENV in captured_env, "env var must still be set (to empty list)" + paths = json.loads(captured_env[_SERVE_STARTUP_PATHS_ENV]) + assert paths == [], f"startup paths must be empty when flow_dir is set, got {paths}" + + +def test_serve_command_warns_when_workers_gt1_without_flow_dir(): + """--workers > 1 without --flow-dir should emit a warning to stderr.""" + import json + import os + import tempfile + from pathlib import Path + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import serve_command + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + stderr_output = [] + + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "flow.json" + p.write_text(json.dumps(flow_data)) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run"), + patch("typer.echo", side_effect=lambda msg, **kw: stderr_output.append(msg) if kw.get("err") else None), + ): + serve_command( + script_paths=[str(p)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=None, + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + assert any("--flow-dir" in msg for msg in stderr_output), ( + f"Expected a warning mentioning --flow-dir, got: {stderr_output}" + ) + + +def test_serve_command_rejects_py_with_multiple_workers(tmp_path): + """.py startup files must be rejected when --workers > 1 (cannot persist to store).""" + import os + from unittest.mock import patch + + from lfx.cli.commands import serve_command + + script = tmp_path / "my_flow.py" + script.write_text("graph = None") + + stderr_output = [] + + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.uvicorn.run"), + patch("typer.echo", side_effect=lambda msg, **kw: stderr_output.append(msg) if kw.get("err") else None), + ): + from click.exceptions import Exit as ClickExit + + with pytest.raises(ClickExit): + serve_command( + script_paths=[str(script)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=tmp_path / "flows", + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + assert any(".py" in msg and "cannot be used" in msg for msg in stderr_output), stderr_output + + +def test_serve_command_allows_py_with_multiple_workers_no_flow_dir(tmp_path): + """.py files + --workers > 1 + no --flow-dir must be allowed. + + Each worker reloads the .py via LFX_SERVE_STARTUP_PATHS since there is no + store. The error only applies when --flow-dir is set (workers would skip + startup paths and warm_from_store, missing the .py flow entirely). + """ + import json + import os + from unittest.mock import AsyncMock, MagicMock, patch + + from lfx.cli.commands import serve_command + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV + + script = tmp_path / "my_flow.py" + script.write_text("graph = None") + + mock_graph = MagicMock() + mock_graph.context = {} + captured_env: dict = {} + + def capture_env(*_a, **_kw): + captured_env.update(os.environ) + + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_graph_from_script", new=AsyncMock(return_value=mock_graph)), + patch("lfx.cli.commands.find_graph_variable", return_value={"type": "assignment", "line": 1}), + patch("lfx.cli.commands.uvicorn.run", side_effect=capture_env), + ): + # Must NOT raise — .py without flow_dir is allowed for multi-worker + serve_command( + script_paths=[str(script)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=None, # no flow_dir — .py is allowed here + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + # LFX_SERVE_STARTUP_PATHS must contain the .py path so workers can reload it + assert _SERVE_STARTUP_PATHS_ENV in captured_env, "LFX_SERVE_STARTUP_PATHS must be set" + paths = json.loads(captured_env[_SERVE_STARTUP_PATHS_ENV]) + assert any("my_flow.py" in p for p in paths), f"expected .py path in STARTUP_PATHS, got {paths}" + + +def test_serve_command_no_warning_when_workers_gt1_with_flow_dir(tmp_path): + """--workers > 1 WITH --flow-dir must not emit the missing-store warning.""" + import json + import os + import tempfile + from pathlib import Path + from unittest.mock import MagicMock, patch + + from lfx.cli.commands import serve_command + + flow_data = {"name": "Test", "description": "", "data": {"nodes": [], "edges": []}} + mock_graph = MagicMock() + mock_graph.context = {} + + stderr_output = [] + + with tempfile.TemporaryDirectory() as tmp: + p = Path(tmp) / "flow.json" + p.write_text(json.dumps(flow_data)) + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}), # pragma: allowlist secret + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + patch("lfx.cli.commands.uvicorn.run"), + patch("typer.echo", side_effect=lambda msg, **kw: stderr_output.append(msg) if kw.get("err") else None), + ): + serve_command( + script_paths=[str(p)], + host="127.0.0.1", + port=9999, + workers=2, + verbose=False, + env_file=None, + log_level="warning", + flow_json=None, + flow_dir=tmp_path / "flows", + stdin=False, + check_variables=False, + no_env_fallback=False, + ) + + assert not any("--flow-dir" in msg for msg in stderr_output) diff --git a/src/lfx/tests/unit/cli/test_serve_app.py b/src/lfx/tests/unit/cli/test_serve_app.py index 6f5c354a47..a53c658496 100644 --- a/src/lfx/tests/unit/cli/test_serve_app.py +++ b/src/lfx/tests/unit/cli/test_serve_app.py @@ -5,13 +5,14 @@ import json import os from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient from lfx.cli.serve_app import ( FlowMeta, + FlowRegistry, create_multi_serve_app, verify_api_key, ) @@ -58,6 +59,881 @@ def allow_custom_components_by_default(monkeypatch): monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) +class TestFlowRegistry: + def _make_meta(self, flow_id: str) -> FlowMeta: + return FlowMeta(id=flow_id, relative_path=f"{flow_id}.json", title=flow_id, description=None) + + def test_add_and_get(self): + registry = FlowRegistry() + graph = MagicMock() + meta = self._make_meta("flow-1") + registry.add(graph, meta) + result = registry.get("flow-1") + assert result is not None + assert result[0] is graph + assert result[1] == meta + + def test_get_missing_returns_none(self): + assert FlowRegistry().get("nonexistent") is None + + def test_list_metas_empty(self): + assert FlowRegistry().list_metas() == [] + + def test_list_metas_multiple(self): + registry = FlowRegistry() + graph = MagicMock() + registry.add(graph, self._make_meta("a")) + registry.add(graph, self._make_meta("b")) + ids = {m.id for m in registry.list_metas()} + assert ids == {"a", "b"} + + def test_duplicate_add_raises_without_overwrite(self): + from lfx.cli.serve_app import FlowAlreadyRegisteredError + + registry = FlowRegistry() + meta = self._make_meta("flow-1") + registry.add(MagicMock(), meta) + with pytest.raises(FlowAlreadyRegisteredError, match="already registered"): + registry.add(MagicMock(), meta) + + def test_duplicate_add_replaces_with_overwrite(self): + registry = FlowRegistry() + g1, g2 = MagicMock(), MagicMock() + meta = self._make_meta("flow-1") + registry.add(g1, meta) + registry.add(g2, meta, overwrite=True) + assert registry.get("flow-1")[0] is g2 + + def test_len(self): + registry = FlowRegistry() + assert len(registry) == 0 + registry.add(MagicMock(), self._make_meta("x")) + assert len(registry) == 1 + + def test_len_counts_store_only_flows(self): + """len() must include flows that are in the store but not yet cache-loaded.""" + raw = {"name": "Store-Only", "data": {}, "id": "store-only-id"} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "store-only-id" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["store-only-id"] + + registry = FlowRegistry(store=StubStore()) + # Nothing in memory yet — but len() must still report 1 via list_metas(). + assert len(registry) == 1 + + def test_len_not_double_counted_memory_and_store(self): + """len() must not double-count a flow that's both in memory and in the store.""" + raw = {"name": "Shared", "data": {}, "id": "shared-id"} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "shared-id" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["shared-id"] + + registry = FlowRegistry(store=StubStore()) + graph = MagicMock() + graph.context = {} + registry.add(graph, self._make_meta("shared-id")) + assert len(registry) == 1 + + def test_warm_from_store_skips_unloadable_flows(self, tmp_path): + """A single corrupt/unloadable store file must not abort warm-up of the rest. + + On a shared/PVC store, one bad ``{id}.json`` would otherwise crash every + worker's startup. warm_from_store must skip it and load the good flows. + """ + from lfx.cli.flow_store import FilesystemFlowStore + + store = FilesystemFlowStore(tmp_path) + # Corrupt file: read() returns None -> get() returns None (skipped silently). + (tmp_path / "corrupt.json").write_text("{ not json ", encoding="utf-8") + # Valid JSON that fails to reconstruct -> get() raises -> warm_from_store must catch. + store.write("unloadable", {"name": "Bad", "data": {"nodes": [], "edges": []}, "id": "unloadable"}) + # A good flow that reconstructs fine. + store.write("good", {"name": "Good", "data": {"nodes": [], "edges": []}, "id": "good"}) + + registry = FlowRegistry(store=store) + good_graph = MagicMock() + good_graph.context = {} + + def fake_load(raw_json): + if raw_json.get("id") == "unloadable": + msg = "component not available in this build" + raise ValueError(msg) + return good_graph + + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=fake_load): + registry.warm_from_store() # must NOT raise despite corrupt + unloadable files + assert registry.get("good") is not None + + def test_remove_existing(self): + registry = FlowRegistry() + meta = self._make_meta("flow-1") + registry.add(MagicMock(), meta) + assert registry.remove("flow-1") is True + assert registry.get("flow-1") is None + + def test_remove_nonexistent(self): + assert FlowRegistry().remove("ghost") is False + + def test_registry_no_env_fallback_stamps_context_on_add(self): + """FlowRegistry(no_env_fallback=True) must set graph.context['no_env_fallback']=True on add.""" + registry = FlowRegistry(no_env_fallback=True) + graph = MagicMock() + graph.context = {} + meta = self._make_meta("test-flow") + registry.add(graph, meta) + assert graph.context.get("no_env_fallback") is True + + def test_registry_default_does_not_stamp_no_env_fallback(self): + """FlowRegistry() (default) must NOT set no_env_fallback on the graph context.""" + registry = FlowRegistry() + graph = MagicMock() + graph.context = {} + meta = self._make_meta("test-flow") + registry.add(graph, meta) + assert "no_env_fallback" not in graph.context + + def test_registry_no_env_fallback_stamps_context_on_overwrite(self): + """Stamp must also be applied when overwrite=True.""" + registry = FlowRegistry(no_env_fallback=True) + first_graph = MagicMock() + first_graph.context = {} + meta = self._make_meta("test-flow") + registry.add(first_graph, meta) + + replacement_graph = MagicMock() + replacement_graph.context = {} + registry.add(replacement_graph, meta, overwrite=True) + assert replacement_graph.context.get("no_env_fallback") is True + + def test_registry_get_misses_to_store(self): + """Cache miss must load from store, cache result, and return it.""" + from unittest.mock import MagicMock, patch + + raw = { + "name": "Stored Flow", + "description": "from disk", + "data": {"nodes": [], "edges": []}, + "id": "flow-from-store", + } + + class StubStore: + def write(self, flow_id, flow_json): + pass + + def read(self, flow_id): + return raw if flow_id == "flow-from-store" else None + + def delete(self, _flow_id): + return False + + def list_ids(self): + return ["flow-from-store"] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + result = registry.get("flow-from-store") + + assert result is not None + graph, meta = result + assert graph is mock_graph + assert meta.id == "flow-from-store" + assert meta.title == "Stored Flow" + # second call must hit in-memory cache, not store + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=AssertionError("should use cache")): + registry.get("flow-from-store") + + def test_list_metas_caches_store_reads(self): + """list_metas() must not re-read a store file it already parsed in a prior call.""" + raw = {"name": "Cached", "data": {"nodes": [], "edges": []}, "id": "cached-id"} + read_count = 0 + + class CountingStore: + def write(self, *_a): + pass + + def read(self, _flow_id): + nonlocal read_count + read_count += 1 + return raw + + def delete(self, *_a): + return False + + def list_ids(self): + return ["cached-id"] + + registry = FlowRegistry(store=CountingStore()) + registry.list_metas() # first call — reads from store, populates cache + registry.list_metas() # second call — must use cache, not read again + assert read_count == 1, f"Expected 1 store read, got {read_count}" + + def test_list_metas_cache_invalidated_after_remove(self): + """After remove(), a subsequent list_metas() must not serve stale cached metadata.""" + raw = {"name": "Gone", "data": {"nodes": [], "edges": []}, "id": "gone-id"} + + class MemStore: + def __init__(self): + self._data = {"gone-id": raw} + + def write(self, fid, v): + self._data[fid] = v + + def read(self, fid): + return self._data.get(fid) + + def delete(self, fid): + return bool(self._data.pop(fid, None)) + + def list_ids(self): + return list(self._data) + + registry = FlowRegistry(store=MemStore()) + metas = registry.list_metas() + assert any(m.id == "gone-id" for m in metas) + + registry.remove("gone-id") + metas_after = registry.list_metas() + assert not any(m.id == "gone-id" for m in metas_after) + + def test_list_metas_skips_store_sourced_flow_deleted_by_other_worker(self): + """list_metas() must not return a flow that another worker deleted from the store.""" + raw = {"name": "Shared", "data": {"nodes": [], "edges": []}, "id": "shared-id"} + + class DeletableStore: + is_persistent = True + + def __init__(self): + self._data = {"shared-id": raw} + + def write(self, fid, v): + self._data[fid] = v + + def read(self, fid): + return self._data.get(fid) + + def delete(self, fid): + return bool(self._data.pop(fid, None)) + + def list_ids(self): + return list(self._data) + + def exists(self, fid): + return fid in self._data + + store = DeletableStore() + registry = FlowRegistry(store=store) + + mock_graph = MagicMock() + mock_graph.context = {} + + # Load flow into registry via get() — simulates warm_from_store on startup + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("shared-id") + + assert any(m.id == "shared-id" for m in registry.list_metas()) + + # Another worker deletes the file — simulate by removing from the store directly + store.delete("shared-id") + + # list_metas() must now exclude the stale entry + assert not any(m.id == "shared-id" for m in registry.list_metas()) + + def test_get_evicts_stale_store_sourced_flow(self): + """get() must return None and evict the entry when the store file was deleted by another worker.""" + raw = {"name": "Evictable", "data": {"nodes": [], "edges": []}, "id": "evict-id"} + + class DeletableStore: + is_persistent = True + + def __init__(self): + self._data = {"evict-id": raw} + + def write(self, fid, v): + self._data[fid] = v + + def read(self, fid): + return self._data.get(fid) + + def delete(self, fid): + return bool(self._data.pop(fid, None)) + + def list_ids(self): + return list(self._data) + + store = DeletableStore() + registry = FlowRegistry(store=store) + + mock_graph = MagicMock() + mock_graph.context = {} + + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + assert registry.get("evict-id") is not None + + # Another worker deletes the file + store.delete("evict-id") + + # get() must now return None and evict the stale entry + assert registry.get("evict-id") is None + assert "evict-id" not in registry._flows + + def test_registry_get_uses_json_id_over_filename_stem(self): + """When JSON has a different id than the filename stem, meta.id uses the JSON id. + + The flow is reachable by that UUID (not just the filename stem). + """ + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = { + "name": "prompt_one", + "id": json_uuid, + "data": {"nodes": [], "edges": []}, + } + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "prompt_one" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["prompt_one"] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + result = registry.get("prompt_one") + + assert result is not None + _, meta = result + assert meta.id == json_uuid, "meta.id must come from the JSON id field" + assert meta.title == "prompt_one" + # flow must also be reachable by UUID without hitting the store again + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=AssertionError("should use cache")): + by_uuid = registry.get(json_uuid) + assert by_uuid is not None + + def test_list_metas_deduplicates_filename_and_json_id_aliases(self): + """list_metas() must not return the same flow twice when it's cached under two keys.""" + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "prompt_one" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["prompt_one"] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("prompt_one") + + metas = registry.list_metas() + assert len(metas) == 1, f"expected 1 meta, got {len(metas)}: {[m.id for m in metas]}" + assert metas[0].id == json_uuid + + def test_overwrite_aliased_flow_removes_old_store_file_and_alias(self): + """add(overwrite=True) on a pre-placed aliased flow must delete the old file. + + Clears the stem alias so new workers don't reconstruct the stale version. + """ + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + old_raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + deleted: list[str] = [] + store_data: dict[str, dict] = {"prompt_one": old_raw} + + class StubStore: + def write(self, fid, data): + store_data[fid] = data + + def read(self, fid): + return store_data.get(fid) + + def delete(self, fid): + existed = fid in store_data + store_data.pop(fid, None) + deleted.append(fid) + return existed + + def list_ids(self): + return list(store_data) + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("prompt_one") # loads and aliases + + # Overwrite via add() (simulating POST /flows/upload/ with replace=True) + new_graph = MagicMock() + new_graph.context = {} + new_meta = FlowMeta(id=json_uuid, relative_path="", title="prompt_one v2", description=None) + new_raw = {"name": "prompt_one v2", "id": json_uuid, "data": {"nodes": [], "edges": []}} + registry.add(new_graph, new_meta, overwrite=True, raw_json=new_raw) + + # Old file must be deleted; new file written under the UUID + assert "prompt_one" in deleted, "old store file must be deleted on overwrite" + assert json_uuid in store_data, "new file must be written under the UUID" + assert "prompt_one" not in store_data, "old file must be gone from store" + + # In-memory: stem alias gone, UUID key has new graph + assert registry.get(json_uuid)[0] is new_graph + assert "prompt_one" not in registry._flows, "stem alias must be cleared" + + # Simulate a new worker: fresh registry warms from store + new_registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=new_graph): + new_registry.warm_from_store() + result = new_registry.get(json_uuid) + assert result is not None + assert result[1].title == "prompt_one v2", "new worker must serve the updated flow, not the old one" + + def test_overwrite_unaliased_preplaced_flow_removes_old_store_file(self): + """add(overwrite=True) must delete a pre-placed file even when get() never ran. + + The replace path in upload_flow skips registry.get(), so _store_keys never + records the stem alias and old_store_key is None. add() must still scan the + store for a file whose JSON "id" matches and delete it — otherwise both + prompt_one.json and {uuid}.json survive. + """ + from unittest.mock import MagicMock + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + old_raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + deleted: list[str] = [] + store_data: dict[str, dict] = {"prompt_one": old_raw} + + class StubStore: + is_persistent = True + + def write(self, fid, data): + store_data[fid] = data + + def read(self, fid): + return store_data.get(fid) + + def delete(self, fid): + existed = fid in store_data + store_data.pop(fid, None) + deleted.append(fid) + return existed + + def list_ids(self): + return list(store_data) + + registry = FlowRegistry(store=StubStore()) + + # No registry.get("prompt_one") here — the replace path skips it, so the + # stem alias is never recorded. + new_graph = MagicMock() + new_graph.context = {} + new_meta = FlowMeta(id=json_uuid, relative_path="", title="prompt_one v2", description=None) + new_raw = {"name": "prompt_one v2", "id": json_uuid, "data": {"nodes": [], "edges": []}} + registry.add(new_graph, new_meta, overwrite=True, raw_json=new_raw) + + assert "prompt_one" in deleted, "stale stem file must be deleted on overwrite" + assert json_uuid in store_data, "new file must be written under the UUID" + assert "prompt_one" not in store_data, "old stem file must be gone — no duplicate" + assert list(store_data) == [json_uuid], "store must hold a single file" + + def test_registry_add_with_raw_json_writes_to_store(self): + """add(raw_json=...) must write to store.""" + from unittest.mock import MagicMock + + written = {} + + class SpyStore: + def write(self, flow_id, flow_json): + written[flow_id] = flow_json + + def read(self, _flow_id): + return None + + def delete(self, _flow_id): + return False + + def list_ids(self): + return [] + + registry = FlowRegistry(store=SpyStore()) + graph = MagicMock() + graph.context = {} + meta = self._make_meta("flow-1") + raw = {"name": "flow-1", "data": {}} + registry.add(graph, meta, raw_json=raw) + + assert written == {"flow-1": raw} + + def test_registry_add_without_raw_json_skips_store(self): + """add() without raw_json must NOT write to store.""" + from unittest.mock import MagicMock + + written = {} + + class SpyStore: + def write(self, flow_id, flow_json): + written[flow_id] = flow_json + + def read(self, _flow_id): + return None + + def delete(self, _flow_id): + return False + + def list_ids(self): + return [] + + registry = FlowRegistry(store=SpyStore()) + graph = MagicMock() + graph.context = {} + registry.add(graph, self._make_meta("flow-1")) + assert written == {} + + def test_registry_remove_deletes_from_store(self): + """remove() must delete from store in addition to clearing in-memory.""" + from unittest.mock import MagicMock + + deleted = [] + + class SpyStore: + def write(self, *_a): + pass + + def read(self, *_a): + return None + + def delete(self, flow_id): + deleted.append(flow_id) + return True + + def list_ids(self): + return [] + + registry = FlowRegistry(store=SpyStore()) + graph = MagicMock() + graph.context = {} + registry.add(graph, self._make_meta("flow-1")) + registry.remove("flow-1") + assert "flow-1" in deleted + + def test_list_metas_includes_store_only_flows(self): + """list_metas() must include flows that are in the store but not yet cached.""" + raw = {"name": "Store-Only Flow", "description": None, "data": {}, "id": "store-only"} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "store-only" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["store-only"] + + registry = FlowRegistry(store=StubStore()) + metas = registry.list_metas() + assert any(m.id == "store-only" for m in metas) + assert any(m.title == "Store-Only Flow" for m in metas) + + def test_registry_remove_store_only_flow_returns_true(self): + """remove() must return True when the flow is in the store but not in memory.""" + raw = {"name": "Store-Only", "data": {}, "id": "store-only"} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "store-only" else None + + def delete(self, flow_id): + return flow_id == "store-only" + + def list_ids(self): + return ["store-only"] + + registry = FlowRegistry(store=StubStore()) + # do NOT call registry.get() first — flow is store-only, not in memory + result = registry.remove("store-only") + assert result is True + + def test_remove_by_uuid_clears_both_aliases(self): + """remove(uuid) must remove the filename-stem alias too, not leave a dangling entry.""" + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + class StubStore: + def __init__(self): + self._deleted: set[str] = set() + + def write(self, *_a): + pass + + def read(self, fid): + return None if fid in self._deleted else (raw if fid == "prompt_one" else None) + + def delete(self, fid): + existed = fid == "prompt_one" and fid not in self._deleted + self._deleted.add(fid) + return existed + + def list_ids(self): + return [f for f in ["prompt_one"] if f not in self._deleted] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("prompt_one") + + result = registry.remove(json_uuid) + assert result is True + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=AssertionError("store already deleted")): + assert registry.get(json_uuid) is None, "flow must be gone by UUID" + assert registry.get("prompt_one") is None, "filename alias must be cleared too" + assert registry.list_metas() == [], "list_metas must not return the removed flow" + + def test_remove_by_stem_clears_both_aliases(self): + """remove(stem) must remove the UUID alias too, not leave a dangling entry.""" + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + deleted: list[str] = [] + + class StubStore: + def __init__(self): + self._deleted: set[str] = set() + + def write(self, *_a): + pass + + def read(self, fid): + return None if fid in self._deleted else (raw if fid == "prompt_one" else None) + + def delete(self, fid): + deleted.append(fid) + existed = fid == "prompt_one" and fid not in self._deleted + self._deleted.add(fid) + return existed + + def list_ids(self): + return [f for f in ["prompt_one"] if f not in self._deleted] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("prompt_one") + + result = registry.remove("prompt_one") + assert result is True + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=AssertionError("store already deleted")): + assert registry.get(json_uuid) is None, "UUID alias must be cleared too" + assert registry.get("prompt_one") is None, "stem must be gone" + assert "prompt_one" in deleted, "must delete the correct file (by stem, not UUID)" + + def test_remove_deletes_both_stem_and_uuid_files_for_cross_worker_propagation(self): + """remove() must delete both the stem file and the UUID file. + + Any worker, regardless of which key its stale check uses, must see the deletion. + + Scenario: flow-dir has my-flow.json (pre-placed) and {uuid}.json (written by add()). + Worker A deletes by UUID; Worker B cached by stem alias. Both files must be gone. + """ + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "my-flow", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + store_data: dict[str, dict] = { + "my-flow": raw, # pre-placed stem file + json_uuid: raw, # UUID file written by add() at startup + } + deleted: list[str] = [] + + class DualKeyStore: + is_persistent = True # required to trigger the stem-scan path in remove() + + def write(self, fid, v): + store_data[fid] = v + + def read(self, fid): + return store_data.get(fid) + + def delete(self, fid): + deleted.append(fid) + existed = fid in store_data + store_data.pop(fid, None) + return existed + + def list_ids(self): + return list(store_data) + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=DualKeyStore()) + + # Worker A: loaded via UUID (no stem alias created) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get(json_uuid) + + # Stale-check key for Worker A is UUID (no alias) + assert registry._store_keys.get(json_uuid) is None + + # Worker A deletes by UUID + result = registry.remove(json_uuid) + assert result is True + + # Both the UUID file AND the stem file must be deleted + assert json_uuid in deleted, "UUID-keyed file must be deleted" + assert "my-flow" in deleted, ( + "stem-keyed file must also be deleted so workers with stem alias don't pass stale check" + ) + assert not store_data, "store must be empty after remove()" + + def test_list_metas_uncached_store_flow_uses_json_id(self): + """list_metas() uncached branch must use the JSON id, not the filename stem.""" + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {}} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "prompt_one" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["prompt_one"] + + registry = FlowRegistry(store=StubStore()) + # Do NOT call registry.get() first — flow stays uncached in the store branch + metas = registry.list_metas() + assert len(metas) == 1 + assert metas[0].id == json_uuid, "uncached list_metas must use JSON id, not filename stem" + + def test_len_not_double_counted_with_aliases(self): + """len() must return 1 when a flow is cached under both filename stem and JSON UUID.""" + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "prompt_one" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["prompt_one"] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.get("prompt_one") + + assert len(registry) == 1, f"expected 1 distinct flow, got {len(registry)}" + + def test_warm_from_store_makes_flow_reachable_by_json_uuid(self): + """After warm_from_store(), a pre-placed file must be reachable by its JSON UUID.""" + from unittest.mock import MagicMock, patch + + json_uuid = "b0529294-e297-41d1-9303-2c2128b7860a" + raw = {"name": "prompt_one", "id": json_uuid, "data": {"nodes": [], "edges": []}} + + class StubStore: + def write(self, *_a): + pass + + def read(self, flow_id): + return raw if flow_id == "prompt_one" else None + + def delete(self, *_a): + return False + + def list_ids(self): + return ["prompt_one"] + + mock_graph = MagicMock() + mock_graph.context = {} + + registry = FlowRegistry(store=StubStore()) + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + registry.warm_from_store() + + result = registry.get(json_uuid) + assert result is not None, "flow must be reachable by its JSON UUID after warm_from_store()" + _, meta = result + assert meta.id == json_uuid + + class TestSecurityFunctions: """Test security-related functions.""" @@ -73,8 +949,8 @@ class TestSecurityFunctions: result = verify_api_key(None, "test-key-123") assert result == "test-key-123" - def test_verify_api_key_header_takes_precedence(self): - """Test that query parameter is used when both are provided.""" + def test_verify_api_key_query_param_takes_precedence(self): + """Query param is checked first; when both are provided the query param value is used.""" with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key-123"}): # pragma: allowlist secret result = verify_api_key("test-key-123", "wrong-key") assert result == "test-key-123" @@ -108,7 +984,6 @@ class TestCreateServeApp: @pytest.fixture def simple_chat_json(self): - """Load the simple chat JSON test data.""" test_data_dir = Path(__file__).parent.parent.parent / "data" json_path = test_data_dir / "simple_chat_no_llm.json" with json_path.open() as f: @@ -116,13 +991,10 @@ class TestCreateServeApp: @pytest.fixture def real_graph(self, simple_chat_json): - """Create a real graph using Graph.from_payload to match serve_app expectations.""" - # Create graph using from_payload with real test data return Graph.from_payload(simple_chat_json, flow_id="00000000-0000-0000-0000-000000000001") @pytest.fixture def mock_meta(self): - """Create mock flow metadata.""" return FlowMeta( id="00000000-0000-0000-0000-000000000001", relative_path="test.json", @@ -131,75 +1003,191 @@ class TestCreateServeApp: ) def test_create_multi_serve_app_single_flow(self, real_graph, mock_meta): - """Test creating app with single flow.""" - graphs = {"00000000-0000-0000-0000-000000000001": real_graph} - metas = {"00000000-0000-0000-0000-000000000001": mock_meta} - verbose_print = Mock() + from lfx.cli.serve_app import FlowRegistry - app = create_multi_serve_app( - root_dir=Path("/test"), - graphs=graphs, - metas=metas, - verbose_print=verbose_print, - ) + registry = FlowRegistry() + registry.add(real_graph, mock_meta) - assert app.title == "LFX Multi-Flow Server (1)" - assert "Use `/flows` to list available IDs" in app.description + app = create_multi_serve_app(registry=registry) - # Check routes - routes = [route.path for route in app.routes] - assert "/health" in routes - assert "/flows" in routes # Multi-flow always has this - assert "/flows/00000000-0000-0000-0000-000000000001/run" in routes # Flow-specific endpoint - - def test_create_multi_serve_app_multiple_flows(self, real_graph, mock_meta, simple_chat_json): - """Test creating app with multiple flows.""" - # Create second real graph using from_payload - graph2 = Graph.from_payload(simple_chat_json, flow_id="flow-2") - - meta2 = FlowMeta( - id="flow-2", - relative_path="flow2.json", - title="Flow 2", - description="Second flow", - ) - - graphs = {"00000000-0000-0000-0000-000000000001": real_graph, "flow-2": graph2} - metas = {"00000000-0000-0000-0000-000000000001": mock_meta, "flow-2": meta2} - verbose_print = Mock() - - app = create_multi_serve_app( - root_dir=Path("/test"), - graphs=graphs, - metas=metas, - verbose_print=verbose_print, - ) - - assert app.title == "LFX Multi-Flow Server (2)" - assert "Use `/flows` to list available IDs" in app.description - - # Check routes routes = [route.path for route in app.routes] assert "/health" in routes assert "/flows" in routes - assert "/flows/00000000-0000-0000-0000-000000000001/run" in routes - assert "/flows/00000000-0000-0000-0000-000000000001/info" in routes - assert "/flows/flow-2/run" in routes - assert "/flows/flow-2/info" in routes + assert "/flows/{flow_id}/run" in routes + assert "/flows/{flow_id}/info" in routes + assert "/flows/upload/" in routes - def test_create_multi_serve_app_mismatched_keys(self, real_graph, mock_meta): - """Test error when graphs and metas have different keys.""" - graphs = {"00000000-0000-0000-0000-000000000001": real_graph} - metas = {"different-id": mock_meta} - verbose_print = Mock() + def test_create_multi_serve_app_multiple_flows(self, real_graph, mock_meta, simple_chat_json): + from lfx.cli.serve_app import FlowRegistry - with pytest.raises(ValueError, match="graphs and metas must contain the same keys"): - create_multi_serve_app( - root_dir=Path("/test"), - graphs=graphs, - metas=metas, - verbose_print=verbose_print, - ) + graph2 = Graph.from_payload(simple_chat_json, flow_id="flow-2") + meta2 = FlowMeta(id="flow-2", relative_path="flow2.json", title="Flow 2", description=None) + + registry = FlowRegistry() + registry.add(real_graph, mock_meta) + registry.add(graph2, meta2) + + app = create_multi_serve_app(registry=registry) + + routes = [route.path for route in app.routes] + assert "/flows/{flow_id}/run" in routes + assert "/flows/{flow_id}/info" in routes + # Single dispatch route covers all flow IDs — no per-flow routes + assert "/flows/00000000-0000-0000-0000-000000000001/run" not in routes + assert "/flows/flow-2/run" not in routes + + +class TestCreateServeAppFactory: + """Tests for the create_serve_app() ASGI factory used by uvicorn workers.""" + + def test_create_serve_app_empty_start(self): + """create_serve_app() with no env vars produces an empty but functional app.""" + import os + from unittest.mock import patch + + from lfx.cli.serve_app import create_serve_app + + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}, clear=False): # pragma: allowlist secret + # Remove startup paths env if present + env = {k: v for k, v in os.environ.items() if not k.startswith("LFX_SERVE_")} + with patch.dict(os.environ, env, clear=True): + os.environ["LANGFLOW_API_KEY"] = "test-key" # pragma: allowlist secret + app = create_serve_app() + + routes = [r.path for r in app.routes] + assert "/health" in routes + assert "/flows" in routes + assert "/flows/upload/" in routes + + def test_create_serve_app_startup_paths_cleaned_from_env_by_caller(self): + """LFX_SERVE_STARTUP_PATHS must be consumed by create_serve_app() but not deleted. + + Cleanup is the caller's (serve_command's) responsibility via the prefix sweep. + """ + import json + import os + from unittest.mock import patch + + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV, create_serve_app + + env_override = { + "LANGFLOW_API_KEY": "test-key", # pragma: allowlist secret + _SERVE_STARTUP_PATHS_ENV: json.dumps([]), + } + + with patch.dict(os.environ, env_override): + # Empty paths list → falls through to else branch + app = create_serve_app() + # The env var should still be present during the call (not deleted inside) + assert _SERVE_STARTUP_PATHS_ENV in os.environ + + assert len(app.state.registry) == 0 + + def test_create_serve_app_with_flow_dir_skips_startup_paths_uses_store(self, tmp_path): + """When flow_dir is set, create_serve_app() must NOT call build_registry_from_paths. + + The parent already persisted startup flows to the store; workers load them via + warm_from_store() only. Re-reading files would cause redundant store writes and + is wrong when the startup files are .py (can't be stored). + """ + import json + import os + from unittest.mock import MagicMock, patch + + from lfx.cli.flow_store import FilesystemFlowStore + from lfx.cli.serve_app import ( + _SERVE_FLOW_DIR_ENV, + _SERVE_STARTUP_PATHS_ENV, + create_serve_app, + ) + + # Pre-place a flow in the store (simulating what the parent did) + store = FilesystemFlowStore(tmp_path) + raw = {"name": "Pre-persisted", "description": None, "data": {"nodes": [], "edges": []}, "id": "pre-id"} + store.write("pre-id", raw) + + mock_graph = MagicMock() + mock_graph.context = {} + + env_override = { + "LANGFLOW_API_KEY": "test-key", # pragma: allowlist secret + _SERVE_FLOW_DIR_ENV: str(tmp_path), + # Startup paths IS set but must be IGNORED since flow_dir is present + _SERVE_STARTUP_PATHS_ENV: json.dumps(["/some/startup/flow.json"]), + } + + with ( + patch.dict(os.environ, env_override), + patch("lfx.cli.commands.build_registry_from_paths") as mock_brfp, + patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph), + ): + app = create_serve_app() + + # build_registry_from_paths must NOT have been called — flow_dir means use the store + mock_brfp.assert_not_called() + # Flow from the store must be in the registry (loaded via warm_from_store) + assert len(app.state.registry) == 1 + + def test_create_serve_app_without_flow_dir_loads_startup_paths_from_files(self, tmp_path): + """When flow_dir is NOT set, create_serve_app() must load startup flows from file paths.""" + import json + import os + from pathlib import Path + from unittest.mock import MagicMock, patch + + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV, create_serve_app + + src = Path(__file__).parent.parent.parent / "data" / "simple_chat_no_llm.json" + flow_path = tmp_path / "flow.json" + flow_path.write_bytes(src.read_bytes()) + + mock_graph = MagicMock() + mock_graph.context = {} + + env_override = { + "LANGFLOW_API_KEY": "test-key", # pragma: allowlist secret + # No LFX_SERVE_FLOW_DIR — no flow_dir + _SERVE_STARTUP_PATHS_ENV: json.dumps([str(flow_path)]), + } + + with ( + patch.dict(os.environ, env_override), + patch("lfx.cli.commands.load_flow_from_json", return_value=mock_graph), + ): + app = create_serve_app() + + assert len(app.state.registry) == 1, "worker must have loaded the startup flow from the file path" + + def test_create_serve_app_startup_path_load_error_propagates(self, tmp_path): + """create_serve_app() must raise (not swallow) if a startup flow file fails to load. + + This exercises the ThreadPoolExecutor path: the coroutine raises inside the + thread, and the exception must propagate back to the worker process so the + worker fails fast instead of silently starting with an empty registry. + """ + import json + import os + + from lfx.cli.serve_app import _SERVE_STARTUP_PATHS_ENV, create_serve_app + + bad_file = tmp_path / "bad_flow.json" + bad_file.write_text(json.dumps({"nodes": [], "edges": []})) + + env_override = { + "LANGFLOW_API_KEY": "test-key", # pragma: allowlist secret + _SERVE_STARTUP_PATHS_ENV: json.dumps([str(bad_file)]), + # No LFX_SERVE_FLOW_DIR — triggers the ThreadPoolExecutor path + } + + with ( + patch.dict(os.environ, env_override), + patch( + "lfx.cli.commands.load_flow_from_json", + side_effect=ValueError("corrupt flow"), + ), + pytest.raises(Exception, match="corrupt flow"), + ): + create_serve_app() class TestServeAppEndpoints: @@ -257,15 +1245,11 @@ class TestServeAppEndpoints: description="A test flow", ) - graphs = {"00000000-0000-0000-0000-000000000001": real_graph_with_async} - metas = {"00000000-0000-0000-0000-000000000001": meta} - verbose_print = Mock() + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) app = create_multi_serve_app( - root_dir=Path("/test"), - graphs=graphs, - metas=metas, - verbose_print=verbose_print, + registry=registry, ) monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) @@ -301,15 +1285,12 @@ class TestServeAppEndpoints: description="Second flow", ) - graphs = {"00000000-0000-0000-0000-000000000001": real_graph_with_async, "flow-2": graph2} - metas = {"00000000-0000-0000-0000-000000000001": meta1, "flow-2": meta2} - verbose_print = Mock() + registry = FlowRegistry() + registry.add(real_graph_with_async, meta1) + registry.add(graph2, meta2) app = create_multi_serve_app( - root_dir=Path("/test"), - graphs=graphs, - metas=metas, - verbose_print=verbose_print, + registry=registry, ) monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) @@ -385,11 +1366,10 @@ class TestServeAppEndpoints: title="Test Flow", description="A test flow", ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) app = create_multi_serve_app( - root_dir=Path("/test"), - graphs={"00000000-0000-0000-0000-000000000001": real_graph_with_async}, - metas={"00000000-0000-0000-0000-000000000001": meta}, - verbose_print=Mock(), + registry=registry, ) headers = {"x-api-key": "test-api-key"} @@ -420,7 +1400,7 @@ class TestServeAppEndpoints: headers=headers, ) - assert response.status_code == 200 + assert response.status_code == 500 assert response.json()["success"] is False assert "custom components are not allowed" in response.json()["result"] mock_execute.assert_not_called() @@ -464,13 +1444,11 @@ class TestServeAppEndpoints: "/flows/00000000-0000-0000-0000-000000000001/run", json=request_data, headers=headers ) - assert response.status_code == 200 # Returns 200 with error in response body + assert response.status_code == 500 data = response.json() assert data["success"] is False - # serve_app error handling returns "Flow execution failed: {error}" assert data["result"] == "Flow execution failed: Flow execution failed" assert data["type"] == "error" - # The error message should be in the logs assert "ERROR: Flow execution failed" in data["logs"] def test_run_endpoint_no_results(self, app_client): @@ -490,7 +1468,7 @@ class TestServeAppEndpoints: "/flows/00000000-0000-0000-0000-000000000001/run", json=request_data, headers=headers ) - assert response.status_code == 200 + assert response.status_code == 500 data = response.json() assert data["result"] == "No response generated" assert data["success"] is False @@ -511,11 +1489,8 @@ class TestServeAppEndpoints: patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), ): - response = app_client.post( - "/flows/00000000-0000-0000-0000-000000000001/run", json=request_data, headers=headers - ) + app_client.post("/flows/00000000-0000-0000-0000-000000000001/run", json=request_data, headers=headers) - assert response.status_code == 200 assert captured["session_id"] == "my-conversation" def test_stream_endpoint_forwards_session_id(self, app_client): @@ -545,7 +1520,8 @@ class TestServeAppEndpoints: def test_list_flows_endpoint(self, multi_flow_client): """Test listing flows in multi-flow mode.""" - response = multi_flow_client.get("/flows") + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}): # pragma: allowlist secret + response = multi_flow_client.get("/flows", headers={"x-api-key": "test-api-key"}) assert response.status_code == 200 flows = response.json() @@ -599,6 +1575,280 @@ class TestServeAppEndpoints: assert response.status_code == 422 # Validation error + def test_run_endpoint_injects_global_vars_into_context(self, real_graph_with_async, monkeypatch): + """global_vars in RunRequest must appear in graph_copy.context['request_variables'].""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) + app = create_multi_serve_app(registry=registry) + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + captured: dict = {} + + async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + captured["request_variables"] = dict(graph.context.get("request_variables") or {}) + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), + TestClient(app) as client, + ): + client.post( + "/flows/00000000-0000-0000-0000-000000000001/run", + json={ + "input_value": "hello", + "global_vars": {"MY_API_KEY": "secret-value"}, # pragma: allowlist secret + }, + headers=headers, + ) + + assert captured["request_variables"] == {"MY_API_KEY": "secret-value"} # pragma: allowlist secret + + def test_run_endpoint_global_vars_do_not_mutate_registry_graph(self, real_graph_with_async, monkeypatch): + """global_vars must only be set on the deepcopy; the registry's original graph must be unchanged.""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) + 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 + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_noop), + TestClient(app) as client, + ): + client.post( + "/flows/00000000-0000-0000-0000-000000000001/run", + json={ + "input_value": "hello", + "global_vars": {"MY_API_KEY": "secret-value"}, # pragma: allowlist secret + }, + headers=headers, + ) + + original_graph = registry.get("00000000-0000-0000-0000-000000000001")[0] + rv = original_graph.context.get("request_variables") or {} + assert "MY_API_KEY" not in rv + + def test_run_endpoint_preserves_no_env_fallback_after_deepcopy(self, real_graph_with_async, monkeypatch): + """no_env_fallback must reach the executed graph_copy, not only the registry graph. + + deepcopy() drops graph.context, so run_flow must re-apply the stamp after the copy. + """ + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry(no_env_fallback=True) + registry.add(real_graph_with_async, meta) + app = create_multi_serve_app(registry=registry) + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + captured: dict = {} + + async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + captured["no_env_fallback"] = graph.context.get("no_env_fallback") + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), + TestClient(app) as client, + ): + client.post( + "/flows/00000000-0000-0000-0000-000000000001/run", + json={"input_value": "hello"}, + headers=headers, + ) + + assert captured["no_env_fallback"] is True + + def test_stream_endpoint_preserves_no_env_fallback_after_deepcopy(self, real_graph_with_async, monkeypatch): + """no_env_fallback must reach the executed graph_copy on the stream path too.""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry(no_env_fallback=True) + registry.add(real_graph_with_async, meta) + app = create_multi_serve_app(registry=registry) + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + captured: dict = {} + + async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + captured["no_env_fallback"] = graph.context.get("no_env_fallback") + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), + TestClient(app) as client, + client.stream( + "POST", + "/flows/00000000-0000-0000-0000-000000000001/stream", + json={"input_value": "hello"}, + headers=headers, + ) as response, + ): + assert response.status_code == 200 + for _ in response.iter_bytes(): + pass + + assert captured["no_env_fallback"] is True + + def test_stream_endpoint_injects_global_vars_into_context(self, real_graph_with_async, monkeypatch): + """global_vars in StreamRequest must appear in graph_copy.context['request_variables'].""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) + app = create_multi_serve_app(registry=registry) + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + captured: dict = {} + + async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + captured["request_variables"] = dict(graph.context.get("request_variables") or {}) + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), + TestClient(app) as client, + client.stream( + "POST", + "/flows/00000000-0000-0000-0000-000000000001/stream", + json={ + "input_value": "hello", + "global_vars": {"STREAM_KEY": "stream-secret"}, # pragma: allowlist secret + }, + headers=headers, + ) as response, + ): + assert response.status_code == 200 + for _ in response.iter_bytes(): + pass + + assert captured["request_variables"] == {"STREAM_KEY": "stream-secret"} # pragma: allowlist secret + + def test_stream_endpoint_no_global_vars_leaves_context_clean(self, real_graph_with_async, monkeypatch): + """Omitting global_vars must not create request_variables in the graph context.""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) + app = create_multi_serve_app(registry=registry) + monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True) + + captured: dict = {} + + async def mock_execute_capture(graph, input_value, session_id=None): # noqa: ARG001 + captured["request_variables"] = graph.context.get("request_variables") + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_capture), + TestClient(app) as client, + client.stream( + "POST", + "/flows/00000000-0000-0000-0000-000000000001/stream", + json={"input_value": "hello"}, # no global_vars key + headers=headers, + ) as response, + ): + assert response.status_code == 200 + for _ in response.iter_bytes(): + pass + + assert not captured.get("request_variables") + + def test_stream_endpoint_global_vars_do_not_mutate_registry_graph(self, real_graph_with_async, monkeypatch): + """global_vars must only be set on the deepcopy; the registry's original graph must be unchanged.""" + from lfx.services.deps import get_settings_service + + meta = FlowMeta( + id="00000000-0000-0000-0000-000000000001", + relative_path="test.json", + title="Test Flow", + description=None, + ) + registry = FlowRegistry() + registry.add(real_graph_with_async, meta) + 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 + return [], "" + + headers = {"x-api-key": "test-api-key"} + with ( + patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-api-key"}), # pragma: allowlist secret + patch("lfx.cli.serve_app.execute_graph_with_capture", mock_execute_noop), + TestClient(app) as client, + client.stream( + "POST", + "/flows/00000000-0000-0000-0000-000000000001/stream", + json={ + "input_value": "hello", + "global_vars": {"STREAM_KEY": "stream-secret"}, # pragma: allowlist secret + }, + headers=headers, + ) as response, + ): + assert response.status_code == 200 + for _ in response.iter_bytes(): + pass + + original_graph = registry.get("00000000-0000-0000-0000-000000000001")[0] + rv = original_graph.context.get("request_variables") or {} + assert "STREAM_KEY" not in rv + def test_flow_execution_with_message_output(self, app_client, real_graph_with_async): """Test flow execution with message-type output.""" @@ -643,3 +1893,295 @@ class TestServeAppEndpoints: data = response.json() assert data["result"] == "Message output" assert data["success"] is True + + +class TestUploadEndpoint: + """Tests for POST /flows/upload/.""" + + @pytest.fixture + def app_with_empty_registry(self): + from lfx.cli.serve_app import FlowRegistry + + registry = FlowRegistry() + app = create_multi_serve_app(registry=registry) + with patch.dict(os.environ, {"LANGFLOW_API_KEY": "test-key"}): # pragma: allowlist secret + yield TestClient(app) + + @pytest.fixture + def valid_flow_data(self): + test_data_dir = Path(__file__).parent.parent.parent / "data" + json_path = test_data_dir / "simple_chat_no_llm.json" + with json_path.open() as f: + return json.load(f) + + @pytest.fixture + def full_export(self): + """Full Langflow export JSON as exported from the UI (name/data/... at top level). + + This is what a user sends when they run: + curl -X POST .../flows/upload/ -d @myflow.json + + body.data will be {"edges": [...], "nodes": [...]} — the inner graph with NO + nested "data" key. A regression that calls load_flow_from_json(body.data) + raises KeyError('data') here; the correct call passes body.model_dump(...). + """ + test_data_dir = Path(__file__).parent.parent.parent / "data" + json_path = test_data_dir / "simple_chat_no_llm.json" + with json_path.open() as f: + return json.load(f) + + def test_upload_full_export_as_body(self, app_with_empty_registry, full_export): + """Uploading a Langflow export JSON directly as the body must succeed. + + Regression test: load_flow_from_json must be called with the full model dict + (which has a top-level "data" key), not body.data alone (which is just the + inner graph and has no "data" key). + """ + response = app_with_empty_registry.post( + "/flows/upload/", + json=full_export, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 201, response.json() + body = response.json() + assert body["name"] == full_export["name"] + assert body["run_url"].startswith("/flows/") + assert body["run_url"].endswith("/run") + + def test_upload_valid_flow(self, app_with_empty_registry, valid_flow_data): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "My Uploaded Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 201 + body = response.json() + assert body["name"] == "My Uploaded Flow" + assert body["run_url"].startswith("/flows/") + assert body["run_url"].endswith("/run") + assert "id" in body + + def test_upload_requires_auth(self, app_with_empty_registry, valid_flow_data): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Flow", "data": valid_flow_data}, + ) + assert response.status_code == 401 + + def test_upload_invalid_flow_data_returns_422(self, app_with_empty_registry): + with patch("lfx.cli.serve_app.load_flow_from_json", side_effect=ValueError("bad flow")): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Bad Flow", "data": {"nodes": [], "edges": []}}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 422 + assert "bad flow" in response.json()["detail"] + + def test_upload_prepare_failure_returns_422(self, app_with_empty_registry): + mock_graph = MagicMock() + mock_graph.prepare.side_effect = RuntimeError("prepare failed") + with patch("lfx.cli.serve_app.load_flow_from_json", return_value=mock_graph): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Bad Flow", "data": {"nodes": [], "edges": []}}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 422 + assert "prepare failed" in response.json()["detail"] + + def test_upload_flow_is_immediately_listed(self, app_with_empty_registry, valid_flow_data): + upload_resp = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Runnable Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert upload_resp.status_code == 201 + flow_id = upload_resp.json()["id"] + + list_resp = app_with_empty_registry.get("/flows", headers={"x-api-key": "test-key"}) + assert any(f["id"] == flow_id for f in list_resp.json()) + + def test_upload_duplicate_without_replace_returns_409(self, app_with_empty_registry, valid_flow_data): + fixed_id = "aaaabbbb-1111-2222-3333-444455556666" + r1 = app_with_empty_registry.post( + "/flows/upload/", + json={"id": fixed_id, "name": "Flow A", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert r1.status_code == 201 + + r2 = app_with_empty_registry.post( + "/flows/upload/", + json={"id": fixed_id, "name": "Flow B", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert r2.status_code == 409 + assert "already exists" in r2.json()["detail"] + + def test_upload_replace_true_overwrites(self, app_with_empty_registry, valid_flow_data): + fixed_id = "bbbbcccc-1111-2222-3333-444455556666" + r1 = app_with_empty_registry.post( + "/flows/upload/", + json={"id": fixed_id, "name": "Original Name", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert r1.status_code == 201 + flow_id = r1.json()["id"] + + r2 = app_with_empty_registry.post( + "/flows/upload/", + json={"id": fixed_id, "name": "Updated Name", "data": valid_flow_data, "replace": True}, + headers={"x-api-key": "test-key"}, + ) + assert r2.status_code == 201 + assert r2.json()["id"] == flow_id + assert r2.json()["name"] == "Updated Name" + + flows = app_with_empty_registry.get("/flows", headers={"x-api-key": "test-key"}).json() + ids = [f["id"] for f in flows] + assert ids.count(flow_id) == 1 + + def test_upload_with_description(self, app_with_empty_registry, valid_flow_data): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Flow", "data": valid_flow_data, "description": "my desc"}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 201 + assert response.json()["description"] == "my desc" + + # ------------------------------------------------------------------ + # Execution failures return HTTP 500 not HTTP 200 + # ------------------------------------------------------------------ + + def test_run_execution_exception_returns_500(self, app_with_empty_registry, valid_flow_data): + """An unhandled exception during graph execution must produce HTTP 500, not 200.""" + upload = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Failing Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert upload.status_code == 201 + flow_id = upload.json()["id"] + + with patch( + "lfx.cli.serve_app.execute_graph_with_capture", + side_effect=RuntimeError("boom"), + ): + response = app_with_empty_registry.post( + f"/flows/{flow_id}/run", + json={"input_value": "hi"}, + headers={"x-api-key": "test-key"}, + ) + + assert response.status_code == 500 + body = response.json() + assert body["success"] is False + assert "boom" in body["result"] + + def test_run_flow_failure_result_returns_500(self, app_with_empty_registry, valid_flow_data): + """A flow that returns success=False in its result must also produce HTTP 500.""" + upload = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Bad Result Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert upload.status_code == 201 + flow_id = upload.json()["id"] + + with ( + patch("lfx.cli.serve_app.execute_graph_with_capture", return_value=([], "")), + patch( + "lfx.cli.serve_app.extract_result_data", + return_value={"success": False, "result": "flow failed", "type": "error"}, + ), + ): + response = app_with_empty_registry.post( + f"/flows/{flow_id}/run", + json={"input_value": "hi"}, + headers={"x-api-key": "test-key"}, + ) + + assert response.status_code == 500 + body = response.json() + assert body["success"] is False + + # ------------------------------------------------------------------ + # GET /flows requires authentication + # ------------------------------------------------------------------ + + def test_list_flows_requires_auth(self, app_with_empty_registry): + response = app_with_empty_registry.get("/flows") + assert response.status_code == 401 + + def test_list_flows_with_auth(self, app_with_empty_registry, valid_flow_data): + upload = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "My Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert upload.status_code == 201 + + response = app_with_empty_registry.get("/flows", headers={"x-api-key": "test-key"}) + assert response.status_code == 200 + assert any(f["id"] == upload.json()["id"] for f in response.json()) + + # ------------------------------------------------------------------ + # Runtime flow removal via DELETE /flows/{flow_id} + # ------------------------------------------------------------------ + + def test_delete_flow_removes_it(self, app_with_empty_registry, valid_flow_data): + upload = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Temp Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + assert upload.status_code == 201 + flow_id = upload.json()["id"] + + delete_resp = app_with_empty_registry.delete(f"/flows/{flow_id}", headers={"x-api-key": "test-key"}) + assert delete_resp.status_code == 204 + + run_resp = app_with_empty_registry.post( + f"/flows/{flow_id}/run", + json={"input_value": "hi"}, + headers={"x-api-key": "test-key"}, + ) + assert run_resp.status_code == 404 + + def test_delete_nonexistent_flow_returns_404(self, app_with_empty_registry): + response = app_with_empty_registry.delete("/flows/does-not-exist", headers={"x-api-key": "test-key"}) + assert response.status_code == 404 + + def test_delete_flow_requires_auth(self, app_with_empty_registry, valid_flow_data): + upload = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Flow", "data": valid_flow_data}, + headers={"x-api-key": "test-key"}, + ) + flow_id = upload.json()["id"] + response = app_with_empty_registry.delete(f"/flows/{flow_id}") + assert response.status_code == 401 + + # ------------------------------------------------------------------ + # body.id must be a valid UUID to prevent path/store injection + # ------------------------------------------------------------------ + + def test_upload_with_valid_uuid_id(self, app_with_empty_registry, valid_flow_data): + explicit_id = "b0529294-e297-41d1-9303-2c2128b7860a" + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Flow", "data": valid_flow_data, "id": explicit_id}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 201 + assert response.json()["id"] == explicit_id + + def test_upload_with_invalid_id_returns_422(self, app_with_empty_registry, valid_flow_data): + response = app_with_empty_registry.post( + "/flows/upload/", + json={"name": "Flow", "data": valid_flow_data, "id": "not-a-uuid"}, + headers={"x-api-key": "test-key"}, + ) + assert response.status_code == 422 diff --git a/src/lfx/tests/unit/cli/test_serve_app_streaming.py b/src/lfx/tests/unit/cli/test_serve_app_streaming.py index 165fa070eb..e28962c6f5 100644 --- a/src/lfx/tests/unit/cli/test_serve_app_streaming.py +++ b/src/lfx/tests/unit/cli/test_serve_app_streaming.py @@ -2,15 +2,13 @@ import asyncio import hashlib -import tempfile -from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, patch import pytest from asgi_lifespan import LifespanManager from httpx import ASGITransport, AsyncClient -from lfx.cli.serve_app import FlowMeta, StreamRequest, create_multi_serve_app +from lfx.cli.serve_app import FlowMeta, FlowRegistry, StreamRequest, create_multi_serve_app from lfx.interface.components import component_cache @@ -124,24 +122,24 @@ def multi_serve_app(mock_graphs, mock_metas, monkeypatch): "Execution completed successfully", ) - with tempfile.TemporaryDirectory() as temp_dir: - app = create_multi_serve_app( - root_dir=Path(temp_dir), graphs=mock_graphs, metas=mock_metas, verbose_print=lambda _: None - ) + registry = FlowRegistry() + for flow_id, graph in mock_graphs.items(): + registry.add(graph, mock_metas[flow_id]) + app = create_multi_serve_app(registry=registry) - # Override the dependency after app creation - def mock_verify_api_key(query_param: str | None = None, header_param: str | None = None) -> str: # noqa: ARG001 - return "test-api-key" + # Override the dependency after app creation + def mock_verify_api_key(query_param: str | None = None, header_param: str | None = None) -> str: # noqa: ARG001 + return "test-api-key" - # Import the original dependency - from lfx.cli.serve_app import verify_api_key + # Import the original dependency + from lfx.cli.serve_app import verify_api_key - app.dependency_overrides[verify_api_key] = mock_verify_api_key + app.dependency_overrides[verify_api_key] = mock_verify_api_key - yield app + yield app - # Clean up - app.dependency_overrides.clear() + # Clean up + app.dependency_overrides.clear() @pytest.fixture @@ -431,12 +429,10 @@ class TestMultiServeStreaming: monkeypatch.setenv("LANGFLOW_API_KEY", "test-api-key") mock_graphs["flow1"].raw_graph_data = _blocked_raw_graph() - app = create_multi_serve_app( - root_dir=Path("/tmp"), - graphs=mock_graphs, - metas=mock_metas, - verbose_print=lambda _: None, - ) + registry = FlowRegistry() + for flow_id, graph in mock_graphs.items(): + registry.add(graph, mock_metas[flow_id]) + app = create_multi_serve_app(registry=registry) from lfx.cli.serve_app import verify_api_key @@ -481,3 +477,48 @@ class TestMultiServeStreaming: assert response.status_code == 200 assert "custom components are not allowed" in response.text mock_run_flow.assert_not_called() + + @pytest.mark.asyncio + async def test_stream_setup_error_returns_valid_json(self, mock_graphs, mock_metas, monkeypatch): + """Error messages with special characters must not break the SSE JSON payload.""" + monkeypatch.setenv("LANGFLOW_API_KEY", "test-api-key") + + registry = FlowRegistry() + for flow_id, graph in mock_graphs.items(): + registry.add(graph, mock_metas[flow_id]) + app = create_multi_serve_app(registry=registry) + + from lfx.cli.serve_app import verify_api_key + + app.dependency_overrides[verify_api_key] = lambda: "test-api-key" + + # Exception message contains characters that would break hand-crafted JSON + nasty_message = 'failed: "quotes" and\nnewlines' + + async with ( + LifespanManager(app, startup_timeout=None, shutdown_timeout=None) as manager, + AsyncClient( + transport=ASGITransport(app=manager.app), + base_url="http://testserver/", + http2=True, + ) as client, + ): + with patch( + "lfx.cli.serve_app.validate_flow_for_current_settings", + side_effect=RuntimeError(nasty_message), + ): + response = await client.post( + "/flows/flow1/stream", + json={"input_value": "test"}, + headers={"x-api-key": "test-api-key"}, + ) + + assert response.status_code == 200 + # Strip the SSE framing and parse as JSON — must not raise + for line in response.text.splitlines(): + if line.startswith("data: "): + import json + + payload = json.loads(line[len("data: ") :]) + assert payload["success"] is False + assert nasty_message in payload["error"] diff --git a/src/lfx/tests/unit/cli/test_serve_components.py b/src/lfx/tests/unit/cli/test_serve_components.py index 961d88375a..f6a0ae79e9 100644 --- a/src/lfx/tests/unit/cli/test_serve_components.py +++ b/src/lfx/tests/unit/cli/test_serve_components.py @@ -12,10 +12,9 @@ from lfx.cli.common import flow_id_from_path, load_graph_from_path, validate_scr from lfx.cli.serve_app import ( ErrorResponse, FlowMeta, + FlowRegistry, RunRequest, RunResponse, - _analyze_graph_structure, - _generate_dynamic_run_description, create_multi_serve_app, ) from lfx.graph import Graph @@ -75,89 +74,6 @@ class TestDataModels: assert error.success is False -class TestGraphAnalysis: - """Test graph analysis functions.""" - - def test_analyze_graph_structure_basic(self): - """Test basic graph structure analysis.""" - # Create a mock graph that matches what _analyze_graph_structure expects - mock_graph = Mock() - - # Create mock node objects with the expected structure - node1 = Mock() - node1.data = { - "type": "ChatInput", - "display_name": "Chat Input", - "description": "Input component", - "template": {"input_value": {"type": "str"}}, - } - - node2 = Mock() - node2.data = { - "type": "ChatOutput", - "display_name": "Chat Output", - "description": "Output component", - "template": {"output_value": {"type": "str"}}, - } - - mock_graph.nodes = {"input-1": node1, "output-1": node2} - - # Create mock edges - edge = Mock() - edge.source = "input-1" - edge.target = "output-1" - mock_graph.edges = [edge] - - analysis = _analyze_graph_structure(mock_graph) - - assert analysis["node_count"] == 2 - assert analysis["edge_count"] == 1 - assert len(analysis["components"]) == 2 - assert isinstance(analysis["input_types"], list) - assert isinstance(analysis["output_types"], list) - - def test_analyze_graph_structure_error_handling(self): - """Test graph analysis with malformed graph.""" - mock_graph = Mock() - mock_graph.nodes = {} - mock_graph.edges = [] - - # Force an exception during analysis - mock_graph.nodes = None - - analysis = _analyze_graph_structure(mock_graph) - - # Should provide fallback values - assert len(analysis["components"]) == 1 - assert analysis["components"][0]["type"] == "Unknown" - assert "text" in analysis["input_types"] - assert "text" in analysis["output_types"] - - def test_generate_dynamic_run_description(self): - """Test dynamic description generation.""" - # Create a mock graph for _generate_dynamic_run_description - mock_graph = Mock() - - # Mock the analyze function to return expected data - with patch("lfx.cli.serve_app._analyze_graph_structure") as mock_analyze: - mock_analyze.return_value = { - "node_count": 2, - "edge_count": 1, - "components": [{"type": "ChatInput"}, {"type": "ChatOutput"}], - "input_types": ["text"], - "output_types": ["text"], - "entry_points": [{"template": {"input_value": {"type": "str"}}}], - "exit_points": [{"template": {"output_value": {"type": "str"}}}], - } - - description = _generate_dynamic_run_description(mock_graph) - - assert "Execute the deployed LFX graph" in description - assert "Authentication Required" in description - assert "Example Request" in description - assert "Example Response" in description - - class TestCommonFunctions: """Test common utility functions.""" @@ -255,58 +171,39 @@ def create_real_graph(): class TestFastAPIAppCreation: """Test FastAPI application creation.""" - def test_create_multi_serve_app_basic(self, tmp_path): + def test_create_multi_serve_app_basic(self): """Test basic multi-serve app creation.""" - root_dir = tmp_path - graphs = {"test-flow": create_real_graph()} - metas = {"test-flow": FlowMeta(id="test-flow", relative_path="test.json", title="Test Flow")} + meta = FlowMeta(id="test-flow", relative_path="test.json", title="Test Flow") + registry = FlowRegistry() + registry.add(create_real_graph(), meta) - def verbose_print(msg): - pass # Real function + app = create_multi_serve_app(registry=registry) - with patch("lfx.cli.serve_app.verify_api_key"): - app = create_multi_serve_app(root_dir=root_dir, graphs=graphs, metas=metas, verbose_print=verbose_print) + assert app.title.startswith("LFX Multi-Flow Server") + assert "1" in app.title # Should show count - assert app.title.startswith("LFX Multi-Flow Server") - assert "1" in app.title # Should show count + def test_create_multi_serve_app_mismatched_keys(self): + """Test app creation — registry never has mismatched keys by design; verify basic creation.""" + meta = FlowMeta(id="test-flow", relative_path="test.json", title="Test Flow") + registry = FlowRegistry() + registry.add(create_real_graph(), meta) - def test_create_multi_serve_app_mismatched_keys(self, tmp_path): - """Test app creation with mismatched graph/meta keys.""" - root_dir = tmp_path - graphs = {"flow1": create_real_graph()} - metas = {"flow2": FlowMeta(id="flow2", relative_path="test.json", title="Test")} - - def verbose_print(msg): - pass # Real function - - with pytest.raises(ValueError, match="graphs and metas must contain the same keys"): - create_multi_serve_app(root_dir=root_dir, graphs=graphs, metas=metas, verbose_print=verbose_print) + # Registry always keeps graphs and metas in sync — no mismatch possible + app = create_multi_serve_app(registry=registry) + assert app is not None class TestFastAPIEndpoints: """Test FastAPI endpoints using TestClient.""" - def setup_method(self, tmp_path): + def setup_method(self, tmp_path=None): # noqa: ARG002 """Set up test client with mock data.""" - self.root_dir = tmp_path self.real_graph = create_real_graph() - self.graphs = {"test-flow": self.real_graph} - self.metas = { - "test-flow": FlowMeta( - id="test-flow", relative_path="test.json", title="Test Flow", description="A test flow" - ) - } + meta = FlowMeta(id="test-flow", relative_path="test.json", title="Test Flow", description="A test flow") - def verbose_print(msg): - pass # Real function - - self.verbose_print = verbose_print - - # Create the app first - with patch("lfx.cli.serve_app.verify_api_key"): - self.app = create_multi_serve_app( - root_dir=self.root_dir, graphs=self.graphs, metas=self.metas, verbose_print=self.verbose_print - ) + registry = FlowRegistry() + registry.add(self.real_graph, meta) + self.app = create_multi_serve_app(registry=registry) # Override the dependency for testing def mock_verify_key(): @@ -389,41 +286,41 @@ class TestFastAPIEndpoints: class TestErrorHandling: """Test error handling in various components.""" - def test_invalid_json_in_request(self, tmp_path): + def test_invalid_json_in_request(self): """Test handling of invalid JSON in requests.""" - with patch("lfx.cli.serve_app.verify_api_key", return_value="test-key"): - app = create_multi_serve_app( - root_dir=tmp_path, - graphs={"test": create_real_graph()}, - metas={"test": FlowMeta(id="test", relative_path="test.json", title="Test")}, - verbose_print=lambda msg: None, # noqa: ARG005 - ) - client = TestClient(app) + from lfx.cli.serve_app import verify_api_key - response = client.post( - "/flows/test/run", - data="invalid json", - headers={"x-api-key": "test-key", "Content-Type": "application/json"}, - ) + meta = FlowMeta(id="test", relative_path="test.json", title="Test") + registry = FlowRegistry() + registry.add(create_real_graph(), meta) + app = create_multi_serve_app(registry=registry) + app.dependency_overrides[verify_api_key] = lambda: "test-key" + client = TestClient(app) - assert response.status_code == 422 # Validation error + response = client.post( + "/flows/test/run", + data="invalid json", + headers={"x-api-key": "test-key", "Content-Type": "application/json"}, + ) - def test_missing_flow_id(self, tmp_path): + assert response.status_code == 422 # Validation error + + def test_missing_flow_id(self): """Test accessing non-existent flow.""" - with patch("lfx.cli.serve_app.verify_api_key", return_value="test-key"): - app = create_multi_serve_app( - root_dir=tmp_path, - graphs={"test": create_real_graph()}, - metas={"test": FlowMeta(id="test", relative_path="test.json", title="Test")}, - verbose_print=lambda msg: None, # noqa: ARG005 - ) - client = TestClient(app) + from lfx.cli.serve_app import verify_api_key - response = client.post( - "/flows/nonexistent/run", json={"input_value": "test"}, headers={"x-api-key": "test-key"} - ) + meta = FlowMeta(id="test", relative_path="test.json", title="Test") + registry = FlowRegistry() + registry.add(create_real_graph(), meta) + app = create_multi_serve_app(registry=registry) + app.dependency_overrides[verify_api_key] = lambda: "test-key" + client = TestClient(app) - assert response.status_code == 404 + response = client.post( + "/flows/nonexistent/run", json={"input_value": "test"}, headers={"x-api-key": "test-key"} + ) + + assert response.status_code == 404 class TestIntegration: @@ -462,21 +359,19 @@ class TestIntegration: ) # Test app creation - with patch("lfx.cli.serve_app.verify_api_key", return_value="test-key"): - app = create_multi_serve_app( - root_dir=flow_path.parent, - graphs={flow_id: loaded_graph}, - metas={flow_id: meta}, - verbose_print=mock_verbose_print, - ) + from lfx.cli.serve_app import verify_api_key as _verify_api_key - client = TestClient(app) + registry = FlowRegistry() + registry.add(loaded_graph, meta) + app = create_multi_serve_app(registry=registry) + app.dependency_overrides[_verify_api_key] = lambda: "test-key" + client = TestClient(app) - # Test endpoints - flows_response = client.get("/flows") - assert flows_response.status_code == 200 - assert len(flows_response.json()) == 1 + # Test endpoints + flows_response = client.get("/flows") + assert flows_response.status_code == 200 + assert len(flows_response.json()) == 1 - health_response = client.get("/health") - assert health_response.status_code == 200 - assert health_response.json()["flow_count"] == 1 + health_response = client.get("/health") + assert health_response.status_code == 200 + assert health_response.json()["flow_count"] == 1 diff --git a/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py b/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py index ad67196645..7d74720690 100644 --- a/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py +++ b/src/lfx/tests/unit/custom/component/test_concurrent_tool_invocation.py @@ -220,6 +220,28 @@ def test_should_isolate_inputs_when_input_has_non_picklable_value(): assert product_ids == {"P1", "P2"} +def test_deepcopy_preserves_component_reference_cycles(): + """Component.__deepcopy__ must register itself in memo before copying _components. + + A feedback loop links components through _components in a cycle (A -> B -> A). + If the memo entry is set only after the recursive deepcopy calls, the cycle is + duplicated into a second copy instead of preserved. + """ + a = SlowLabelComponent(_id="a") + b = SlowLabelComponent(_id="b") + a._components = [b] + b._components = [a] + + a_copy = deepcopy(a) + b_copy = a_copy._components[0] + + # The cycle must round-trip back to the same copied objects. + assert b_copy._components[0] is a_copy + # And the copies must be distinct from the originals. + assert a_copy is not a + assert b_copy is not b + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/src/lfx/tests/unit/interface/test_loading_no_env_fallback.py b/src/lfx/tests/unit/interface/test_loading_no_env_fallback.py new file mode 100644 index 0000000000..2e64895f0b --- /dev/null +++ b/src/lfx/tests/unit/interface/test_loading_no_env_fallback.py @@ -0,0 +1,59 @@ +"""Tests for the no_env_fallback guard in load_from_env_vars.""" + +from __future__ import annotations + +import os +from unittest.mock import patch + +from lfx.interface.initialize.loading import load_from_env_vars + + +class TestLoadFromEnvVarsNoFallback: + def test_env_fallback_used_when_flag_absent(self): + """Without the flag, missing request_variables falls back to os.environ.""" + params = {"api_key": "MY_SECRET"} + with patch.dict(os.environ, {"MY_SECRET": "env-value"}): + result = load_from_env_vars(params, ["api_key"], context=None) + assert result["api_key"] == "env-value" + + def test_env_fallback_used_when_flag_false(self): + """Explicit no_env_fallback=False behaves the same as absent.""" + params = {"api_key": "MY_SECRET"} + with patch.dict(os.environ, {"MY_SECRET": "env-value"}): + result = load_from_env_vars(params, ["api_key"], context={"no_env_fallback": False}) + assert result["api_key"] == "env-value" + + def test_env_fallback_skipped_when_flag_true(self): + """With no_env_fallback=True, os.environ is never consulted even if the var is set.""" + params = {"api_key": "MY_SECRET"} + with ( + patch.dict(os.environ, {"MY_SECRET": "env-value"}), + patch("lfx.interface.initialize.loading.os.getenv") as mock_getenv, + ): + result = load_from_env_vars(params, ["api_key"], context={"no_env_fallback": True}) + assert result["api_key"] is None + # Only the credential variable must never be looked up — logger internals may call getenv + credential_lookups = [c for c in mock_getenv.call_args_list if c.args and c.args[0] == "MY_SECRET"] + assert not credential_lookups, f"os.getenv('MY_SECRET') must not be called, got: {credential_lookups}" + + def test_request_variables_win_even_with_flag_true(self): + """request_variables always takes priority, even when no_env_fallback=True.""" + params = {"api_key": "MY_SECRET"} + context = { + "no_env_fallback": True, + "request_variables": {"MY_SECRET": "override-value"}, + } + with patch.dict(os.environ, {"MY_SECRET": "env-value"}): + result = load_from_env_vars(params, ["api_key"], context=context) + assert result["api_key"] == "override-value" + + def test_request_variables_win_when_flag_false(self): + """request_variables takes priority over os.environ even when no_env_fallback=False.""" + params = {"api_key": "MY_SECRET"} + context = { + "no_env_fallback": False, + "request_variables": {"MY_SECRET": "override-value"}, + } + with patch.dict(os.environ, {"MY_SECRET": "env-value"}): + result = load_from_env_vars(params, ["api_key"], context=context) + assert result["api_key"] == "override-value" diff --git a/src/lfx/tests/unit/services/test_minimal_services.py b/src/lfx/tests/unit/services/test_minimal_services.py index 0087df5b6d..314452d813 100644 --- a/src/lfx/tests/unit/services/test_minimal_services.py +++ b/src/lfx/tests/unit/services/test_minimal_services.py @@ -299,6 +299,55 @@ class TestVariableService: finally: del os.environ["LANGFLOW_REQUEST_VARIABLES"] + async def test_global_alias_from_request_scope(self, variables): + """x-langflow-global-var-* aliases resolve from request-scoped variables.""" + from lfx.services.variable.request_scope import activate_request_variables, reset_request_variables + + token = activate_request_variables({"x-langflow-global-var-access-token": "alias-token"}) + try: + assert await variables.get_variable("access_token") == "alias-token" + finally: + reset_request_variables(token) + + async def test_request_scope_overrides_env(self, variables): + """A request-scoped variable wins over an env var of the same name (core feature).""" + from lfx.services.variable.request_scope import activate_request_variables, reset_request_variables + + os.environ["SHARED_VAR"] = "env-value" + token = activate_request_variables({"SHARED_VAR": "request-value"}) + try: + assert await variables.get_variable("SHARED_VAR") == "request-value" + finally: + reset_request_variables(token) + del os.environ["SHARED_VAR"] + + async def test_request_scope_alias_overrides_env(self, variables): + """A request-scoped x-langflow-global-var-* alias wins over an env var of the same name. + + Ensures a caller's per-request credential is never shadowed by an ambient process + credential, regardless of which form (exact name or alias) the caller supplied. + """ + from lfx.services.variable.request_scope import activate_request_variables, reset_request_variables + + os.environ["ACCESS_TOKEN"] = "env-token" # noqa: S105 + token = activate_request_variables({"x-langflow-global-var-access-token": "alias-token"}) + try: + assert await variables.get_variable("ACCESS_TOKEN") == "alias-token" + finally: + reset_request_variables(token) + del os.environ["ACCESS_TOKEN"] + + async def test_in_memory_overrides_request_scope(self, variables): + """An in-memory variable wins over a request-scoped variable of the same name.""" + from lfx.services.variable.request_scope import activate_request_variables, reset_request_variables + + variables.set_variable("SHARED_VAR", "memory-value") + token = activate_request_variables({"SHARED_VAR": "request-value"}) + try: + assert await variables.get_variable("SHARED_VAR") == "memory-value" + finally: + reset_request_variables(token) + @pytest.mark.asyncio async def test_langflow_request_variables_invalid_json_falls_back(self, variables): """Test invalid request variable JSON does not break env fallback.""" diff --git a/src/lfx/tests/unit/services/test_request_scope_isolation.py b/src/lfx/tests/unit/services/test_request_scope_isolation.py new file mode 100644 index 0000000000..14c5e24f89 --- /dev/null +++ b/src/lfx/tests/unit/services/test_request_scope_isolation.py @@ -0,0 +1,216 @@ +"""Isolation guarantees for request-scoped variables (lfx serve concurrency). + +These tests verify the two properties that matter for running multiple flows in a +single ``lfx serve`` process: + +1. **Per-request isolation** - one request's ``global_vars`` must never be visible + to a concurrently-running request. This is provided by the ContextVar in + :mod:`lfx.services.variable.request_scope`: each asyncio task gets its own + context copy, so activating a scope in one task cannot leak into another. +2. **No accidental env fallback for supplied credentials** - when a request scope + provides a variable, resolution must use it and must not consult ``os.environ`` + for that name (so a stale/foreign env var can never shadow a request credential). +""" + +from __future__ import annotations + +import asyncio +import os +from unittest.mock import patch + +from lfx.services.variable.request_scope import ( + activate_no_env_fallback, + activate_request_variables, + get_active_request_variables, + reset_no_env_fallback, + reset_request_variables, +) +from lfx.services.variable.service import VariableService + + +async def test_concurrent_request_scopes_are_isolated(): + """Two concurrent requests each see only their own request-scoped variable.""" + service = VariableService() + observed: dict[str, str | None] = {} + + async def handle(request_id: str, secret: str) -> None: + token = activate_request_variables({"access_token": secret}) + try: + # Yield so the sibling task activates its own scope before we read; + # if the ContextVar leaked across tasks this read would see it. + await asyncio.sleep(0.01) + observed[request_id] = await service.get_variable("access_token") + # Read again after another yield to catch any cross-task overwrite. + await asyncio.sleep(0.01) + observed[f"{request_id}-recheck"] = await service.get_variable("access_token") + finally: + reset_request_variables(token) + + await asyncio.gather(handle("A", "secret-A"), handle("B", "secret-B")) + + assert observed["A"] == "secret-A" + assert observed["A-recheck"] == "secret-A" + assert observed["B"] == "secret-B" + assert observed["B-recheck"] == "secret-B" + + +async def test_unscoped_request_does_not_see_a_concurrent_scope(): + """A request with no active scope must not observe another request's live scope.""" + service = VariableService() + scope_active = asyncio.Event() + scope_may_exit = asyncio.Event() + observed: dict[str, str | None] = {} + + async def scoped() -> None: + token = activate_request_variables({"access_token": "secret-A"}) + try: + scope_active.set() + await scope_may_exit.wait() # hold the scope open while the other task reads + finally: + reset_request_variables(token) + + async def unscoped() -> None: + await scope_active.wait() # read while scoped()'s scope is definitely active + observed["value"] = await service.get_variable("access_token") + scope_may_exit.set() + + await asyncio.gather(scoped(), unscoped()) + + assert observed["value"] is None + + +async def test_scope_does_not_leak_after_reset(): + """After reset, the variable is no longer resolvable (request teardown is clean).""" + service = VariableService() + + token = activate_request_variables({"access_token": "secret"}) + assert await service.get_variable("access_token") == "secret" + reset_request_variables(token) + + assert get_active_request_variables() is None + assert await service.get_variable("access_token") is None + + +async def test_scoped_value_used_without_consulting_env_for_that_name(): + """A request-supplied variable is returned without falling back to os.environ. + + Guards against a regression where the env tier runs first (or in addition), + which would let a process env var shadow a per-request credential. + """ + service = VariableService() + token = activate_request_variables({"access_token": "from-request"}) + with ( + patch.dict(os.environ, {"access_token": "from-env"}), + patch("lfx.services.variable.service.os.getenv") as mock_getenv, + ): + try: + assert await service.get_variable("access_token") == "from-request" + finally: + reset_request_variables(token) + + name_lookups = [c for c in mock_getenv.call_args_list if c.args and c.args[0] == "access_token"] + assert not name_lookups, f"os.getenv('access_token') must not be consulted, got: {name_lookups}" + + +async def test_env_fallback_enabled_by_default(): + """Without the no-env-fallback flag, env resolution still works (default preserved).""" + service = VariableService() + with patch.dict(os.environ, {"DEFAULT_SECRET": "from-env"}): + assert await service.get_variable("DEFAULT_SECRET") == "from-env" + + +async def test_no_env_fallback_suppresses_env_in_variable_service(): + """With env fallback disabled, a variable only in os.environ resolves to None. + + Closes the gap where VariableService (model/KB credential path) ignored + ``no_env_fallback`` and leaked process env credentials even under --no-env-fallback. + """ + service = VariableService() + scope_token = activate_request_variables({"access_token": "from-request"}) + flag_token = activate_no_env_fallback(disabled=True) + with ( + patch.dict(os.environ, {"OTHER_SECRET": "from-env"}), + patch("lfx.services.variable.service.os.getenv") as mock_getenv, + ): + try: + # Supplied via the request scope -> still resolves. + assert await service.get_variable("access_token") == "from-request" + # Only in process env -> must NOT resolve when fallback is disabled. + assert await service.get_variable("OTHER_SECRET") is None + finally: + reset_no_env_fallback(flag_token) + reset_request_variables(scope_token) + + assert not mock_getenv.call_args_list, f"os.getenv must not be consulted, got: {mock_getenv.call_args_list}" + + +async def test_no_env_fallback_suppresses_langflow_request_variables_env(): + """Env fallback disabled + no active scope must not leak the process LANGFLOW_REQUEST_VARIABLES blob. + + Regression: _get_request_variables() read os.getenv("LANGFLOW_REQUEST_VARIABLES") + before the no-env-fallback gate, so an empty-global_vars request under + --no-env-fallback (which binds the scope to None) still resolved credentials + from the process env blob, violating the "never reads os.environ" contract. + """ + service = VariableService() + scope_token = activate_request_variables(None) # None == empty global_vars in serve + flag_token = activate_no_env_fallback(disabled=True) + with patch.dict(os.environ, {"LANGFLOW_REQUEST_VARIABLES": '{"leaked_token": "SHOULD-NOT-LEAK"}'}): + try: + assert await service.get_variable("leaked_token") is None + finally: + reset_no_env_fallback(flag_token) + reset_request_variables(scope_token) + + +async def test_langflow_request_variables_env_blob_used_when_fallback_enabled(): + """Env blob resolves when fallback is enabled: null dropped, structured -> valid JSON string.""" + service = VariableService() + scope_token = activate_request_variables(None) + blob = '{"shared_token": "from-env-blob", "null_cred": null, "nested": {"a": 1}}' + with patch.dict(os.environ, {"LANGFLOW_REQUEST_VARIABLES": blob}): + try: + assert await service.get_variable("shared_token") == "from-env-blob" + # null is dropped (not the truthy string "None"); env/None fallthrough -> None. + assert await service.get_variable("null_cred") is None + # Structured value serialized as valid JSON, round-trippable via json.loads. + assert await service.get_variable("nested") == '{"a": 1}' + finally: + reset_request_variables(scope_token) + + +async def test_no_env_fallback_flag_is_isolated_per_request(): + """The no-env-fallback flag is per-request. + + One request disabling it must not affect a concurrent request that allows + env fallback. + """ + service = VariableService() + observed: dict[str, str | None] = {} + + async def strict_request() -> None: + # Active (non-None) scope so resolution never reads LANGFLOW_REQUEST_VARIABLES. + scope = activate_request_variables({"unused": "x"}) + flag = activate_no_env_fallback(disabled=True) + try: + await asyncio.sleep(0.01) # let the lenient request run concurrently + observed["strict"] = await service.get_variable("SHARED_ENV") + finally: + reset_no_env_fallback(flag) + reset_request_variables(scope) + + async def lenient_request() -> None: + scope = activate_request_variables({"unused": "y"}) + flag = activate_no_env_fallback(disabled=False) + try: + await asyncio.sleep(0.01) + observed["lenient"] = await service.get_variable("SHARED_ENV") + finally: + reset_no_env_fallback(flag) + reset_request_variables(scope) + + with patch.dict(os.environ, {"SHARED_ENV": "from-env"}): + await asyncio.gather(strict_request(), lenient_request()) + + assert observed["strict"] is None # env fallback disabled for this request + assert observed["lenient"] == "from-env" # env fallback allowed for this request diff --git a/src/lfx/tests/unit/utils/test_async_helpers.py b/src/lfx/tests/unit/utils/test_async_helpers.py new file mode 100644 index 0000000000..1924744801 --- /dev/null +++ b/src/lfx/tests/unit/utils/test_async_helpers.py @@ -0,0 +1,44 @@ +"""Tests for ``run_until_complete`` contextvars propagation. + +``run_until_complete`` has two branches: ``asyncio.run`` (no running loop) and a +worker-thread fallback (a loop is already running). The fallback must carry the +caller's contextvars into the new thread — lfx serve resolves request-scoped +variables and the no-env-fallback flag via ContextVars, and a bare +``ThreadPoolExecutor`` would otherwise reset them to their defaults. +""" + +from __future__ import annotations + +import contextvars + +from lfx.utils.async_helpers import run_until_complete + +_probe_var: contextvars.ContextVar[str] = contextvars.ContextVar("probe_var", default="default") + + +async def test_run_until_complete_propagates_contextvars_in_thread_fallback(): + """When a loop is already running, the worker-thread branch keeps the caller's ContextVars.""" + + async def _read() -> str: + return _probe_var.get() + + token = _probe_var.set("propagated") + try: + # pytest-asyncio runs this test inside a loop, so run_until_complete takes the + # thread-hop branch. Before the fix this returned "default" (context reset). + assert run_until_complete(_read()) == "propagated" + finally: + _probe_var.reset(token) + + +def test_run_until_complete_propagates_contextvars_in_run_branch(): + """With no running loop, the asyncio.run branch also sees the caller's ContextVars.""" + + async def _read() -> str: + return _probe_var.get() + + token = _probe_var.set("sync-path") + try: + assert run_until_complete(_read()) == "sync-path" + finally: + _probe_var.reset(token)