Merge branch 'release-1.10.0' into fix/mask-global-variable-output-exposure

This commit is contained in:
Eric Hare
2026-04-27 14:32:30 -07:00
committed by GitHub
96 changed files with 7820 additions and 376 deletions

View File

@ -8480,5 +8480,5 @@
}
]
},
"generated_at": "2026-04-23T21:12:19Z"
"generated_at": "2026-04-24T03:31:30Z"
}

View File

@ -39,9 +39,28 @@ This step also usually lasts about a week.
After QA and bugfixing are complete for both OSS and Desktop:
* Final releases are cut from their respective RC branches.
* Release timing is coordinated with Langflows DevRel team.
* Release timing is coordinated with Langflow's DevRel team.
* For at least 24 hours after release, Discord, GitHub, and other support channels should be monitored for critical bug reports.
### 4. Release Artifacts
The release workflow automatically publishes the following artifacts:
* **PyPI Packages:**
* `langflow` - Main package with all integrations
* `langflow-base` - Core framework without integrations
* `lfx` - Lightweight executor CLI
* `langflow-sdk` - SDK for programmatic access (when updated)
* **Docker Images:**
* `langflowai/langflow` - Full Langflow image
* `langflowai/langflow-backend` - Backend-only image (published independently)
* `langflowai/langflow-frontend` - Frontend-only image (published independently)
* `langflowai/langflow-ep` - Enterprise edition image (published independently)
* `langflowai/langflow-base` - Base image without integrations
**Note:** Backend, frontend, and enterprise images are published separately from the main image and will be built even if the main version already exists on Docker Hub.
## Branch Model
| Branch | Purpose | Merge Policy |
@ -101,6 +120,10 @@ git merge --ff-only release-X.Y.Z # Fast-forward main to include RC changes
* Follows [Semantic Versioning](https://semver.org): `MAJOR.MINOR.PATCH`.
* RC tags use `-rc.N`, e.g. `v1.8.0-rc.1`.
* **All tags MUST start with `v` prefix** (e.g., `v1.9.1`, not `1.9.1`).
* The release workflow validates this format and rejects tags without the `v` prefix.
* Duplicate tags (e.g., both `1.8.3` and `v1.8.3`) cause GitHub's release notes generation to use the wrong base comparison, resulting in incomplete changelogs.
* The workflow automatically checks for and prevents duplicate tags.
## Roles

View File

@ -131,6 +131,8 @@ langflow:
enabled: true
driver:
value: "postgresql"
host:
value: "postgresql-svc.langflow.svc.cluster.local"
port:
value: "5432"
user:

View File

@ -131,6 +131,8 @@ langflow:
enabled: true
driver:
value: "postgresql"
host:
value: "postgresql-svc.langflow.svc.cluster.local"
port:
value: "5432"
user:

View File

@ -131,6 +131,8 @@ langflow:
enabled: true
driver:
value: "postgresql"
host:
value: "postgresql-svc.langflow.svc.cluster.local"
port:
value: "5432"
user:

View File

@ -0,0 +1,235 @@
"""add_memory_base_schema
Consolidates all Memory Base schema changes into a single migration:
- job.dedupe_key (nullable String) + ix_job_dedupe_key
- message.run_id (nullable UUID) + ix_message_run_id
- message.is_output (bool, default false)
- memory_base table + ix_memory_base_flow_id + ix_memory_base_user_id
- memory_base_session table + three indexes
- message_ingestion_record table + three indexes
- memory_base_workflow_run table + two indexes
Phase: EXPAND
Revision ID: mb00a1b2c3d4
Revises: d306e5c17c41
Create Date: 2026-04-14 00:00:00.000000
"""
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
from langflow.utils import migration
# revision identifiers, used by Alembic.
revision: str = "mb00a1b2c3d4" # pragma: allowlist secret
down_revision: str | None = "d306e5c17c41" # pragma: allowlist secret
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
conn = op.get_bind()
# ------------------------------------------------------------------ #
# job.dedupe_key #
# ------------------------------------------------------------------ #
inspector = sa.inspect(conn)
existing_job_indexes = {idx["name"] for idx in inspector.get_indexes("job")}
with op.batch_alter_table("job", schema=None) as batch_op:
if not migration.column_exists("job", "dedupe_key", conn):
batch_op.add_column(sa.Column("dedupe_key", sa.String(), nullable=True))
if "ix_job_dedupe_key" not in existing_job_indexes:
batch_op.create_index(batch_op.f("ix_job_dedupe_key"), ["dedupe_key"], unique=False)
# ------------------------------------------------------------------ #
# message.run_id + message.is_output #
# ------------------------------------------------------------------ #
with op.batch_alter_table("message", schema=None) as batch_op:
if not migration.column_exists("message", "run_id", conn):
batch_op.add_column(sa.Column("run_id", sa.Uuid(), nullable=True))
if not migration.column_exists("message", "is_output", conn):
batch_op.add_column(sa.Column("is_output", sa.Boolean(), nullable=False, server_default=sa.text("false")))
existing_message_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("message")}
if "ix_message_run_id" not in existing_message_indexes:
op.create_index("ix_message_run_id", "message", ["run_id"])
# ------------------------------------------------------------------ #
# memory_base #
# ------------------------------------------------------------------ #
if not migration.table_exists("memory_base", conn):
op.create_table(
"memory_base",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column("name", sa.String(), nullable=False),
sa.Column("flow_id", sa.Uuid(), nullable=False),
sa.Column("user_id", sa.Uuid(), nullable=False),
sa.Column("threshold", sa.Integer(), nullable=False, server_default=sa.text("50")),
sa.Column("auto_capture", sa.Boolean(), nullable=False, server_default=sa.text("true")),
sa.Column("embedding_model", sa.String(), nullable=False, server_default=sa.text("''")),
sa.Column("preprocessing", sa.Boolean(), nullable=False, server_default=sa.text("false")),
sa.Column("preproc_model", sa.String(), nullable=True),
sa.Column("preproc_instructions", sa.String(), nullable=True),
sa.Column("kb_name", sa.String(), nullable=False),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint("user_id", "name", name="uq_memory_base_user_name"),
sa.Index("ix_memory_base_flow_id", "flow_id"),
sa.Index("ix_memory_base_user_id", "user_id"),
)
# ------------------------------------------------------------------ #
# memory_base_session #
# ------------------------------------------------------------------ #
if not migration.table_exists("memory_base_session", conn):
op.create_table(
"memory_base_session",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column(
"memory_base_id",
sa.Uuid(),
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
),
sa.Column("session_id", sa.String(), nullable=False),
sa.Column("cursor_id", sa.Uuid(), nullable=True),
sa.Column("total_processed", sa.Integer(), nullable=False, server_default=sa.text("0")),
sa.Column("last_sync_at", sa.DateTime(timezone=True), nullable=True),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint("memory_base_id", "session_id", name="uq_memory_base_session"),
)
op.create_index("ix_memory_base_session_memory_base_id", "memory_base_session", ["memory_base_id"])
op.create_index("ix_memory_base_session_session_id", "memory_base_session", ["session_id"])
op.create_index(
"ix_memory_base_session_lookup",
"memory_base_session",
["memory_base_id", "session_id"],
)
# ------------------------------------------------------------------ #
# message_ingestion_record #
# ------------------------------------------------------------------ #
if not migration.table_exists("message_ingestion_record", conn):
op.create_table(
"message_ingestion_record",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column(
"message_id",
sa.Uuid(),
sa.ForeignKey("message.id", ondelete="CASCADE"),
nullable=False,
),
sa.Column(
"memory_base_id",
sa.Uuid(),
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
),
sa.Column(
"job_id",
sa.Uuid(),
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
sa.Column("session_id", sa.String(), nullable=False),
sa.Column("ingested_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"message_id",
"session_id",
"memory_base_id",
name="uq_mir_message_session_mb",
),
)
op.create_index("ix_mir_message_id", "message_ingestion_record", ["message_id"])
op.create_index("ix_mir_job_id", "message_ingestion_record", ["job_id"])
op.create_index(
"ix_mir_memory_base_session",
"message_ingestion_record",
["memory_base_id", "session_id"],
)
# ------------------------------------------------------------------ #
# memory_base_workflow_run #
# ------------------------------------------------------------------ #
if not migration.table_exists("memory_base_workflow_run", conn):
op.create_table(
"memory_base_workflow_run",
sa.Column("id", sa.Uuid(), nullable=False),
sa.Column(
"memory_base_id",
sa.Uuid(),
sa.ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
),
sa.Column("session_id", sa.String(), nullable=False),
sa.Column(
"workflow_job_id",
sa.Uuid(),
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
sa.Column(
"ingestion_job_id",
sa.Uuid(),
sa.ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
sa.Column("recorded_at", sa.DateTime(timezone=True), nullable=False),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"memory_base_id",
"session_id",
"workflow_job_id",
name="uq_mbwr_mb_session_wf_job",
),
)
op.create_index("ix_mbwr_mb_session", "memory_base_workflow_run", ["memory_base_id", "session_id"])
op.create_index("ix_mbwr_ingestion_job_id", "memory_base_workflow_run", ["ingestion_job_id"])
def downgrade() -> None:
conn = op.get_bind()
# Children first (FK dependencies) ----------------------------------- #
if migration.table_exists("memory_base_workflow_run", conn):
op.drop_index("ix_mbwr_ingestion_job_id", table_name="memory_base_workflow_run")
op.drop_index("ix_mbwr_mb_session", table_name="memory_base_workflow_run")
op.drop_table("memory_base_workflow_run")
if migration.table_exists("message_ingestion_record", conn):
op.drop_index("ix_mir_memory_base_session", table_name="message_ingestion_record")
op.drop_index("ix_mir_job_id", table_name="message_ingestion_record")
op.drop_index("ix_mir_message_id", table_name="message_ingestion_record")
op.drop_table("message_ingestion_record")
if migration.table_exists("memory_base_session", conn):
op.drop_index("ix_memory_base_session_lookup", table_name="memory_base_session")
op.drop_index("ix_memory_base_session_session_id", table_name="memory_base_session")
op.drop_index("ix_memory_base_session_memory_base_id", table_name="memory_base_session")
op.drop_table("memory_base_session")
if migration.table_exists("memory_base", conn):
op.drop_index("ix_memory_base_user_id", table_name="memory_base")
op.drop_index("ix_memory_base_flow_id", table_name="memory_base")
op.drop_table("memory_base")
# Message column/index ----------------------------------------------- #
existing_message_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("message")}
if "ix_message_run_id" in existing_message_indexes:
op.drop_index("ix_message_run_id", table_name="message")
with op.batch_alter_table("message", schema=None) as batch_op:
if migration.column_exists("message", "is_output", conn):
batch_op.drop_column("is_output")
if migration.column_exists("message", "run_id", conn):
batch_op.drop_column("run_id")
# Job column/index --------------------------------------------------- #
with op.batch_alter_table("job", schema=None) as batch_op:
existing_job_indexes = {idx["name"] for idx in sa.inspect(conn).get_indexes("job")}
if "ix_job_dedupe_key" in existing_job_indexes:
batch_op.drop_index(batch_op.f("ix_job_dedupe_key"))
if migration.column_exists("job", "dedupe_key", conn):
batch_op.drop_column("dedupe_key")

View File

@ -29,7 +29,15 @@ from langflow.exceptions.component import ComponentBuildError
from langflow.schema.message import ErrorMessage
from langflow.schema.schema import OutputValue
from langflow.services.database.models.flow.model import Flow
from langflow.services.deps import get_chat_service, get_telemetry_service, session_scope
from langflow.services.database.models.jobs.model import JobType
from langflow.services.deps import (
get_chat_service,
get_job_service,
get_memory_base_service,
get_task_service,
get_telemetry_service,
session_scope,
)
from langflow.services.job_queue.service import JobQueueNotFoundError, JobQueueService
from langflow.services.telemetry.schema import ComponentInputsPayload, ComponentPayload, PlaygroundPayload
@ -527,35 +535,81 @@ async def generate_flow_events(
event_manager.on_error(data=error_message.data)
raise
# Create a WORKFLOW job record so memory-base on_flow_output can track this run.
# Best-effort: failures here must never break the build path.
_build_job_svc = None
_build_run_id: uuid.UUID | None = None
try:
_build_run_id = uuid.UUID(graph.run_id) if graph.run_id else None
if _build_run_id is not None:
_build_job_svc = get_job_service()
await _build_job_svc.create_job(
job_id=_build_run_id,
flow_id=flow_id,
user_id=current_user.id,
job_type=JobType.WORKFLOW,
)
except Exception: # noqa: BLE001
await logger.awarning(
"Failed to create workflow job for /build — memory base tracking disabled for flow %s",
flow_id,
exc_info=True,
)
_build_job_svc = None
event_manager.on_vertices_sorted(data={"ids": ids, "to_run": vertices_to_run})
vertex_timedeltas: list[float] = []
event_manager.on_build_start(data={})
tasks = []
for vertex_id in ids:
task = asyncio.create_task(build_vertices(vertex_id, graph, event_manager, vertex_timedeltas))
tasks.append(task)
try:
await asyncio.gather(*tasks)
except asyncio.CancelledError:
background_tasks.add_task(graph.end_all_traces_in_context())
raise
except Exception as e:
await logger.aerror(f"Error building vertices: {e}")
custom_component = graph.get_vertex(vertex_id).custom_component
trace_name = getattr(custom_component, "trace_name", None)
error_message = ErrorMessage(
flow_id=flow_id,
exception=e,
session_id=graph.session_id,
trace_name=trace_name,
)
event_manager.on_error(data=error_message.data)
raise
async def _run_vertex_build() -> None:
tasks = []
for vertex_id in ids:
task = asyncio.create_task(build_vertices(vertex_id, graph, event_manager, vertex_timedeltas))
tasks.append(task)
try:
await asyncio.gather(*tasks)
except asyncio.CancelledError:
background_tasks.add_task(graph.end_all_traces_in_context())
raise
except Exception as e:
await logger.aerror(f"Error building vertices: {e}")
custom_component = graph.get_vertex(vertex_id).custom_component
trace_name = getattr(custom_component, "trace_name", None)
error_message = ErrorMessage(
flow_id=flow_id,
exception=e,
session_id=graph.session_id,
trace_name=trace_name,
)
event_manager.on_error(data=error_message.data)
raise
if _build_job_svc and _build_run_id:
await _build_job_svc.execute_with_status(_build_run_id, _run_vertex_build)
else:
await _run_vertex_build()
build_duration = sum(vertex_timedeltas)
event_manager.on_end(data={"build_duration": build_duration})
await graph.end_all_traces()
# Fire memory-base auto-capture hook — non-blocking background effect.
# Must use fire_and_forget_task (not background_tasks.add_task) because
# generate_flow_events runs as an asyncio task; by the time the flow
# finishes, FastAPI has already drained the background_tasks queue and any
# tasks added after that point are silently dropped.
try:
_run_id_uuid = uuid.UUID(graph.run_id) if graph.run_id else None # type-cast only; same run_id set on graph
await get_task_service().fire_and_forget_task(
get_memory_base_service().on_flow_output,
flow_id=flow_id,
session_id=graph.session_id or str(flow_id),
job_id=_run_id_uuid,
)
except (RuntimeError, ValueError, OSError):
await logger.awarning("Memory base hook scheduling failed for flow %s", flow_id, exc_info=True)
await event_manager.queue.put((None, None, time.time()))

View File

@ -15,6 +15,7 @@ from langflow.api.v1 import (
login_router,
mcp_projects_router,
mcp_router,
memories_router,
model_options_router,
models_router,
monitor_router,
@ -68,6 +69,7 @@ router_v1.include_router(folders_router)
router_v1.include_router(projects_router)
router_v1.include_router(starter_projects_router)
router_v1.include_router(knowledge_bases_router)
router_v1.include_router(memories_router)
router_v1.include_router(mcp_router)
router_v1.include_router(voice_mode_router)
router_v1.include_router(mcp_projects_router)

View File

@ -445,7 +445,7 @@ class KBIngestionHelper:
splitter_kwargs["separators"] = [resolved_separator]
text_splitter = RecursiveCharacterTextSplitter(**splitter_kwargs)
embeddings = await KBIngestionHelper._build_embeddings(embedding_provider, embedding_model, current_user)
embeddings = await KBIngestionHelper.build_embeddings(embedding_provider, embedding_model, current_user)
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
chroma = Chroma(
@ -462,40 +462,30 @@ class KBIngestionHelper:
continue
chunks = text_splitter.split_text(content)
for i in range(0, len(chunks), INGESTION_BATCH_SIZE):
if await KBIngestionHelper._is_job_cancelled(job_service, task_job_id):
raise IngestionCancelledError
docs = [
Document(
page_content=c,
metadata={
"source": source_name or file_name,
"file_name": file_name,
"chunk_index": i,
"total_chunks": len(chunks),
"ingested_at": datetime.now(timezone.utc).isoformat(),
"job_id": job_id_str,
},
)
for i, c in enumerate(chunks)
]
batch = chunks[i : i + INGESTION_BATCH_SIZE]
docs = [
Document(
page_content=c,
metadata={
"source": source_name or file_name,
"file_name": file_name,
"chunk_index": i + j,
"total_chunks": len(chunks),
"ingested_at": datetime.now(timezone.utc).isoformat(),
"job_id": job_id_str,
},
)
for j, c in enumerate(batch)
]
for attempt in range(MAX_RETRY_ATTEMPTS):
if await KBIngestionHelper._is_job_cancelled(job_service, task_job_id):
raise IngestionCancelledError
try:
await chroma.aadd_documents(docs)
break
except Exception as e:
if attempt == MAX_RETRY_ATTEMPTS - 1:
raise
wait = (attempt + 1) * EXPONENTIAL_BACKOFF_MULTIPLIER
await logger.awarning("Write failed, retrying in %ds: %s", wait, e)
await asyncio.sleep(wait)
await asyncio.sleep(0.01)
written = await KBIngestionHelper.write_documents_to_chroma(
documents=docs,
chroma=chroma,
task_job_id=task_job_id,
job_service=job_service,
)
if written < len(docs):
# Job was cancelled mid-file
raise IngestionCancelledError
total_chunks_created += len(chunks)
processed_files.append(file_name)
@ -555,13 +545,69 @@ class KBIngestionHelper:
KBStorageHelper.release_chroma_resources(kb_path)
@staticmethod
async def _is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool:
async def write_documents_to_chroma(
*,
documents: list[Document],
chroma: Chroma,
task_job_id: uuid.UUID,
job_service: JobService,
) -> int:
"""Write pre-built Documents into an open Chroma collection.
This is the shared primitive used by both file-based KB ingestion
(``perform_ingestion``) and message-based Memory Base ingestion.
Documents must already be chunked and have their metadata populated
by the caller — this method only handles the batched write, cancellation
checking, and retry logic.
Args:
documents: LangChain Document objects ready for embedding.
chroma: An already-constructed ``Chroma`` instance pointing at the
target collection.
task_job_id: Job ID used to poll for cancellation.
job_service: Service for checking job status.
Returns:
Number of documents successfully written. If the job is cancelled
mid-batch this will be less than ``len(documents)``.
Raises:
Exception: Re-raises any non-cancellation write failure after the
retry budget is exhausted.
"""
written = 0
for i in range(0, len(documents), INGESTION_BATCH_SIZE):
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
return written
batch = documents[i : i + INGESTION_BATCH_SIZE]
for attempt in range(MAX_RETRY_ATTEMPTS):
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
return written
try:
await chroma.aadd_documents(batch)
break
except Exception as e:
if attempt == MAX_RETRY_ATTEMPTS - 1:
raise
wait = (attempt + 1) * EXPONENTIAL_BACKOFF_MULTIPLIER
await logger.awarning("Write failed, retrying in %ds: %s", wait, e)
await asyncio.sleep(wait)
written += len(batch)
await asyncio.sleep(0.01)
return written
@staticmethod
async def is_job_cancelled(job_service: JobService, job_id: uuid.UUID) -> bool:
"""Internal helper to check if a job has been cancelled."""
job = await job_service.get_job_by_job_id(job_id)
return job is not None and job.status == JobStatus.CANCELLED
@staticmethod
async def _build_embeddings(provider: str, model: str, current_user):
async def build_embeddings(provider: str, model: str, current_user):
"""Internal helper to build embeddings object."""
options = get_embedding_model_options(user_id=current_user.id)
selected_option = next((o for o in options if o["provider"] == provider and o["name"] == model), None)

View File

@ -10,6 +10,7 @@ from langflow.api.v1.knowledge_bases import router as knowledge_bases_router
from langflow.api.v1.login import router as login_router
from langflow.api.v1.mcp import router as mcp_router
from langflow.api.v1.mcp_projects import router as mcp_projects_router
from langflow.api.v1.memories import router as memories_router
from langflow.api.v1.model_options import router as model_options_router
from langflow.api.v1.models import router as models_router
from langflow.api.v1.monitor import router as monitor_router
@ -36,6 +37,7 @@ __all__ = [
"login_router",
"mcp_projects_router",
"mcp_router",
"memories_router",
"model_options_router",
"models_router",
"monitor_router",

View File

@ -6,7 +6,7 @@ import time
from collections.abc import AsyncGenerator
from http import HTTPStatus
from typing import TYPE_CHECKING, Annotated
from uuid import uuid4
from uuid import UUID, uuid4
import orjson
import sqlalchemy as sa
@ -62,8 +62,17 @@ from langflow.services.auth.utils import (
from langflow.services.cache.utils import save_uploaded_file
from langflow.services.database.models.flow.model import Flow, FlowRead
from langflow.services.database.models.flow.utils import get_all_webhook_components_in_flow
from langflow.services.database.models.jobs.model import JobType
from langflow.services.database.models.user.model import User, UserRead
from langflow.services.deps import get_auth_service, get_session_service, get_settings_service, get_telemetry_service
from langflow.services.deps import (
get_auth_service,
get_job_service,
get_memory_base_service,
get_session_service,
get_settings_service,
get_task_service,
get_telemetry_service,
)
from langflow.services.event_manager import create_webhook_event_manager, webhook_event_manager
from langflow.services.telemetry.schema import RunPayload
from langflow.utils.compression import compress_response
@ -173,8 +182,8 @@ async def simple_run_flow(
graph = Graph.from_payload(
graph_data, flow_id=flow_id_str, user_id=str(user_id), flow_name=flow.name, context=context
)
if run_id is None:
run_id = str(uuid4())
run_id_uuid = uuid4() if run_id is None else UUID(run_id)
run_id = str(run_id_uuid)
graph.set_run_id(run_id)
inputs = None
if input_request.input_value is not None:
@ -197,15 +206,57 @@ async def simple_run_flow(
and (input_request.output_type == "any" or input_request.output_type in vertex.id.lower()) # type: ignore[operator]
)
]
task_result, session_id = await run_graph_internal(
graph=graph,
flow_id=flow_id_str,
session_id=input_request.session_id,
inputs=inputs,
outputs=outputs,
stream=stream,
event_manager=event_manager,
)
# Create a WORKFLOW job record so memory-base on_flow_output can track this run.
if user_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Authentication required to run flows.",
)
try:
_job_svc = get_job_service()
await _job_svc.create_job(
job_id=run_id_uuid,
flow_id=flow.id,
user_id=user_id,
job_type=JobType.WORKFLOW,
)
task_result, session_id = await _job_svc.execute_with_status(
run_id_uuid,
run_graph_internal,
graph=graph,
flow_id=flow_id_str,
session_id=input_request.session_id,
inputs=inputs,
outputs=outputs,
stream=stream,
event_manager=event_manager,
)
except Exception as exc:
await logger.aerror(
"Workflow job execution failed for flow %s: %s",
flow.id,
str(exc),
exc_info=True,
)
raise APIException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
exception=exc,
flow=flow,
) from exc
# Fire memory-base auto-capture hook — non-blocking background effect.
try:
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only
await get_task_service().fire_and_forget_task(
get_memory_base_service().on_flow_output,
flow_id=flow.id,
session_id=session_id,
job_id=_run_id_uuid,
)
except (RuntimeError, ValueError, OSError):
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
return RunResponse(outputs=task_result, session_id=session_id)
@ -970,6 +1021,18 @@ async def experimental_run_flow(
except Exception as exc:
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(exc)) from exc
# Fire memory-base auto-capture hook — non-blocking background effect.
try:
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only
await get_task_service().fire_and_forget_task(
get_memory_base_service().on_flow_output,
flow_id=flow.id,
session_id=session_id,
job_id=_run_id_uuid,
)
except (RuntimeError, ValueError, OSError):
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
return RunResponse(outputs=task_result, session_id=session_id)

View File

@ -74,6 +74,30 @@ def _resolve_kb_path(kb_name: str, current_user: CurrentActiveUser) -> Path:
return kb_path
def _is_memory_base_associated(metadata: dict[str, Any]) -> bool:
"""Return True if the KB metadata indicates an association with a Memory Base."""
source_types = metadata.get("source_types")
return isinstance(source_types, list) and "memory" in source_types
def _check_memory_base_association(kb_name: str, current_user: CurrentActiveUser) -> None:
"""Raise 403 if the KB is associated with a Memory Base.
Designed as a FastAPI dependency for per-KB routes — FastAPI injects
``kb_name`` from the path parameter and ``current_user`` via its own
dependency. The list endpoint filters memory KBs inline using
``_is_memory_base_associated`` directly.
"""
kb_path = _resolve_kb_path(kb_name, current_user)
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
if _is_memory_base_associated(metadata):
raise HTTPException(
status_code=403,
detail=f"Access denied: knowledge base '{kb_name}' is managed by a Memory Base.",
)
@router.post("", status_code=HTTPStatus.CREATED)
@router.post("/", status_code=HTTPStatus.CREATED)
async def create_knowledge_base(
@ -275,7 +299,7 @@ async def preview_chunks(
return {"files": file_previews}
@router.post("/{kb_name}/ingest", status_code=HTTPStatus.OK)
@router.post("/{kb_name}/ingest", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
async def ingest_files_to_knowledge_base(
kb_name: str,
current_user: CurrentActiveUser,
@ -427,6 +451,8 @@ async def list_knowledge_bases(
try:
# Use deep update (fast=False) to ensure legacy KBs are migrated on first view
metadata = KBAnalysisHelper.get_metadata(kb_dir, fast=False)
if _is_memory_base_associated(metadata):
continue # Skip KBs that are associated with a Memory Base
# Extract KB ID from metadata (stored as string, convert to UUID)
kb_id_str = metadata.get("id")
@ -503,7 +529,7 @@ async def list_knowledge_bases(
return knowledge_bases
@router.get("/{kb_name}", status_code=HTTPStatus.OK)
@router.get("/{kb_name}", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
async def get_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> KnowledgeBaseInfo:
"""Get detailed information about a specific knowledge base."""
try:
@ -543,7 +569,7 @@ async def get_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> K
raise HTTPException(status_code=500, detail="Error getting knowledge base.") from e
@router.get("/{kb_name}/chunks", status_code=HTTPStatus.OK)
@router.get("/{kb_name}/chunks", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
async def get_knowledge_base_chunks(
kb_name: str,
current_user: CurrentActiveUser,
@ -636,7 +662,7 @@ async def get_knowledge_base_chunks(
KBStorageHelper.release_chroma_resources(kb_path)
@router.delete("/{kb_name}", status_code=HTTPStatus.OK)
@router.delete("/{kb_name}", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
async def delete_knowledge_base(kb_name: str, current_user: CurrentActiveUser) -> dict[str, str]:
"""Delete a specific knowledge base."""
try:
@ -703,7 +729,7 @@ async def delete_knowledge_bases_bulk(request: BulkDeleteRequest, current_user:
return result
@router.post("/{kb_name}/cancel", status_code=HTTPStatus.OK)
@router.post("/{kb_name}/cancel", status_code=HTTPStatus.OK, dependencies=[Depends(_check_memory_base_association)])
async def cancel_ingestion(
kb_name: str,
current_user: CurrentActiveUser,

View File

@ -0,0 +1,345 @@
"""REST API for Memory Base management.
Endpoints:
POST /memories - Create
GET /memories - List (current user, paginated)
GET /memories/{id} - Get one
GET /memories/{id}/sessions - List sessions (tracked + untracked from MessageTable)
PATCH /memories/{id} - Update (name / threshold / auto_capture / preprocessing)
DELETE /memories/{id} - Delete (cancels active tasks + removes KB from disk)
POST /memories/{id}/flush - Manual flush / trigger ingestion
POST /memories/{id}/regenerate - Regenerate from mismatch
Edge cases enforced:
409 Conflict - name already in use for this user (on create).
409 Conflict - active ingestion task already running for same (mb, session).
404 Not Found - memory base does not belong to the current user.
422 Unprocessable - preprocessing=true but preproc_model missing.
"""
from __future__ import annotations
import uuid
from datetime import datetime
from http import HTTPStatus
from typing import Annotated, Any
from fastapi import APIRouter, Body, Depends, HTTPException
from fastapi_pagination import Page, Params
from fastapi_pagination.ext.sqlmodel import apaginate
from pydantic import BaseModel
from sqlmodel import col, select
from langflow.api.utils import CurrentActiveUser
from langflow.services.database.models.memory_base.model import (
MemoryBase,
MemoryBaseCreate,
MemoryBaseRead,
MemoryBaseSessionRead,
MemoryBaseUpdate,
)
from langflow.services.database.models.message.model import MessageTable
from langflow.services.deps import get_memory_base_service, session_scope
from langflow.services.jobs import DuplicateJobError
router = APIRouter(tags=["Memories"], prefix="/memories", include_in_schema=False)
# ------------------------------------------------------------------ #
# Request / Response schemas #
# ------------------------------------------------------------------ #
class MessageReadResponse(BaseModel):
"""Slim message projection for Memory Base session message listings.
Only messages that have been ingested into the requested Memory Base are returned.
``job_id`` and ``ingested_at`` are sourced from MessageIngestionRecord.
"""
model_config = {"from_attributes": True}
id: uuid.UUID
timestamp: datetime | None = None
sender: str
sender_name: str
session_id: str
text: str
content_blocks: list[dict[str, Any]] = []
job_id: uuid.UUID | None = None
ingested_at: datetime | None = None
class FlushRequest(BaseModel):
session_id: str
class MismatchResponse(BaseModel):
mismatch_detected: bool
class RegenerateResponse(BaseModel):
job_ids: list[str]
# ------------------------------------------------------------------ #
# CRUD #
# ------------------------------------------------------------------ #
@router.post("", status_code=HTTPStatus.CREATED)
@router.post("/", status_code=HTTPStatus.CREATED)
async def create_memory_base(
current_user: CurrentActiveUser,
payload: Annotated[MemoryBaseCreate, Body(embed=False)] = ...,
) -> MemoryBaseRead:
"""Create a new Memory Base.
- kb_name is auto-generated as `{sanitized_name}_{8hex}`.
- KB directory and embedding_metadata.json are created on disk immediately.
- Returns 409 if a Memory Base with the same name already exists for this user.
- Returns 422 if preprocessing=true but preproc_model is missing.
"""
try:
mb = await get_memory_base_service().create(payload, user_id=current_user.id)
except PermissionError as exc:
# Flow not found or belongs to another user — return 404 to avoid info-leak
raise HTTPException(status_code=404, detail=str(exc)) from exc
except ValueError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
return MemoryBaseRead.model_validate(mb)
@router.get("", status_code=HTTPStatus.OK)
@router.get("/", status_code=HTTPStatus.OK)
async def list_memory_bases(
current_user: CurrentActiveUser,
params: Annotated[Params, Depends()],
flow_id: uuid.UUID | None = None,
) -> Page[MemoryBaseRead]:
"""List all Memory Bases owned by the current user (paginated) for a flow_id.
Query params (from fastapi-pagination):
page - 1-based page number (default 1)
size - page size (default 50)
"""
async with session_scope() as db:
stmt = get_memory_base_service().list_for_user_stmt(user_id=current_user.id, flow_id=flow_id)
return await apaginate(
db, stmt, params=params, transformer=lambda items: [MemoryBaseRead.model_validate(m) for m in items]
)
@router.get("/{memory_base_id}", status_code=HTTPStatus.OK)
async def get_memory_base(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
) -> MemoryBaseRead:
"""Get details for a specific Memory Base."""
mb = await get_memory_base_service().get(memory_base_id, user_id=current_user.id)
if mb is None:
raise HTTPException(status_code=404, detail="Memory base not found")
return MemoryBaseRead.model_validate(mb)
@router.get("/{memory_base_id}/sessions", status_code=HTTPStatus.OK)
async def list_sessions(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
params: Annotated[Params, Depends()],
) -> Page[MemoryBaseSessionRead]:
"""List persisted sessions for this Memory Base (paginated).
Only sessions that have been synced at least once (i.e. have a
MemoryBaseSession row) are returned. Results are ordered by
last_sync_at descending.
Each item includes ``pending_count``: the number of completed flow runs
remaining before the next auto-capture ingestion is triggered.
"""
from langflow.services.memory_base.ingestion import count_pending_messages
async with session_scope() as db:
try:
mb = await get_memory_base_service().get_memory_base_or_404(db, memory_base_id, current_user.id)
except ValueError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
stmt = get_memory_base_service().sessions_stmt(memory_base_id, current_user.id)
raw_page = await apaginate(db, stmt, params=params)
items: list[MemoryBaseSessionRead] = []
for s in raw_page.items:
pending_count = await count_pending_messages(db, mb, s)
read = MemoryBaseSessionRead.model_validate(s)
read.pending_count = pending_count
items.append(read)
return raw_page.model_copy(update={"items": items})
@router.get("/{memory_base_id}/sessions/{session_id}/messages", status_code=HTTPStatus.OK)
async def list_session_messages(
memory_base_id: uuid.UUID,
session_id: str,
current_user: CurrentActiveUser,
params: Annotated[Params, Depends()],
) -> Page[MessageReadResponse]:
"""List messages ingested into this Memory Base session (paginated).
Only messages that have been successfully ingested into the requested Memory Base
are returned. Messages are ordered by timestamp ascending.
Each item includes ``job_id`` and ``ingested_at`` from the MessageIngestionRecord.
Returns 404 if the Memory Base does not belong to the current user.
"""
from sqlalchemy import and_
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
async with session_scope() as db:
mb_stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == current_user.id)
result = await db.exec(mb_stmt)
if result.first() is None:
raise HTTPException(status_code=404, detail="Memory base not found")
# INNER JOIN — only messages that were actually ingested into this MB/session pair.
# No extra WHERE filters needed:
# - mir.session_id == session_id in the JOIN guarantees msg.session_id == session_id
# (session_id is denormalized from the message at ingestion time — immutable).
# - flow_id is implicitly correct: ingestion only ever touches messages from mb.flow_id,
# and MB ownership is already verified above.
msg_stmt = (
select(MessageTable, MessageIngestionRecord)
.join(
MessageIngestionRecord,
and_(
MessageIngestionRecord.message_id == MessageTable.id,
MessageIngestionRecord.memory_base_id == memory_base_id,
MessageIngestionRecord.session_id == session_id,
),
)
.order_by(col(MessageTable.timestamp).asc())
)
return await apaginate(
db,
msg_stmt,
params=params,
transformer=lambda rows: [
MessageReadResponse(
id=msg.id,
timestamp=msg.timestamp,
sender=msg.sender,
sender_name=msg.sender_name,
session_id=msg.session_id,
text=msg.text,
content_blocks=msg.content_blocks or [],
job_id=mir.job_id,
ingested_at=mir.ingested_at,
)
for msg, mir in rows
],
)
@router.patch("/{memory_base_id}", status_code=HTTPStatus.OK)
async def update_memory_base(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
patch: Annotated[MemoryBaseUpdate, Body(embed=False)] = ...,
) -> MemoryBaseRead:
"""Update mutable parameters (threshold, auto_capture, preprocessing, etc.).
Threshold changes only take effect at the next auto-capture trigger.
Any already-running ingestion task continues with its original arguments.
"""
mb = await get_memory_base_service().update(memory_base_id, user_id=current_user.id, patch=patch)
if mb is None:
raise HTTPException(status_code=404, detail="Memory base not found")
return MemoryBaseRead.model_validate(mb)
@router.delete("/{memory_base_id}", status_code=HTTPStatus.NO_CONTENT)
async def delete_memory_base(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
) -> None:
"""Delete a Memory Base.
Active ingestion tasks are forcefully cancelled before the DB record is
removed. The associated KB directory is deleted from disk afterwards
(best-effort — a disk failure will not affect the 204 response).
"""
deleted = await get_memory_base_service().delete(memory_base_id, user_id=current_user.id)
if not deleted:
raise HTTPException(status_code=404, detail="Memory base not found")
# ------------------------------------------------------------------ #
# Ingestion trigger #
# ------------------------------------------------------------------ #
@router.post("/{memory_base_id}/flush", status_code=HTTPStatus.ACCEPTED)
async def flush_memory_base(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
body: Annotated[FlushRequest, Body(embed=False)] = ...,
) -> dict:
"""Manually trigger an ingestion / sync job regardless of the threshold.
Returns 409 Conflict if an ingestion task is already in progress for the
given (memory_base_id, session_id) pair to prevent concurrent indexing.
"""
try:
job_id = await get_memory_base_service().trigger_ingestion(
memory_base_id=memory_base_id,
user_id=current_user.id,
session_id=body.session_id,
)
except ValueError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
except DuplicateJobError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
except RuntimeError as exc:
raise HTTPException(status_code=409, detail=str(exc)) from exc
return {"job_id": job_id}
# ------------------------------------------------------------------ #
# Mismatch detection & regeneration #
# ------------------------------------------------------------------ #
@router.get("/{memory_base_id}/mismatch", status_code=HTTPStatus.OK)
async def check_mismatch(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
) -> MismatchResponse:
"""Detect if the vector store is empty while metadata records processed messages.
The UI should surface a "Mismatch Detected" warning and offer a Regenerate button.
"""
try:
detected = await get_memory_base_service().check_mismatch(memory_base_id, user_id=current_user.id)
except ValueError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
return MismatchResponse(mismatch_detected=detected)
@router.post("/{memory_base_id}/regenerate", status_code=HTTPStatus.ACCEPTED)
async def regenerate_memory_base(
memory_base_id: uuid.UUID,
current_user: CurrentActiveUser,
) -> RegenerateResponse:
"""Regenerate the Knowledge Base by resetting all session cursors and re-ingesting.
Use this to recover from external Chroma directory deletions or vector DB corruption.
All MemoryBaseSession.cursor_id values are set to None before re-running ingestion.
"""
try:
job_ids = await get_memory_base_service().regenerate(memory_base_id, user_id=current_user.id)
except ValueError as exc:
raise HTTPException(status_code=404, detail=str(exc)) from exc
return RegenerateResponse(job_ids=job_ids)

View File

@ -29,6 +29,7 @@ from uuid import UUID, uuid4
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, Request, status
from fastapi.responses import StreamingResponse
from lfx.graph.graph.base import Graph
from lfx.log.logger import logger
from lfx.schema.workflow import (
WORKFLOW_EXECUTION_RESPONSES,
WORKFLOW_STATUS_RESPONSES,
@ -65,7 +66,7 @@ from langflow.services.auth.utils import api_key_security
from langflow.services.database.models.flow.model import FlowRead
from langflow.services.database.models.jobs.model import JobType
from langflow.services.database.models.user.model import UserRead
from langflow.services.deps import get_job_service, get_task_service
from langflow.services.deps import get_job_service, get_memory_base_service, get_task_service
# Configuration constants
EXECUTION_TIMEOUT = 300 # 5 minutes default timeout for sync execution
@ -391,6 +392,18 @@ async def execute_sync_workflow(
stream=False,
)
# Fire memory-base auto-capture hook — non-blocking background effect.
try:
_run_id_uuid = UUID(graph.run_id) if graph.run_id else None # type-cast only; same run_id set on graph
await get_task_service().fire_and_forget_task(
get_memory_base_service().on_flow_output,
flow_id=flow.id,
session_id=execution_session_id,
job_id=_run_id_uuid,
)
except (RuntimeError, ValueError, OSError):
await logger.awarning("Memory base hook scheduling failed for flow %s", flow.id, exc_info=True)
# Build RunResponse
run_response = RunResponse(outputs=task_result, session_id=execution_session_id)
# Convert to WorkflowExecutionResponse
@ -470,10 +483,39 @@ async def execute_workflow_background(
user_id=api_key_user.id,
)
# Closure captures flow identity for the memory-base hook.
# run_id is the same as job_id — graph.set_run_id(job_id) was called above.
_hook_flow_id = flow.id
_hook_run_id = job_id
async def _run_and_notify(**kwargs):
"""Thin wrapper: execute graph then fire memory-base hook as a background effect.
The hook is dispatched non-blocking after graph completion. Any failure in
the hook is swallowed so it never affects the job status of the graph run.
"""
result = await run_graph_internal(**kwargs)
_, _effective_session_id = result
try:
# Direct await — we are already inside a background task; awaiting here
# is non-blocking from the client's perspective and avoids the race
# condition that arises when dispatching a second fire_and_forget from
# within an already-running fire_and_forget task.
await get_memory_base_service().on_flow_output(
flow_id=_hook_flow_id,
session_id=_effective_session_id,
job_id=_hook_run_id,
)
except Exception: # noqa: BLE001
await logger.awarning(
"Memory base hook failed for flow %s, but workflow succeeded.", _hook_flow_id, exc_info=True
)
return result
await task_service.fire_and_forget_task(
job_service.execute_with_status,
job_id=job_id,
run_coro_func=run_graph_internal,
run_coro_func=_run_and_notify,
graph=graph,
flow_id=flow_id_str,
session_id=session_id,
@ -673,7 +715,7 @@ async def stop_workflow(
task_service = get_task_service()
try:
# 1. Fetch Job
# 1. Fetch Job and verify ownership
job = await job_service.get_job_by_job_id(job_id, user_id=api_key_user.id)
except Exception as exc:
raise HTTPException(

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View File

@ -8,6 +8,7 @@ from .flow_version import FlowVersion
from .flow_version_deployment_attachment import FlowVersionDeploymentAttachment
from .folder import Folder
from .jobs import Job
from .memory_base import MemoryBase, MemoryBaseSession, MemoryBaseWorkflowRun, MessageIngestionRecord
from .message import MessageTable
from .traces.model import SpanTable, TraceTable
from .transactions import TransactionTable
@ -24,6 +25,10 @@ __all__ = [
"FlowVersionDeploymentAttachment",
"Folder",
"Job",
"MemoryBase",
"MemoryBaseSession",
"MemoryBaseWorkflowRun",
"MessageIngestionRecord",
"MessageTable",
"SSOConfig",
"SSOUserProfile",

View File

@ -63,6 +63,9 @@ class JobBase(SQLModel):
asset_type: str | None = Field(
index=False, nullable=True
) # Polymorphic: records if job is related to an entity like a KB, workflow, etc.
dedupe_key: str | None = Field(
index=True, nullable=True
) # Optional idempotency key to prevent duplicate jobs for the same asset and operation.
class Job(JobBase, table=True): # type: ignore[call-arg]

View File

@ -0,0 +1,21 @@
from langflow.services.database.models.memory_base.model import (
MemoryBase,
MemoryBaseCreate,
MemoryBaseRead,
MemoryBaseSession,
MemoryBaseSessionRead,
MemoryBaseUpdate,
MemoryBaseWorkflowRun,
MessageIngestionRecord,
)
__all__ = [
"MemoryBase",
"MemoryBaseCreate",
"MemoryBaseRead",
"MemoryBaseSession",
"MemoryBaseSessionRead",
"MemoryBaseUpdate",
"MemoryBaseWorkflowRun",
"MessageIngestionRecord",
]

View File

@ -0,0 +1,199 @@
from datetime import datetime, timezone
from uuid import UUID, uuid4
import sqlalchemy as sa
from pydantic import model_validator
from sqlalchemy import Column, DateTime, ForeignKey, Index, UniqueConstraint
from sqlmodel import Field, Relationship, SQLModel
class MemoryBaseBase(SQLModel):
name: str = Field(index=False)
flow_id: UUID = Field(index=True)
user_id: UUID = Field(index=True)
threshold: int = Field(default=50)
auto_capture: bool = Field(default=True)
# Preprocessing config — accepted in payload but logic deferred to future scope
embedding_model: str = Field(default="")
preprocessing: bool = Field(default=False)
preproc_model: str | None = Field(default=None)
preproc_instructions: str | None = Field(default=None)
class MemoryBase(MemoryBaseBase, table=True): # type: ignore[call-arg]
__tablename__ = "memory_base"
__table_args__ = (UniqueConstraint("user_id", "name", name="uq_memory_base_user_name"),)
id: UUID = Field(default_factory=uuid4, primary_key=True)
# kb_name is auto-generated at creation time — not user-supplied
kb_name: str = Field(default="")
created_at: datetime = Field(
default_factory=lambda: datetime.now(timezone.utc),
sa_column=Column(DateTime(timezone=True), nullable=False),
)
sessions: list["MemoryBaseSession"] = Relationship(
back_populates="memory_base",
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
)
class MemoryBaseCreate(MemoryBaseBase):
user_id: UUID | None = None # Derived from auth token in the endpoint; not required in request body
@model_validator(mode="after")
def preproc_model_required_when_preprocessing(self) -> "MemoryBaseCreate":
if self.preprocessing and not self.preproc_model:
msg = "preproc_model is required when preprocessing is enabled"
raise ValueError(msg)
return self
class MemoryBaseUpdate(SQLModel):
name: str | None = None
threshold: int | None = None
auto_capture: bool | None = None
preprocessing: bool | None = None
preproc_model: str | None = None
preproc_instructions: str | None = None
class MemoryBaseRead(MemoryBaseBase):
id: UUID
kb_name: str
created_at: datetime
class MemoryBaseSessionBase(SQLModel):
"""Fields shared between the table class and response schemas."""
session_id: str = Field(index=True)
cursor_id: UUID | None = Field(default=None)
total_processed: int = Field(default=0)
last_sync_at: datetime | None = Field(
default=None,
sa_column=Column(DateTime(timezone=True), nullable=True),
)
class MemoryBaseSession(MemoryBaseSessionBase, table=True): # type: ignore[call-arg]
__tablename__ = "memory_base_session"
__table_args__ = (
UniqueConstraint("memory_base_id", "session_id", name="uq_memory_base_session"),
Index("ix_memory_base_session_lookup", "memory_base_id", "session_id"),
)
id: UUID = Field(default_factory=uuid4, primary_key=True)
# FK defined via sa_column so Alembic sees the same shape as the migration:
# inline ForeignKey on the column with ondelete="CASCADE".
# This matches the pattern used by the File model (ForeignKey on sa_column).
memory_base_id: UUID = Field(
sa_column=Column(
sa.Uuid(),
ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
index=True,
)
)
memory_base: MemoryBase = Relationship(back_populates="sessions")
class MemoryBaseSessionRead(MemoryBaseSessionBase):
id: UUID
memory_base_id: UUID # Explicit — not in base to keep base free of DB-layer FK
pending_count: int = Field(default=0)
class MemoryBaseWorkflowRun(SQLModel, table=True): # type: ignore[call-arg]
"""Tracks WORKFLOW job runs per (memory_base, session) for threshold-based ingestion.
One row per WORKFLOW job, per session, per memory base.
- ``workflow_job_id``: the WORKFLOW job that produced this run (SET NULL on job deletion).
- ``ingestion_job_id``: set only after the ingestion job that processed this run completes
successfully. NULL means the run is still pending (not yet counted toward an ingestion).
Count pending = COUNT(*) WHERE ingestion_job_id IS NULL for a given (memory_base_id, session_id).
"""
__tablename__ = "memory_base_workflow_run"
__table_args__ = (
UniqueConstraint("memory_base_id", "session_id", "workflow_job_id", name="uq_mbwr_mb_session_wf_job"),
Index("ix_mbwr_mb_session", "memory_base_id", "session_id"),
Index("ix_mbwr_ingestion_job_id", "ingestion_job_id"),
)
id: UUID = Field(default_factory=uuid4, primary_key=True)
memory_base_id: UUID = Field(
sa_column=Column(
sa.Uuid(),
ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
)
)
session_id: str = Field(sa_column=Column(sa.String(), nullable=False))
workflow_job_id: UUID | None = Field(
default=None,
sa_column=Column(
sa.Uuid(),
ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
)
ingestion_job_id: UUID | None = Field(
default=None,
sa_column=Column(
sa.Uuid(),
ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
)
recorded_at: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))
class MessageIngestionRecord(SQLModel, table=True): # type: ignore[call-arg]
"""M-N join table recording which messages were ingested into which Memory Base by which job.
One record per (message, session, memory_base) — enforced by the unique constraint.
Records are written only after a confirmed successful Chroma write (write-on-success).
On regenerate, all records for the memory_base are deleted atomically alongside the
cursor reset so that re-ingestion starts clean.
"""
__tablename__ = "message_ingestion_record"
__table_args__ = (
UniqueConstraint("message_id", "session_id", "memory_base_id", name="uq_mir_message_session_mb"),
Index("ix_mir_message_id", "message_id"),
Index("ix_mir_job_id", "job_id"),
Index("ix_mir_memory_base_session", "memory_base_id", "session_id"),
)
id: UUID = Field(default_factory=uuid4, primary_key=True)
message_id: UUID = Field(
sa_column=Column(
sa.Uuid(),
ForeignKey("message.id", ondelete="CASCADE"),
nullable=False,
)
)
memory_base_id: UUID = Field(
sa_column=Column(
sa.Uuid(),
ForeignKey("memory_base.id", ondelete="CASCADE"),
nullable=False,
)
)
job_id: UUID | None = Field(
default=None,
sa_column=Column(
sa.Uuid(),
ForeignKey("job.job_id", ondelete="SET NULL"),
nullable=True,
),
)
# Denormalized from MessageTable.session_id — immutable, avoids JOIN on the hot query path
session_id: str = Field(sa_column=Column(sa.String(), nullable=False))
ingested_at: datetime = Field(sa_column=Column(DateTime(timezone=True), nullable=False))

View File

@ -32,6 +32,7 @@ class MessageBase(SQLModel):
properties: Properties = Field(default_factory=Properties)
category: str = Field(default="message")
content_blocks: list[ContentBlock] = Field(default_factory=list)
session_metadata: dict | None = Field(default=None)
@field_serializer("timestamp")
def serialize_timestamp(self, value):
@ -59,7 +60,7 @@ class MessageBase(SQLModel):
return value
@classmethod
def from_message(cls, message: "Message", flow_id: str | UUID | None = None):
def from_message(cls, message: "Message", flow_id: str | UUID | None = None, run_id: str | UUID | None = None):
if message.text is None or not message.sender or not message.sender_name:
msg = "The message does not have the required fields (text, sender, sender_name)."
raise ValueError(msg)
@ -114,6 +115,13 @@ class MessageBase(SQLModel):
msg = f"Flow ID {flow_id} is not a valid UUID"
raise ValueError(msg) from exc
if isinstance(run_id, str):
try:
run_id = UUID(run_id)
except ValueError as exc:
msg = f"Run ID {run_id} is not a valid UUID"
raise ValueError(msg) from exc
return cls(
sender=message.sender,
sender_name=message.sender_name,
@ -123,6 +131,7 @@ class MessageBase(SQLModel):
files=message.files or [],
timestamp=timestamp,
flow_id=flow_id,
run_id=run_id,
properties=properties,
category=message.category,
content_blocks=content_blocks,
@ -149,6 +158,8 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
id: UUID = Field(default_factory=uuid4, primary_key=True)
flow_id: UUID | None = Field(default=None)
run_id: UUID | None = Field(default=None, index=True)
is_output: bool = Field(default=False)
files: list[str] = Field(sa_column=Column(JSON))
properties: dict | Properties = Field( # type: ignore[assignment]
@ -207,11 +218,9 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
@field_serializer("properties", "content_blocks", "session_metadata")
@classmethod
def serialize_properties_or_content_blocks(cls, value) -> dict | list[dict] | None:
def serialize_properties_or_content_blocks(cls, value) -> dict | list[dict]:
# Redundant sanitization here acts as a defensive measure for rows
# already in the database that might contain NaN/Infinity values.
if value is None:
return None
if isinstance(value, list):
value = [cls.serialize_properties_or_content_blocks(item) for item in value]
elif hasattr(value, "model_dump"):
@ -224,8 +233,9 @@ class MessageTable(MessageBase, table=True): # type: ignore[call-arg]
class MessageRead(MessageBase):
id: UUID
flow_id: UUID | None = Field()
flow_id: UUID | None = None
session_metadata: dict | None = None
run_id: UUID | None = None
class MessageCreate(MessageBase):
@ -243,3 +253,5 @@ class MessageUpdate(SQLModel):
error: bool | None = None
properties: Properties | None = None
session_metadata: dict | None = None
category: str | None = None
content_blocks: list[ContentBlock] | None = None

View File

@ -269,3 +269,10 @@ def get_flow_events_service():
from langflow.services.flow_events.factory import FlowEventsServiceFactory
return get_service(ServiceType.FLOW_EVENTS_SERVICE, FlowEventsServiceFactory())
def get_memory_base_service():
"""Retrieves the MemoryBaseService instance from the service manager."""
from langflow.services.memory_base.factory import MemoryBaseServiceFactory
return get_service(ServiceType.MEMORY_BASE_SERVICE, MemoryBaseServiceFactory())

View File

@ -1,5 +1,6 @@
"""Job service package."""
from langflow.services.jobs.exceptions import DuplicateJobError
from langflow.services.jobs.service import JobService
__all__ = ["JobService"]
__all__ = ["DuplicateJobError", "JobService"]

View File

@ -0,0 +1,16 @@
"""Domain exceptions for the jobs service."""
from __future__ import annotations
class JobError(RuntimeError):
"""Base exception for job-domain errors."""
class DuplicateJobError(JobError):
"""Raised by create_job() when a non-retryable job with the same dedupe_key already exists.
(QUEUED, IN_PROGRESS, or COMPLETED).
FAILED and CANCELLED are retryable and do not trigger this error.
Extends RuntimeError so existing except RuntimeError callers keep working.
"""

View File

@ -10,15 +10,16 @@ if TYPE_CHECKING:
from datetime import datetime, timezone
from uuid import UUID
from sqlmodel import col, func, select
from langflow.services.base import Service
from langflow.services.database.models.jobs.crud import (
get_job_by_job_id,
get_jobs_by_flow_id,
get_latest_jobs_by_asset_ids,
update_job_status,
)
from langflow.services.database.models.jobs.model import Job, JobStatus, JobType
from langflow.services.deps import session_scope
from langflow.services.jobs.exceptions import DuplicateJobError
class JobService(Service):
@ -30,11 +31,14 @@ class JobService(Service):
"""Initialize the job service."""
self.set_ready()
async def get_jobs_by_flow_id(self, flow_id: UUID | str, page: int = 1, page_size: int = 10) -> list[Job]:
"""Get jobs for a specific flow with pagination.
async def get_jobs_by_flow_id(
self, flow_id: UUID | str, user_id: UUID, page: int = 1, page_size: int = 10
) -> list[Job]:
"""Get jobs for a specific flow with pagination, filtered by user.
Args:
flow_id: The flow ID to filter jobs by
user_id: The user ID to enforce ownership
page: Page number (1-indexed)
page_size: Number of jobs per page
@ -45,7 +49,16 @@ class JobService(Service):
flow_id = UUID(flow_id)
async with session_scope() as session:
return await get_jobs_by_flow_id(session, flow_id, page=page, size=page_size)
stmt = (
select(Job)
.where(Job.flow_id == flow_id)
.where((Job.user_id == user_id) | (Job.user_id.is_(None)))
.order_by(col(Job.created_at).desc())
.offset((page - 1) * page_size)
.limit(page_size)
)
result = await session.exec(stmt)
return list(result.all())
async def get_job_by_job_id(self, job_id: UUID | str, user_id: UUID | None = None) -> Job | None:
"""Get job for a specific job ID.
@ -62,7 +75,11 @@ class JobService(Service):
job_id = UUID(job_id)
async with session_scope() as session:
return await get_job_by_job_id(session, job_id, user_id=user_id)
stmt = select(Job).where(Job.job_id == job_id)
if user_id:
stmt = stmt.where((Job.user_id == user_id) | (Job.user_id.is_(None)))
result = await session.exec(stmt)
return result.first()
async def create_job(
self,
@ -72,16 +89,19 @@ class JobService(Service):
asset_id: UUID | None = None,
asset_type: str | None = None,
user_id: UUID | None = None,
dedupe_key: str | None = None,
) -> Job:
"""Create a new job record with QUEUED status.
Args:
job_id: The job ID
flow_id: The flow ID
user_id: The user ID
job_type: The job type
asset_id: The asset ID
asset_type: The asset type
user_id: The user ID who owns this job
dedupe_key: Optional idempotency key to prevent duplicate jobs for the same batch
Returns:
Created Job object
@ -93,6 +113,18 @@ class JobService(Service):
flow_id = UUID(flow_id)
async with session_scope() as session:
if dedupe_key is not None:
stmt = (
select(func.count())
.select_from(Job)
.where(Job.dedupe_key == dedupe_key)
.where(col(Job.status).in_([JobStatus.QUEUED, JobStatus.IN_PROGRESS, JobStatus.COMPLETED]))
)
result = await session.exec(stmt)
if result.one() > 0:
msg = f"A non-retryable job with dedupe_key={dedupe_key!r} already exists"
raise DuplicateJobError(msg)
job = Job(
job_id=job_id,
flow_id=flow_id,
@ -101,6 +133,7 @@ class JobService(Service):
asset_id=asset_id,
asset_type=asset_type,
user_id=user_id,
dedupe_key=dedupe_key,
)
session.add(job)
await session.flush()
@ -210,3 +243,18 @@ class JobService(Service):
await logger.ainfo(f"Job {job_id} completed successfully")
await self.update_job_status(job_id, JobStatus.COMPLETED, finished_timestamp=True)
return result
async def _validate_ownership(self, job_id: UUID, user_id: UUID) -> Job:
"""Verify that a job exists and belongs to the specified user.
Raises:
ValueError: If the job is not found or is NOT owned by the user.
"""
job = await self.get_job_by_job_id(job_id)
if job is None:
msg = f"Job {job_id} not found"
raise ValueError(msg)
if job.user_id is not None and job.user_id != user_id:
msg = f"Access denied for job {job_id}"
raise ValueError(msg)
return job

View File

@ -0,0 +1,3 @@
from langflow.services.memory_base.service import MemoryBaseService
__all__ = ["MemoryBaseService"]

View File

@ -0,0 +1,139 @@
"""Document building and KB metadata sync helpers for Memory Base ingestion.
Extracted from task.py to separate "document shaping" from "ingestion orchestration".
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from lfx.log.logger import logger
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
if TYPE_CHECKING:
from pathlib import Path
from langchain_chroma import Chroma
from langflow.services.database.models.message.model import MessageTable
# Chunk size for splitting long messages before embedding
MESSAGE_CHUNK_SIZE = 1000
MESSAGE_CHUNK_OVERLAP = 100
def extract_content_block_text(content_blocks: list) -> str:
"""Extract embeddable text from content blocks of type text, code, and json.
Blocks of any other type (tool_use, error, media, etc.) are skipped.
Each extracted piece is separated by a blank line so chunk boundaries
remain readable in the vector store.
"""
parts: list[str] = []
for block in content_blocks:
# content_blocks are stored as JSON; each block is a dict at runtime.
contents: list = block.get("contents", []) if isinstance(block, dict) else []
for entry in contents:
if not isinstance(entry, dict):
continue
entry_type = entry.get("type")
if entry_type == "text":
fragment = (entry.get("text") or "").strip()
elif entry_type == "code":
lang = entry.get("language") or ""
code = (entry.get("code") or "").strip()
fragment = f"```{lang}\n{code}\n```" if code else ""
elif entry_type == "json":
data = entry.get("data")
fragment = json.dumps(data, ensure_ascii=False) if data is not None else ""
else:
continue
if fragment:
parts.append(fragment)
return "\n\n".join(parts)
def build_documents_from_messages(
messages: list[MessageTable],
*,
session_id: str,
flow_id: str,
job_id: str = "",
) -> list[Document]:
"""Convert MessageTable rows into LangChain Documents.
Each message's embeddable text is the concatenation of msg.text and any
content-block fragments whose type is text, code, or json. Other block
types (tool_use, error, media, ...) are ignored. Long combined texts are
split by RecursiveCharacterTextSplitter before embedding.
"""
splitter = RecursiveCharacterTextSplitter(
chunk_size=MESSAGE_CHUNK_SIZE,
chunk_overlap=MESSAGE_CHUNK_OVERLAP,
)
docs: list[Document] = []
for msg in messages:
parts: list[str] = []
if msg.text and msg.text.strip():
parts.append(msg.text.strip())
cb_text = extract_content_block_text(msg.content_blocks or [])
if cb_text:
parts.append(cb_text)
text = "\n\n".join(parts)
if not text:
continue
chunks = splitter.split_text(text)
for i, chunk in enumerate(chunks):
docs.append(
Document(
page_content=chunk,
metadata={
"message_id": str(msg.id),
"session_id": session_id,
"flow_id": flow_id,
"sender": msg.sender,
"sender_name": msg.sender_name,
"timestamp": msg.timestamp.isoformat() if msg.timestamp else "",
"run_id": str(msg.run_id) if msg.run_id else "",
"chunk_index": i,
"total_chunks": len(chunks),
"source": f"memory_base/{session_id}",
"job_id": job_id,
},
)
)
return docs
def sync_kb_metadata(*, kb_path: Path, chroma: Chroma) -> None:
"""Update embedding_metadata.json after a successful Memory Base ingestion.
Mirrors the post-write metadata sync in ``KBIngestionHelper.perform_ingestion``:
- Refreshes chunk / word / character counts from the live Chroma collection.
- Updates on-disk size.
- Stamps ``is_memory_base: true`` (required for Knowledge Retrieval filtering).
- Sets ``source_types: ["memory"]`` to distinguish from file-based KBs.
Called while the Chroma client is still open so that ``update_text_metrics``
can query the collection directly without opening a second client.
"""
try:
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
KBAnalysisHelper.update_text_metrics(kb_path, metadata, chroma=chroma)
metadata["size"] = KBStorageHelper.get_directory_size(kb_path)
metadata["is_memory_base"] = True
# Preserve any existing source_types but always include "memory"
existing = set(metadata.get("source_types") or [])
existing.add("memory")
metadata["source_types"] = sorted(existing)
(kb_path / "embedding_metadata.json").write_text(json.dumps(metadata, indent=2))
except (OSError, json.JSONDecodeError, ValueError):
# Metadata sync is best-effort; a failure here must not block the cursor advance.
# Note: this runs inside asyncio.to_thread so we use sync logging here.
# The lfx logger's sync .warning() method goes through the same structured pipeline.
logger.warning("KB metadata sync failed for kb_path=%s", kb_path, exc_info=True)

View File

@ -0,0 +1,26 @@
"""Embedding provider inference for MemoryBase.
Extracted from MemoryBaseService to keep single-responsibility per file.
"""
from __future__ import annotations
# Provider inference map — mirrors provider_patterns in KBAnalysisHelper._detect_embedding_provider
# so we can derive the provider from a model name string without filesystem access.
_MODEL_TO_PROVIDER: list[tuple[list[str], str]] = [
(["text-embedding", "ada-", "gpt-"], "OpenAI"),
(["embed-english", "embed-multilingual"], "Cohere"),
(["sentence-transformers", "bert-", "huggingface"], "HuggingFace"),
(["palm", "gecko", "google"], "Google"),
(["ollama"], "Ollama"),
(["azure"], "Azure OpenAI"),
]
def infer_embedding_provider(embedding_model: str) -> str:
"""Derive embedding provider name from a model string."""
lower = embedding_model.lower()
for patterns, provider in _MODEL_TO_PROVIDER:
if any(p in lower for p in patterns):
return provider
return "OpenAI" # Safe default — matches _resolve_embedding fallback

View File

@ -0,0 +1,14 @@
"""Factory for creating MemoryBaseService instances."""
from langflow.services.factory import ServiceFactory
from langflow.services.memory_base.service import MemoryBaseService
class MemoryBaseServiceFactory(ServiceFactory):
"""Factory for creating MemoryBaseService instances."""
def __init__(self):
super().__init__(MemoryBaseService)
def create(self):
return MemoryBaseService()

View File

@ -0,0 +1,421 @@
"""Ingestion orchestration for MemoryBase — auto-capture, regeneration, and mismatch detection.
Extracted from MemoryBaseService to keep single-responsibility per file.
The MemoryBaseService delegates to these functions for all ingestion-related work.
"""
from __future__ import annotations
import asyncio
import uuid
from typing import TYPE_CHECKING
from lfx.log.logger import logger
from sqlmodel import col, func, select
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
from langflow.services.database.models.jobs.model import Job, JobStatus, JobType
from langflow.services.database.models.memory_base.model import (
MemoryBase,
MemoryBaseSession,
MemoryBaseWorkflowRun,
)
from langflow.services.deps import get_job_service, get_task_service, session_scope
from langflow.services.jobs import DuplicateJobError
from langflow.services.memory_base.kb_path_helpers import (
hash_session_id,
resolve_embedding,
resolve_kb_username,
resolve_kb_username_by_user_id,
validate_kb_path,
)
from langflow.services.memory_base.task import IngestionRequest, ingest_memory_task
if TYPE_CHECKING:
from sqlmodel.ext.asyncio.session import AsyncSession
async def trigger_ingestion(
memory_base_id: uuid.UUID,
user_id: uuid.UUID,
session_id: str,
*,
get_mb_or_raise,
get_or_create_session,
) -> str:
"""Manually trigger (or auto-trigger) an ingestion sync.
Returns:
job_id string for the newly created job.
Raises:
ValueError: If MemoryBase not found.
RuntimeError: If a job is already active (caller should return 409).
"""
async with session_scope() as db:
mb = await get_mb_or_raise(db, memory_base_id, user_id)
# Ensure a session record exists
mbs = await get_or_create_session(db, memory_base_id, session_id)
# Snapshot the cursor NOW (immutable arg for the task)
cursor_id_snapshot = mbs.cursor_id
# Build dedupe_key from the latest uncovered WORKFLOW run for idempotency.
latest_job_id = await _get_latest_pending_workflow_job_id(db, mb, mbs)
dedupe_key: str | None = None
if latest_job_id is not None:
dedupe_key = f"ingestion:{memory_base_id}:{session_id}:{latest_job_id}"
kb_username = await resolve_kb_username(db, mb.user_id)
embedding_provider, embedding_model = resolve_embedding(mb.kb_name, kb_username)
# Create tracking job
job_service = get_job_service()
job_id = uuid.uuid4()
await job_service.create_job(
job_id=job_id,
flow_id=mb.flow_id,
user_id=mb.user_id,
job_type=JobType.INGESTION,
asset_id=memory_base_id,
asset_type="memory_base",
dedupe_key=dedupe_key,
)
task_service = get_task_service()
await task_service.fire_and_forget_task(
job_service.execute_with_status,
job_id=job_id,
run_coro_func=ingest_memory_task,
request=IngestionRequest(
memory_base_id=memory_base_id,
session_id=session_id,
flow_id=mb.flow_id,
kb_name=mb.kb_name,
kb_username=kb_username,
user_id=mb.user_id,
embedding_provider=embedding_provider,
embedding_model=embedding_model,
cursor_id=cursor_id_snapshot,
task_job_id=job_id,
job_service=job_service,
),
)
return str(job_id)
async def on_flow_output(
flow_id: uuid.UUID,
session_id: str,
job_id: uuid.UUID | None,
*,
get_or_create_session,
) -> None:
"""Called after a flow run completes.
For every MemoryBase watching this flow with auto_capture=True:
1. Record the workflow run in the tracking table (inside _maybe_trigger).
2. Count uncovered WORKFLOW runs for this session.
3. If count >= threshold, fire ingestion task.
"""
async with session_scope() as db:
stmt = (
select(MemoryBase).where(MemoryBase.flow_id == flow_id).where(MemoryBase.auto_capture == True) # noqa: E712
)
result = await db.exec(stmt)
memory_bases = list(result.all())
hashed_sid = hash_session_id(session_id)
for mb in memory_bases:
try:
await logger.adebug(
"Auto-capture check | memory_base=%s threshold=%s session=%s",
mb.id,
mb.threshold,
hashed_sid,
)
await _maybe_trigger(
mb=mb, session_id=session_id, job_id=job_id, get_or_create_session=get_or_create_session
)
except (RuntimeError, ValueError, OSError):
await logger.aerror("Auto-capture failed for memory_base=%s session=%s", mb.id, hashed_sid, exc_info=True)
async def _maybe_trigger(
*,
mb: MemoryBase,
session_id: str,
job_id: uuid.UUID | None,
get_or_create_session,
) -> None:
async with session_scope() as db:
mbs = await get_or_create_session(db, mb.id, session_id)
# Record this workflow run before evaluating the threshold.
await _insert_workflow_run(db, mb.id, session_id, job_id)
pending = await count_pending_messages(db, mb, mbs)
if pending < mb.threshold:
return
cursor_id_snapshot = mbs.cursor_id
# Build dedupe_key from the latest pending WORKFLOW run for idempotency.
latest_wf_job_id = await _get_latest_pending_workflow_job_id(db, mb, mbs)
dedupe_key: str | None = None
if latest_wf_job_id is not None:
dedupe_key = f"ingestion:{mb.id}:{session_id}:{latest_wf_job_id}"
kb_username = await resolve_kb_username(db, mb.user_id)
embedding_provider, embedding_model = resolve_embedding(mb.kb_name, kb_username)
job_service = get_job_service()
job_id = uuid.uuid4()
try:
await job_service.create_job(
job_id=job_id,
flow_id=mb.flow_id,
user_id=mb.user_id,
job_type=JobType.INGESTION,
asset_id=mb.id,
asset_type="memory_base",
dedupe_key=dedupe_key,
)
except DuplicateJobError:
await logger.adebug("Auto-capture: duplicate job for dedupe_key=%s - skipping.", dedupe_key)
return
task_service = get_task_service()
await task_service.fire_and_forget_task(
job_service.execute_with_status,
job_id=job_id,
run_coro_func=ingest_memory_task,
request=IngestionRequest(
memory_base_id=mb.id,
session_id=session_id,
flow_id=mb.flow_id,
kb_name=mb.kb_name,
kb_username=kb_username,
user_id=mb.user_id,
embedding_provider=embedding_provider,
embedding_model=embedding_model,
cursor_id=cursor_id_snapshot,
task_job_id=job_id,
job_service=job_service,
),
)
async def check_mismatch(
memory_base_id: uuid.UUID,
user_id: uuid.UUID,
*,
get_mb_or_raise,
) -> bool:
"""Return True if metadata claims processed rows but vector store is empty."""
async with session_scope() as db:
mb = await get_mb_or_raise(db, memory_base_id, user_id)
stmt = select(func.sum(MemoryBaseSession.total_processed)).where(
MemoryBaseSession.memory_base_id == memory_base_id
)
result = await db.exec(stmt)
total_processed: int = result.first() or 0
if total_processed == 0:
return False
kb_username = await resolve_kb_username_by_user_id(user_id)
kb_root = KBStorageHelper.get_root_path()
if not kb_root:
return False
kb_path = kb_root / kb_username / mb.kb_name
validate_kb_path(kb_root, kb_path)
if not await asyncio.to_thread(kb_path.exists):
return True
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
return int(metadata.get("chunks", 0)) == 0
async def regenerate(
memory_base_id: uuid.UUID,
user_id: uuid.UUID,
*,
get_mb_or_raise,
trigger_ingestion_fn,
) -> list[str]:
"""Reset all session cursors to None and re-trigger ingestion per session.
Used to recover from FS / Vector DB mismatch (Chroma dir deleted externally).
Returns list of newly created job IDs.
Also deletes all MessageIngestionRecord rows for this memory base atomically
with the cursor reset so that re-ingestion starts clean without hitting the
unique constraint.
"""
from sqlalchemy import delete as sa_delete
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
async with session_scope() as db:
await get_mb_or_raise(db, memory_base_id, user_id)
stmt = select(MemoryBaseSession).where(MemoryBaseSession.memory_base_id == memory_base_id)
result = await db.exec(stmt)
sessions = list(result.all())
for s in sessions:
s.cursor_id = None
db.add(s)
# Delete existing ingestion records so re-ingestion inserts fresh rows
await db.exec( # type: ignore[call-overload]
sa_delete(MessageIngestionRecord).where(MessageIngestionRecord.memory_base_id == memory_base_id)
)
await db.commit()
job_ids: list[str] = []
for s in sessions:
try:
jid = await trigger_ingestion_fn(memory_base_id, user_id, s.session_id)
job_ids.append(jid)
except DuplicateJobError:
await logger.awarning(
"Regenerate: duplicate batch already ingested for session %s - skipped.",
hash_session_id(s.session_id),
)
except RuntimeError:
await logger.awarning(
"Regenerate: active job exists for session %s - reset cursor but skipped trigger.",
hash_session_id(s.session_id),
)
return job_ids
async def cancel_active_jobs(*, memory_base_id: uuid.UUID, db: AsyncSession) -> None:
"""Cancel all IN_PROGRESS or QUEUED jobs for this memory base."""
stmt = (
select(Job)
.where(Job.asset_id == memory_base_id)
.where(Job.asset_type == "memory_base")
.where(col(Job.status).in_([JobStatus.IN_PROGRESS, JobStatus.QUEUED]))
)
result = await db.exec(stmt)
active_jobs = list(result.all())
task_service = get_task_service()
job_service = get_job_service()
for job in active_jobs:
try:
await task_service.revoke_task(job.job_id)
await job_service.update_job_status(job.job_id, JobStatus.CANCELLED)
await logger.ainfo("Cancelled job %s for memory_base %s", job.job_id, memory_base_id)
except (RuntimeError, ValueError, OSError):
await logger.awarning(
"Could not cancel job %s for memory_base %s", job.job_id, memory_base_id, exc_info=True
)
# ------------------------------------------------------------------ #
# Shared query helpers (public — used by service.py and memories.py) #
# ------------------------------------------------------------------ #
async def count_pending_messages(db: AsyncSession, mb: MemoryBase, mbs: MemoryBaseSession) -> int:
"""Count WORKFLOW runs for this (memory_base, session) not yet covered by a completed ingestion.
A row in memory_base_workflow_run with ingestion_job_id IS NULL means the run
has not been processed by any ingestion job. Count pending = number of such rows.
This is session-scoped and time-independent; job failures leave rows NULL so they
are correctly re-counted on the next threshold check.
"""
stmt = (
select(func.count())
.select_from(MemoryBaseWorkflowRun)
.where(MemoryBaseWorkflowRun.memory_base_id == mb.id)
.where(MemoryBaseWorkflowRun.session_id == mbs.session_id)
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
)
try:
result = await db.exec(stmt)
row = result.first()
if row is None:
return 0
return int(row)
except (TypeError, ValueError, OSError) as e:
await logger.aerror("Error counting pending workflow runs: %s", e)
return 0
async def _insert_workflow_run(
db: AsyncSession,
memory_base_id: uuid.UUID,
session_id: str,
job_id: uuid.UUID | None,
) -> None:
"""Record a WORKFLOW job run for (memory_base_id, session_id).
Verifies that job_id refers to a WORKFLOW type job before inserting.
Uses dialect-specific INSERT ... ON CONFLICT DO NOTHING for idempotency —
safe to call multiple times with the same arguments.
Skips silently if job_id is None or the job is not of WORKFLOW type.
"""
from datetime import datetime, timezone
hashed_sid = hash_session_id(session_id)
if job_id is None:
await logger.awarning(
"on_flow_output called with no job_id for memory_base=%s session=%s — run not recorded.",
memory_base_id,
hashed_sid,
)
return
job_result = await db.exec(select(Job).where(Job.job_id == job_id).where(Job.type == JobType.WORKFLOW))
if job_result.first() is None:
await logger.awarning(
"job_id=%s is not a WORKFLOW job — skipping workflow run record for memory_base=%s session=%s.",
job_id,
memory_base_id,
hashed_sid,
)
return
row = {
"id": uuid.uuid4(),
"memory_base_id": memory_base_id,
"session_id": session_id,
"workflow_job_id": job_id,
"ingestion_job_id": None,
"recorded_at": datetime.now(timezone.utc),
}
conn = await db.connection()
if conn.dialect.name == "postgresql":
from sqlalchemy.dialects.postgresql import insert as pg_insert
stmt = pg_insert(MemoryBaseWorkflowRun).values([row]).on_conflict_do_nothing()
else:
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
stmt = sqlite_insert(MemoryBaseWorkflowRun).values([row]).on_conflict_do_nothing()
await db.exec(stmt) # type: ignore[call-overload]
await db.commit()
async def _get_latest_pending_workflow_job_id(
db: AsyncSession, mb: MemoryBase, mbs: MemoryBaseSession
) -> uuid.UUID | None:
"""Return the workflow_job_id of the most recent uncovered workflow run for this session."""
stmt = (
select(MemoryBaseWorkflowRun.workflow_job_id)
.where(MemoryBaseWorkflowRun.memory_base_id == mb.id)
.where(MemoryBaseWorkflowRun.session_id == mbs.session_id)
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
.order_by(col(MemoryBaseWorkflowRun.recorded_at).desc())
.limit(1)
)
result = await db.exec(stmt)
return result.first()

View File

@ -0,0 +1,149 @@
"""KB path resolution and username helpers for MemoryBase.
Extracted from MemoryBaseService to keep single-responsibility per file.
"""
from __future__ import annotations
import asyncio
import hashlib
import json
import re
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING
from lfx.log.logger import logger
from sqlmodel import select
from langflow.api.utils.kb_helpers import KBAnalysisHelper, KBStorageHelper
from langflow.services.deps import session_scope
if TYPE_CHECKING:
from pathlib import Path
from sqlmodel.ext.asyncio.session import AsyncSession
def validate_kb_path(kb_root: Path, kb_path: Path) -> None:
"""Assert that kb_path is contained within kb_root (path traversal guard).
Prevents crafted usernames with '..' segments from escaping the KB root directory.
Follows the same pattern as services/storage/local.py:save_file.
"""
kb_root_resolved = kb_root.resolve()
kb_path_resolved = kb_path.resolve()
if not kb_path_resolved.is_relative_to(kb_root_resolved):
msg = "KB path escapes root directory"
raise ValueError(msg)
def hash_session_id(session_id: str) -> str:
"""Return a truncated SHA-256 hash for safe logging of session IDs."""
return hashlib.sha256(session_id.encode()).hexdigest()[:12]
def sanitize_kb_name(name: str) -> str:
"""Lowercase, replace spaces/hyphens with underscores, strip non-alphanum."""
sanitized = name.strip().lower()
sanitized = re.sub(r"[\s\-]+", "_", sanitized)
sanitized = re.sub(r"[^\w]", "", sanitized)
return sanitized or "memory"
async def resolve_kb_username(db: AsyncSession, user_id: uuid.UUID) -> str:
"""Look up the username for a user_id within an existing DB session."""
from langflow.services.database.models.user.model import User
stmt = select(User.username).where(User.id == user_id)
result = await db.exec(stmt)
username = result.first()
if not username:
msg = f"User {user_id} not found"
raise ValueError(msg)
return username
async def resolve_kb_username_by_user_id(user_id: uuid.UUID) -> str:
"""Look up the username for a user_id using a fresh DB session."""
async with session_scope() as db:
return await resolve_kb_username(db, user_id)
def resolve_embedding(kb_name: str, kb_username: str) -> tuple[str, str]:
"""Read embedding provider/model from KB metadata.json, with sane defaults."""
kb_root = KBStorageHelper.get_root_path()
if not kb_root:
return "OpenAI", "text-embedding-3-small"
kb_path: Path = kb_root / kb_username / kb_name
metadata = KBAnalysisHelper.get_metadata(kb_path, fast=True)
provider = metadata.get("embedding_provider") or "OpenAI"
model = metadata.get("embedding_model") or "text-embedding-3-small"
return provider, model
async def initialize_kb(
*,
kb_name: str,
kb_username: str,
embedding_provider: str,
embedding_model: str,
) -> None:
"""Create KB directory, initialize Chroma, and write embedding_metadata.json.
Mirrors the logic in knowledge_bases.py:create_knowledge_base so Memory Base
KBs are immediately visible with the correct metadata (including is_memory_base: true).
"""
import chromadb
kb_root = KBStorageHelper.get_root_path()
if not kb_root:
await logger.awarning("KB root path not configured — Memory Base KB will not be initialized on disk.")
return
kb_path: Path = kb_root / kb_username / kb_name
validate_kb_path(kb_root, kb_path)
await asyncio.to_thread(kb_path.mkdir, parents=True, exist_ok=True)
# Initialize Chroma collection so the directory is non-empty and readable
try:
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
client.create_collection(name=kb_name)
except (OSError, ValueError, chromadb.errors.ChromaError) as exc:
await logger.awarning("Initial Chroma setup for %s failed: %s", kb_name, exc)
finally:
client = None # type: ignore[assignment]
KBStorageHelper.release_chroma_resources(kb_path)
embedding_metadata = {
"id": str(uuid.uuid4()),
"embedding_provider": embedding_provider,
"embedding_model": embedding_model,
"is_memory_base": True,
"created_at": datetime.now(timezone.utc).isoformat(),
"chunks": 0,
"words": 0,
"characters": 0,
"avg_chunk_size": 0.0,
"size": 0,
"source_types": ["memory"],
}
await asyncio.to_thread(
(kb_path / "embedding_metadata.json").write_text,
json.dumps(embedding_metadata, indent=2),
)
async def delete_kb(*, kb_name: str, kb_username: str) -> None:
"""Remove the KB directory from disk. Logs on failure, does not raise."""
if not kb_name:
return
kb_root = KBStorageHelper.get_root_path()
if not kb_root:
return
kb_path = kb_root / kb_username / kb_name
validate_kb_path(kb_root, kb_path)
try:
await asyncio.to_thread(KBStorageHelper.delete_storage, kb_path, kb_name)
except (OSError, ValueError):
await logger.awarning("Could not delete KB '%s' from disk after Memory Base deletion.", kb_name, exc_info=True)

View File

@ -0,0 +1,288 @@
"""MemoryBase service — CRUD and session state management.
Ingestion orchestration, KB path helpers, and embedding inference are in
separate modules (ingestion.py, kb_path_helpers.py, embedding_helpers.py)
to keep this file focused on data access and business-rule enforcement.
Edge cases handled:
- Name uniqueness per user: 409 if a Memory Base with the same name already exists.
- Deletion during sync: cancels active tasks before DB deletion.
- KB deletion on delete: removes the associated KB directory from disk.
- Concurrent task prevention: returns 409 if a job is already IN_PROGRESS.
- Threshold updates: deferred; does not re-evaluate pending count immediately.
"""
from __future__ import annotations
import uuid
from typing import TYPE_CHECKING
from sqlmodel import col, select
from langflow.services.base import Service
from langflow.services.database.models.memory_base.model import (
MemoryBase,
MemoryBaseCreate,
MemoryBaseSession,
MemoryBaseUpdate,
)
from langflow.services.deps import session_scope
from langflow.services.memory_base.embedding_helpers import infer_embedding_provider
from langflow.services.memory_base.ingestion import (
cancel_active_jobs,
)
from langflow.services.memory_base.ingestion import (
check_mismatch as _check_mismatch,
)
from langflow.services.memory_base.ingestion import (
on_flow_output as _on_flow_output,
)
from langflow.services.memory_base.ingestion import (
regenerate as _regenerate,
)
from langflow.services.memory_base.ingestion import (
trigger_ingestion as _trigger_ingestion,
)
from langflow.services.memory_base.kb_path_helpers import (
delete_kb,
initialize_kb,
resolve_kb_username,
sanitize_kb_name,
)
if TYPE_CHECKING:
from sqlmodel.ext.asyncio.session import AsyncSession
class MemoryBaseService(Service):
"""Service layer for MemoryBase CRUD and session state management."""
name = "memory_base_service"
# ------------------------------------------------------------------ #
# CRUD #
# ------------------------------------------------------------------ #
async def create(self, payload: MemoryBaseCreate, user_id: uuid.UUID) -> MemoryBase:
# 1. Verify that the referenced flow belongs to this user.
async with session_scope() as db:
from langflow.services.database.models.flow.model import Flow
flow_result = await db.exec(select(Flow).where(Flow.id == payload.flow_id).where(Flow.user_id == user_id))
if flow_result.first() is None:
msg = f"Flow {payload.flow_id} not found"
raise PermissionError(msg)
# 2. Resolve username — needed for the KB path.
async with session_scope() as db:
kb_username = await resolve_kb_username(db, user_id)
# 3. Auto-generate kb_name: sanitized_name_<8hex>
kb_name = f"{sanitize_kb_name(payload.name)}_{uuid.uuid4().hex[:8]}"
# 4. Create KB directory and embedding_metadata.json on disk.
embedding_provider = infer_embedding_provider(payload.embedding_model)
await initialize_kb(
kb_name=kb_name,
kb_username=kb_username,
embedding_provider=embedding_provider,
embedding_model=payload.embedding_model,
)
# 5. Uniqueness check + insert.
from sqlalchemy.exc import IntegrityError
async with session_scope() as db:
existing = await db.exec(
select(MemoryBase).where(MemoryBase.user_id == user_id).where(MemoryBase.name == payload.name)
)
if existing.first() is not None:
msg = f"A Memory Base named '{payload.name}' already exists for this user"
raise ValueError(msg)
mb = MemoryBase(
**payload.model_dump(exclude={"user_id"}),
user_id=user_id,
kb_name=kb_name,
)
db.add(mb)
try:
await db.commit()
except IntegrityError:
msg = f"A Memory Base named '{payload.name}' already exists for this user"
raise ValueError(msg) from None
await db.refresh(mb)
return mb
async def list_for_user(self, user_id: uuid.UUID) -> list[MemoryBase]:
async with session_scope() as db:
stmt = select(MemoryBase).where(MemoryBase.user_id == user_id)
result = await db.exec(stmt)
return list(result.all())
def list_for_user_stmt(self, user_id: uuid.UUID, flow_id: uuid.UUID | None = None): # type: ignore[return]
"""Return the SQLModel select statement for pagination at the API layer."""
stmt = select(MemoryBase).where(MemoryBase.user_id == user_id)
if flow_id is not None:
stmt = stmt.where(MemoryBase.flow_id == flow_id)
return stmt
async def get(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> MemoryBase | None:
async with session_scope() as db:
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
result = await db.exec(stmt)
return result.first()
async def update(
self,
memory_base_id: uuid.UUID,
user_id: uuid.UUID,
patch: MemoryBaseUpdate,
) -> MemoryBase | None:
"""Update mutable fields.
Threshold changes take effect on the NEXT auto-capture trigger; any
already-running ingestion task ignores the change (immutable args).
"""
async with session_scope() as db:
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
result = await db.exec(stmt)
mb = result.first()
if mb is None:
return None
for field, value in patch.model_dump(exclude_unset=True).items():
setattr(mb, field, value)
db.add(mb)
await db.commit()
await db.refresh(mb)
return mb
async def delete(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> bool:
"""Delete a MemoryBase and its associated KB directory."""
async with session_scope() as db:
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
result = await db.exec(stmt)
mb = result.first()
if mb is None:
return False
kb_name = mb.kb_name
kb_username = await resolve_kb_username(db, user_id)
# Cancel active ingestion jobs before removing the DB record
await cancel_active_jobs(memory_base_id=memory_base_id, db=db)
await db.delete(mb)
await db.commit()
# Delete the corresponding KB from disk (best-effort — DB already committed)
await delete_kb(kb_name=kb_name, kb_username=kb_username)
return True
# ------------------------------------------------------------------ #
# Sessions #
# ------------------------------------------------------------------ #
async def verify_ownership(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> None:
"""Raise ValueError if the Memory Base does not belong to user_id."""
async with session_scope() as db:
await self.get_memory_base_or_404(db, memory_base_id, user_id)
def sessions_stmt(self, memory_base_id: uuid.UUID, user_id: uuid.UUID): # type: ignore[return]
"""Return the select statement for persisted sessions, for use with apaginate.
Inline-joins MemoryBase to verify ownership in the SQL itself, so a
caller that forgets a pre-check cannot leak other users' sessions.
"""
return (
select(MemoryBaseSession)
.join(MemoryBase, MemoryBase.id == MemoryBaseSession.memory_base_id)
.where(MemoryBaseSession.memory_base_id == memory_base_id)
.where(MemoryBase.user_id == user_id)
.order_by(col(MemoryBaseSession.last_sync_at).desc())
)
# ------------------------------------------------------------------ #
# Ingestion delegation #
# ------------------------------------------------------------------ #
async def trigger_ingestion(
self,
memory_base_id: uuid.UUID,
user_id: uuid.UUID,
session_id: str,
) -> str:
return await _trigger_ingestion(
memory_base_id,
user_id,
session_id,
get_mb_or_raise=self.get_memory_base_or_404,
get_or_create_session=self._get_or_create_session,
)
async def on_flow_output(
self,
flow_id: uuid.UUID,
session_id: str,
job_id: uuid.UUID | None,
) -> None:
await _on_flow_output(
flow_id,
session_id,
job_id,
get_or_create_session=self._get_or_create_session,
)
async def check_mismatch(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> bool:
return await _check_mismatch(
memory_base_id,
user_id,
get_mb_or_raise=self.get_memory_base_or_404,
)
async def regenerate(self, memory_base_id: uuid.UUID, user_id: uuid.UUID) -> list[str]:
return await _regenerate(
memory_base_id,
user_id,
get_mb_or_raise=self.get_memory_base_or_404,
trigger_ingestion_fn=self.trigger_ingestion,
)
# ------------------------------------------------------------------ #
# Public query helpers #
# ------------------------------------------------------------------ #
async def get_memory_base_or_404(
self, db: AsyncSession, memory_base_id: uuid.UUID, user_id: uuid.UUID
) -> MemoryBase:
"""Fetch a MemoryBase or raise ValueError (mapped to 404 at the API layer)."""
stmt = select(MemoryBase).where(MemoryBase.id == memory_base_id).where(MemoryBase.user_id == user_id)
result = await db.exec(stmt)
mb = result.first()
if mb is None:
msg = f"MemoryBase {memory_base_id} not found"
raise ValueError(msg)
return mb
# ------------------------------------------------------------------ #
# Internal helpers #
# ------------------------------------------------------------------ #
async def _get_or_create_session(
self, db: AsyncSession, memory_base_id: uuid.UUID, session_id: str
) -> MemoryBaseSession:
stmt = (
select(MemoryBaseSession)
.where(MemoryBaseSession.memory_base_id == memory_base_id)
.where(MemoryBaseSession.session_id == session_id)
)
result = await db.exec(stmt)
mbs = result.first()
if mbs is None:
mbs = MemoryBaseSession(memory_base_id=memory_base_id, session_id=session_id)
db.add(mbs)
await db.commit()
await db.refresh(mbs)
return mbs

View File

@ -0,0 +1,444 @@
"""Background task for Memory Base ingestion.
Design principles enforced here:
- Cursor atomicity: cursor_id is NEVER updated before ingestion confirms success.
- Retry safety: If a job fails, cursor_id remains at the last known good position.
- Serialization: A per-(memory_base_id, session_id) distributed lock prevents concurrent
jobs from racing to write the same messages into Chroma. Uses PostgreSQL advisory locks
for cross-worker safety, with an in-process asyncio.Lock fallback for SQLite (dev/test).
The lock is acquired before any DB or Chroma access and released in a finally block.
- Live cursor: After acquiring the lock, the current cursor_id is re-read from the DB
(not the dispatch-time snapshot) so the pending message fetch always starts from the
true latest position, even if a prior job advanced the cursor while this job waited.
- Path safety: kb_path is validated against kb_root before any filesystem operation.
The actual Chroma write logic is shared with KB file ingestion via
``KBIngestionHelper.write_documents_to_chroma`` — no duplicate batching/retry code here.
Document building and KB metadata sync live in ``document_builders.py``.
"""
from __future__ import annotations
import asyncio
import hashlib
import types
import weakref
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import TYPE_CHECKING
from langchain_chroma import Chroma
from lfx.log.logger import logger
from sqlalchemy import text
from sqlmodel import Session, col, select
from langflow.api.utils.kb_helpers import KBIngestionHelper, KBStorageHelper
from langflow.services.database.models.memory_base.model import MemoryBaseSession, MemoryBaseWorkflowRun
from langflow.services.database.models.message.model import MessageTable
from langflow.services.deps import get_settings_service, session_scope
from langflow.services.memory_base.document_builders import build_documents_from_messages, sync_kb_metadata
from langflow.services.memory_base.kb_path_helpers import hash_session_id, validate_kb_path
if TYPE_CHECKING:
import uuid
from pathlib import Path
from langflow.services.jobs.service import JobService
@dataclass(frozen=True, slots=True)
class IngestionRequest:
"""Typed parameter bundle for ``ingest_memory_task``.
All fields needed to run an ingestion job are grouped here so callers
construct one object instead of threading 11+ loose kwargs.
"""
memory_base_id: uuid.UUID
session_id: str
flow_id: uuid.UUID
kb_name: str
kb_username: str
user_id: uuid.UUID
embedding_provider: str
embedding_model: str
cursor_id: uuid.UUID | None
task_job_id: uuid.UUID
job_service: JobService
# The ingestion lock timeout is read from settings (max_ingestion_timeout_secs).
# If the timeout expires before the lock is acquired, an asyncio.TimeoutError is raised.
# ---------------------------------------------------------------------------
# Distributed locking: PostgreSQL advisory locks with in-process fallback
# ---------------------------------------------------------------------------
# In multi-worker deployments, an asyncio.Lock is process-local and cannot
# serialize across workers. We use PostgreSQL session-level advisory locks
# keyed on a hash of (memory_base_id, session_id). For SQLite (dev/test) we
# fall back to the in-process asyncio.Lock which is sufficient for a single worker.
_session_ingestion_locks: weakref.WeakValueDictionary[tuple, asyncio.Lock] = weakref.WeakValueDictionary()
def _get_or_create_session_lock(key: tuple) -> asyncio.Lock:
"""Return the asyncio.Lock for the given key (SQLite fallback only)."""
lock = _session_ingestion_locks.get(key)
if lock is None:
lock = asyncio.Lock()
_session_ingestion_locks[key] = lock
return lock
def _compute_advisory_key(memory_base_id: uuid.UUID, session_id: str) -> int:
"""Compute a stable int64 advisory lock key from (memory_base_id, session_id)."""
raw = f"{memory_base_id}:{session_id}".encode()
return int(hashlib.sha256(raw).hexdigest()[:16], 16) % (2**63 - 1)
async def _is_postgres() -> bool:
"""Return True if the database backend is PostgreSQL."""
from langflow.services.deps import get_db_service
db_service = get_db_service()
return db_service.engine.dialect.name == "postgresql"
async def _pg_advisory_lock(db: Session, key: int) -> None:
"""Acquire a PostgreSQL session-level advisory lock with retry and timeout.
The lock is held on the specific connection of the shared 'db' session.
"""
timeout = get_settings_service().settings.max_ingestion_timeout_secs
deadline = asyncio.get_event_loop().time() + timeout
backoff = 0.1
max_backoff = 5.0
while True:
conn = await db.connection()
result = await conn.execute(text(f"SELECT pg_try_advisory_lock({key})"))
acquired = result.scalar()
if acquired:
return
remaining = deadline - asyncio.get_event_loop().time()
if remaining <= 0:
raise asyncio.TimeoutError
await asyncio.sleep(min(backoff, remaining))
backoff = min(backoff * 2, max_backoff)
async def _pg_advisory_unlock(db: Session, key: int) -> None:
"""Release a PostgreSQL session-level advisory lock on the shared session."""
conn = await db.connection()
await conn.execute(text(f"SELECT pg_advisory_unlock({key})"))
async def _acquire_session_lock(db: Session, memory_base_id: uuid.UUID, session_id: str) -> int | asyncio.Lock:
"""Acquire the distributed ingestion lock. Returns the key (PG) or Lock (SQLite)."""
timeout = get_settings_service().settings.max_ingestion_timeout_secs
if await _is_postgres():
key = _compute_advisory_key(memory_base_id, session_id)
await _pg_advisory_lock(db, key)
return key
lock = _get_or_create_session_lock((memory_base_id, session_id))
await asyncio.wait_for(lock.acquire(), timeout=timeout)
return lock
async def _release_session_lock(db: Session, lock_handle: int | asyncio.Lock) -> None:
"""Release the distributed ingestion lock."""
if isinstance(lock_handle, int):
await _pg_advisory_unlock(db, lock_handle)
else:
lock_handle.release()
async def _read_live_cursor(db: Session, memory_base_id: uuid.UUID, session_id: str) -> uuid.UUID | None:
"""Read current cursor_id from shared 'db' session inside the serialization lock."""
stmt = (
select(MemoryBaseSession.cursor_id)
.where(MemoryBaseSession.memory_base_id == memory_base_id)
.where(MemoryBaseSession.session_id == session_id)
)
result = await db.exec(stmt)
return result.first()
async def ingest_memory_task(*, request: IngestionRequest) -> dict:
"""Ingest pending output messages from a session into the target Knowledge Base.
Accepts a single ``IngestionRequest`` dataclass that bundles all required parameters.
Serialization: acquires a per-(memory_base_id, session_id) distributed lock before
any DB or Chroma access. Uses PostgreSQL advisory locks for cross-worker
serialization (multi-worker safe) with an in-process asyncio.Lock fallback for
SQLite. Concurrent jobs for the same session wait up to max_ingestion_timeout_secs;
if the lock cannot be acquired in time, asyncio.TimeoutError is re-raised so
execute_with_status records JobStatus.TIMED_OUT.
Live cursor: after acquiring the lock, the current cursor_id is re-read from the DB.
``cursor_id`` on the request is the dispatch-time snapshot kept only for logging.
"""
# Unpack for readability within the function body
memory_base_id = request.memory_base_id
session_id = request.session_id
flow_id = request.flow_id
kb_name = request.kb_name
kb_username = request.kb_username
user_id = request.user_id
embedding_provider = request.embedding_provider
embedding_model = request.embedding_model
cursor_id = request.cursor_id
task_job_id = request.task_job_id
job_service = request.job_service
hashed_sid = hash_session_id(session_id)
await logger.adebug(
"Ingestion job started | memory_base=%s session=%s dispatch_cursor=%s job=%s",
memory_base_id,
hashed_sid,
cursor_id,
task_job_id,
)
kb_root = KBStorageHelper.get_root_path()
if not kb_root:
msg = "Knowledge base root path is not configured"
raise RuntimeError(msg)
kb_path: Path = kb_root / kb_username / kb_name
# ---- Path traversal guard ----
validate_kb_path(kb_root, kb_path)
# ---- 0. Acquire per-session serialization lock ----
async with session_scope() as db:
try:
lock_handle = await _acquire_session_lock(db, memory_base_id, session_id)
except asyncio.TimeoutError:
await logger.awarning(
"Ingestion lock wait timeout | memory_base=%s session=%s job=%s.",
memory_base_id,
hashed_sid,
task_job_id,
)
raise
try:
# ---- 0b. Re-read live cursor inside the lock ----
live_cursor_id = await _read_live_cursor(db, memory_base_id, session_id)
await logger.adebug(
"Ingestion lock acquired | memory_base=%s session=%s live_cursor=%s job=%s",
memory_base_id,
hashed_sid,
live_cursor_id,
task_job_id,
)
# ---- 1. Fetch pending output messages for this session ----
messages = await _fetch_pending_messages(
db,
flow_id=flow_id,
session_id=session_id,
cursor_id=live_cursor_id,
)
if not messages:
await logger.ainfo(
"MemoryBase %s / session %s: no pending messages, skipping.", memory_base_id, hashed_sid
)
return {"message": "No pending messages", "ingested": 0}
# ---- 2. Build documents from messages ----
job_id_str = str(task_job_id)
documents = build_documents_from_messages(
messages, session_id=session_id, flow_id=str(flow_id), job_id=job_id_str
)
if not documents:
return {"message": "No non-empty messages to ingest", "ingested": 0}
# ---- 3. Check cancellation before touching the vector store ----
if await KBIngestionHelper.is_job_cancelled(job_service, task_job_id):
return {"message": "Job cancelled before ingestion", "ingested": 0}
# ---- 4. Open Chroma, write, then sync KB metadata ----
user_stub = types.SimpleNamespace(id=user_id)
embeddings = await KBIngestionHelper.build_embeddings(embedding_provider, embedding_model, user_stub)
client = KBStorageHelper.get_fresh_chroma_client(kb_path)
written = 0
try:
chroma = Chroma(client=client, embedding_function=embeddings, collection_name=kb_name)
written = await KBIngestionHelper.write_documents_to_chroma(
documents=documents,
chroma=chroma,
task_job_id=task_job_id,
job_service=job_service,
)
if written == len(documents):
await asyncio.to_thread(sync_kb_metadata, kb_path=kb_path, chroma=chroma)
except Exception:
await logger.aerror(
"Ingestion write failed | memory_base=%s session=%s job=%s. Rolling back partial writes...",
memory_base_id,
hashed_sid,
task_job_id,
)
await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name)
raise
finally:
KBStorageHelper.release_chroma_resources(kb_path)
if written < len(documents):
await logger.awarning("Ingestion job %s was cancelled. Cleaning up partial data...", task_job_id)
await KBIngestionHelper.cleanup_chroma_chunks_by_job(task_job_id, kb_path, kb_name)
return {"message": "Job cancelled during ingestion", "ingested": 0}
# ---- 5. Bulk-stamp ingestion metadata ----
await _mark_messages_ingested(db, messages=messages, job_id=task_job_id, memory_base_id=memory_base_id)
# ---- 6. Update cursor atomically ONLY after confirmed success ----
last_message_id = messages[-1].id
ingested_count = len(messages)
await _advance_cursor(
db,
memory_base_id=memory_base_id,
session_id=session_id,
new_cursor_id=last_message_id,
ingested_count=ingested_count,
task_job_id=task_job_id,
)
await logger.ainfo(
"Ingestion job finished | memory_base=%s session=%s job=%s ingested=%d",
memory_base_id,
hashed_sid,
task_job_id,
ingested_count,
)
return {"message": "Success", "ingested": ingested_count}
finally:
await _release_session_lock(db, lock_handle)
async def _fetch_pending_messages(
db: Session,
*,
flow_id: uuid.UUID,
session_id: str,
cursor_id: uuid.UUID | None,
) -> list[MessageTable]:
"""Fetch all messages for this session that come after cursor_id using shared session."""
from sqlalchemy import and_, or_
stmt = (
select(MessageTable)
.where(MessageTable.flow_id == flow_id)
.where(MessageTable.session_id == session_id)
.order_by(col(MessageTable.timestamp).asc(), col(MessageTable.id).asc())
)
if cursor_id is not None:
cursor_stmt = select(MessageTable.timestamp, MessageTable.id).where(MessageTable.id == cursor_id)
result = await db.exec(cursor_stmt)
cursor_row = result.first()
if cursor_row:
cursor_ts, c_id = cursor_row
stmt = stmt.where(
or_(
col(MessageTable.timestamp) > cursor_ts,
and_(
col(MessageTable.timestamp) == cursor_ts,
col(MessageTable.id) > c_id,
),
)
)
result = await db.exec(stmt)
return list(result.all())
async def _mark_messages_ingested(
db: Session,
*,
messages: list[MessageTable],
job_id: uuid.UUID,
memory_base_id: uuid.UUID,
) -> None:
"""Batch-insert ingestion records for all successfully ingested messages using shared session."""
from uuid import uuid4 as _uuid4
from langflow.services.database.models.memory_base.model import MessageIngestionRecord
ingested_at = datetime.now(timezone.utc)
rows = [
{
"id": _uuid4(),
"message_id": msg.id,
"memory_base_id": memory_base_id,
"job_id": job_id,
"session_id": msg.session_id,
"ingested_at": ingested_at,
}
for msg in messages
]
conn = await db.connection()
if conn.dialect.name == "postgresql":
from sqlalchemy.dialects.postgresql import insert as pg_insert
stmt = pg_insert(MessageIngestionRecord).values(rows).on_conflict_do_nothing()
else:
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
stmt = sqlite_insert(MessageIngestionRecord).values(rows).on_conflict_do_nothing()
await db.exec(stmt) # type: ignore[call-overload]
await db.commit()
async def _advance_cursor(
db: Session,
*,
memory_base_id: uuid.UUID,
session_id: str,
new_cursor_id: uuid.UUID,
ingested_count: int,
task_job_id: uuid.UUID,
) -> None:
"""Atomically advance the cursor using the shared 'db' session."""
from sqlalchemy import update as sa_update
stmt = (
select(MemoryBaseSession)
.where(MemoryBaseSession.memory_base_id == memory_base_id)
.where(MemoryBaseSession.session_id == session_id)
)
result = await db.exec(stmt)
mbs = result.first()
if mbs is None:
await logger.awarning(
"MemoryBaseSession for (%s, %s) vanished before cursor update.",
memory_base_id,
hash_session_id(session_id),
)
return
mbs.cursor_id = new_cursor_id
mbs.total_processed += ingested_count
mbs.last_sync_at = datetime.now(timezone.utc)
db.add(mbs)
# Stamp all pending workflow run rows for this session.
await db.exec( # type: ignore[call-overload]
sa_update(MemoryBaseWorkflowRun)
.where(MemoryBaseWorkflowRun.memory_base_id == memory_base_id)
.where(MemoryBaseWorkflowRun.session_id == session_id)
.where(MemoryBaseWorkflowRun.ingestion_job_id == None) # noqa: E711
.values(ingestion_job_id=task_job_id)
)
await db.commit()

View File

@ -22,3 +22,4 @@ class ServiceType(str, Enum):
MCP_COMPOSER_SERVICE = "mcp_composer_service"
JOB_SERVICE = "jobs_service"
FLOW_EVENTS_SERVICE = "flow_events_service"
MEMORY_BASE_SERVICE = "memory_base_service"

View File

@ -94,6 +94,12 @@ async def setup_superuser(settings_service: SettingsService, session: AsyncSessi
)
if user is not None:
await logger.adebug("Superuser created successfully.")
# When the default superuser is recreated (e.g. after a DB reset in
# AUTO_LOGIN mode) the per-user MCP servers config file saved under the
# previous UUID becomes orphaned on disk. Best-effort recover it so
# users don't lose their MCP server configuration across restarts.
if is_default and settings_service.auth_settings.AUTO_LOGIN:
await migrate_orphaned_mcp_servers_config(session, settings_service, user)
except Exception as exc:
logger.exception(exc)
msg = "Could not create superuser. Please create a superuser manually."
@ -103,6 +109,149 @@ async def setup_superuser(settings_service: SettingsService, session: AsyncSessi
settings_service.auth_settings.reset_credentials()
async def migrate_orphaned_mcp_servers_config(
session: AsyncSession,
settings_service: SettingsService,
current_user,
) -> bool:
"""Best-effort recovery of MCP servers config files orphaned by a DB reset.
The MCP servers config is persisted on disk at
``{config_dir}/{user_id}/_mcp_servers_{user_id}.json`` and tracked in the DB
via a ``File`` row. When Langflow starts with a fresh database but the same
config directory (common in containerized deployments without a persisted
DB volume), the default superuser is recreated with a new UUID and the
previously saved MCP config files become unreachable.
Recovery rules (intentionally conservative to avoid importing another user's
config — MCP server entries can contain ``env`` and ``headers`` auth material):
* If the new user already has an MCP config file on disk without a matching
``File`` row, re-register the row (self-heal a partial previous migration).
* If exactly one orphaned ``_mcp_servers_{uuid}.json`` is found in the config
directory, migrate it.
* If multiple orphans are found, skip and log — we can't safely identify the
previous default superuser's file without extra metadata, so leave manual
recovery to the operator.
Returns True when a file was migrated or a missing DB row was restored.
"""
from pathlib import Path
from uuid import UUID
import aiofiles
import anyio
from langflow.services.database.models.file.model import File as UserFile
try:
config_dir_value = settings_service.settings.config_dir
if not config_dir_value:
return False
config_dir = Path(config_dir_value)
if not config_dir.exists() or not config_dir.is_dir():
return False
# The current user's DB record is fresh; nothing to migrate if they
# somehow already have an MCP config row (defensive guard).
name_without_ext = f"_mcp_servers_{current_user.id}"
existing_stmt = (
select(UserFile).where(UserFile.user_id == current_user.id).where(UserFile.name == name_without_ext)
)
if (await session.exec(existing_stmt)).first() is not None:
return False
current_user_dir = str(current_user.id)
new_dir = config_dir / current_user_dir
new_filename = f"_mcp_servers_{current_user.id}.json"
new_file_path = new_dir / new_filename
db_path = f"{current_user.id}/{new_filename}"
# Case 1: a previous migration attempt copied the file but failed before
# committing the DB row. Re-register the existing file instead of
# returning early and leaving the user with an invisible config.
if new_file_path.exists():
try:
size = new_file_path.stat().st_size
except OSError as exc:
await logger.awarning(
"Cannot stat existing MCP config %s while self-healing DB row: %s",
new_file_path,
exc,
)
return False
session.add(UserFile(user_id=current_user.id, name=name_without_ext, path=db_path, size=size))
await session.commit()
await logger.ainfo(
"Restored missing MCP servers config DB row for user %s from existing file %s",
current_user.id,
new_file_path,
)
return True
def _find_orphans() -> list[tuple[float, Path]]:
orphans: list[tuple[float, Path]] = []
for entry in config_dir.iterdir():
if not entry.is_dir() or entry.name == current_user_dir:
continue
try:
UUID(entry.name)
except ValueError:
continue
mcp_path = entry / f"_mcp_servers_{entry.name}.json"
if mcp_path.is_file():
try:
mtime = mcp_path.stat().st_mtime
except OSError:
continue
orphans.append((mtime, mcp_path))
return orphans
orphans = await anyio.to_thread.run_sync(_find_orphans)
if not orphans:
return False
# Ambiguous: more than one candidate could belong to different users.
# Refuse to migrate rather than risk importing unrelated auth material.
if len(orphans) > 1:
orphan_paths = ", ".join(str(p) for _, p in sorted(orphans, key=lambda i: i[0], reverse=True))
await logger.awarning(
"Found %d orphaned MCP servers config files in %s; skipping automatic "
"migration to avoid restoring the wrong one. Move the intended file to "
"%s to recover. Candidates: %s",
len(orphans),
config_dir,
new_file_path,
orphan_paths,
)
return False
_, orphan_path = orphans[0]
async with aiofiles.open(str(orphan_path), "rb") as src:
data = await src.read()
await anyio.to_thread.run_sync(lambda: new_dir.mkdir(parents=True, exist_ok=True))
async with aiofiles.open(str(new_file_path), "wb") as dst:
await dst.write(data)
session.add(UserFile(user_id=current_user.id, name=name_without_ext, path=db_path, size=len(data)))
await session.commit()
await logger.ainfo(
"Migrated orphaned MCP servers config from %s to user %s",
orphan_path,
current_user.id,
)
except Exception as exc: # noqa: BLE001
# Never let migration failure block startup.
await logger.awarning("Failed to migrate orphaned MCP servers config: %s", exc)
return False
else:
return True
async def teardown_superuser(settings_service, session: AsyncSession) -> None:
"""Teardown the superuser."""
# If AUTO_LOGIN is True, we will remove the default superuser

View File

@ -39,6 +39,7 @@ class TestAgentComponent(ComponentTestBaseWithoutClient):
return {
"_type": "Agent",
"add_current_date_tool": True,
"add_calculator_tool": True,
"agent_description": "A helpful agent",
"model": MockLanguageModel(),
"handle_parsing_errors": True,
@ -450,6 +451,199 @@ class TestAgentComponent(ComponentTestBaseWithoutClient):
# Note: The provider-specific field name mapping happens inside get_llm,
# so we just verify max_tokens is passed correctly
async def test_should_append_calculator_tool_when_add_calculator_toggle_is_true(
self, component_class, default_kwargs
):
"""Calculator tool is appended when the toggle is enabled.
Given add_calculator_tool=True, When get_agent_requirements runs,
Then self.tools contains a StructuredTool derived from CalculatorComponent.
"""
from unittest.mock import AsyncMock
from langchain_core.tools import StructuredTool
default_kwargs["add_calculator_tool"] = True
default_kwargs["add_current_date_tool"] = False # isolate: only calculator
component = await self.component_setup(component_class, default_kwargs)
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
component.get_memory_data = AsyncMock(return_value=[])
component._get_shared_callbacks = list
component.set_tools_callbacks = lambda *_: None
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
mock_get_llm.return_value = MockLanguageModel()
_, _, tools = await component.get_agent_requirements()
assert len(tools) == 1
assert isinstance(tools[0], StructuredTool)
assert "evaluate" in tools[0].name.lower(), f"Expected a Calculator-derived tool; got name={tools[0].name!r}"
async def test_should_not_append_calculator_tool_when_add_calculator_toggle_is_false(
self, component_class, default_kwargs
):
"""Calculator tool is skipped when the toggle is disabled.
Given add_calculator_tool=False, When get_agent_requirements runs,
Then no Calculator tool is appended to self.tools.
"""
from unittest.mock import AsyncMock
default_kwargs["add_calculator_tool"] = False
default_kwargs["add_current_date_tool"] = False
component = await self.component_setup(component_class, default_kwargs)
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
component.get_memory_data = AsyncMock(return_value=[])
component._get_shared_callbacks = list
component.set_tools_callbacks = lambda *_: None
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
mock_get_llm.return_value = MockLanguageModel()
_, _, tools = await component.get_agent_requirements()
assert tools == []
def test_should_replace_current_date_and_model_name_when_both_placeholders_present(self, component_class):
"""Unit test: helper replaces both placeholders with concrete values."""
component = component_class()
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
prompt = "Today is {current_date}. You are powered by {model_name}."
result = component._inject_dynamic_prompt_values(prompt)
assert "{current_date}" not in result
assert "{model_name}" not in result
assert "gpt-4o" in result
def test_should_leave_literal_braces_untouched_when_prompt_has_no_known_placeholders(self, component_class):
"""Adversarial: prompts with literal JSON like {"key": 1} must not raise and must stay intact."""
component = component_class()
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
prompt = 'Respond with JSON: {"key": 1, "nested": {"a": [1, 2]}}.'
result = component._inject_dynamic_prompt_values(prompt)
assert result == prompt
def test_should_return_empty_when_prompt_is_empty(self, component_class):
"""Edge case: empty/None prompt is returned as-is without raising."""
component = component_class()
assert component._inject_dynamic_prompt_values("") == ""
assert component._inject_dynamic_prompt_values(None) is None
async def test_should_inject_dynamic_values_into_system_prompt_when_message_response_runs(
self, component_class, default_kwargs
):
"""Integration: message_response must call self.set with the resolved system_prompt."""
from unittest.mock import AsyncMock, MagicMock
default_kwargs["system_prompt"] = "Powered by {model_name}."
default_kwargs["add_calculator_tool"] = False
default_kwargs["add_current_date_tool"] = False
component = await self.component_setup(component_class, default_kwargs)
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
component.get_memory_data = AsyncMock(return_value=[])
component._get_shared_callbacks = list
component.set_tools_callbacks = lambda *_: None
captured: dict = {}
def fake_set(**kwargs):
captured.update(kwargs)
return component
component.set = fake_set
component.create_agent_runnable = MagicMock(return_value=MagicMock())
component.run_agent = AsyncMock(return_value=MagicMock())
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
mock_get_llm.return_value = MockLanguageModel()
await component.message_response()
assert captured.get("system_prompt") == "Powered by gpt-4o."
async def test_should_not_mutate_format_instructions_when_json_response_runs(self, component_class, default_kwargs):
"""Regression: injection must only touch agent_instructions, not format_instructions.
Ensures literal {current_date}/{model_name} tokens in user-authored
format_instructions survive intact while the main system_prompt is
still replaced by the helper.
"""
from unittest.mock import AsyncMock, MagicMock
default_kwargs["system_prompt"] = "Powered by {model_name}."
default_kwargs["format_instructions"] = "Return JSON with fields {current_date} and {model_name} preserved."
default_kwargs["add_calculator_tool"] = False
default_kwargs["add_current_date_tool"] = False
component = await self.component_setup(component_class, default_kwargs)
component.model = [{"name": "gpt-4o", "provider": "OpenAI", "metadata": {}}]
component.get_memory_data = AsyncMock(return_value=[])
component._get_shared_callbacks = list
component.set_tools_callbacks = lambda *_: None
captured: dict = {}
def fake_set(**kwargs):
captured.update(kwargs)
return component
component.set = fake_set
component.create_agent_runnable = MagicMock(return_value=MagicMock())
component.run_agent = AsyncMock(return_value=MagicMock(content="{}"))
with patch("lfx.components.models_and_agents.agent.get_llm") as mock_get_llm:
mock_get_llm.return_value = MockLanguageModel()
await component.json_response()
prompt = captured.get("system_prompt") or ""
assert "Powered by gpt-4o." in prompt, "agent_instructions should have placeholders replaced"
assert "{current_date}" in prompt, "format_instructions literal braces must survive"
assert "{model_name} preserved" in prompt, "format_instructions literal braces must survive"
async def test_should_accept_add_calculator_tool_in_default_keys(self, component_class, default_kwargs):
"""update_build_config's default_keys validation must include add_calculator_tool."""
from lfx.schema.dotdict import dotdict
with patch("lfx.components.models_and_agents.agent.get_language_model_options") as mock_opts:
mock_opts.return_value = [
{
"name": "gpt-4o",
"provider": "OpenAI",
"icon": "OpenAI",
"metadata": {
"model_class": "ChatOpenAI",
"model_name_param": "model",
"api_key_param": "api_key",
},
}
]
component = await self.component_setup(component_class, default_kwargs)
frontend_node = component.to_frontend_node()
build_config = frontend_node["data"]["node"]["template"]
# add_calculator_tool must be present in the build_config already; if not,
# update_build_config will error listing it as missing.
assert "add_calculator_tool" in build_config
updated_config = await component.update_build_config(
dotdict(build_config), mock_opts.return_value, field_name="model"
)
assert "add_calculator_tool" in updated_config
def test_should_have_placeholders_in_default_system_prompt(self, component_class):
"""Default system_prompt ships with placeholders for the dynamic injection.
Ensures {current_date} and {model_name} are visible on a fresh agent so
that the dynamic injection has an observable effect out-of-the-box.
"""
prompt_input = next(
(inp for inp in component_class.inputs if getattr(inp, "name", None) == "system_prompt"),
None,
)
assert prompt_input is not None
assert "{current_date}" in prompt_input.value
assert "{model_name}" in prompt_input.value
class TestAgentComponentWithClient(ComponentTestBaseWithClient):
@pytest.fixture

View File

@ -1,5 +1,6 @@
from pathlib import Path
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from lfx.components.chroma import ChromaVectorStoreComponent
@ -8,6 +9,37 @@ from lfx.schema.data import Data
from tests.base import ComponentTestBaseWithoutClient, VersionComponentMapping
def test_remote_chroma_server_uses_http_client() -> None:
mock_client = MagicMock()
mock_chroma = MagicMock()
mock_chroma.get.return_value = {"ids": [], "documents": [], "metadatas": []}
with (
patch("chromadb.HttpClient", return_value=mock_client) as mock_http_client,
patch("langchain_chroma.Chroma", return_value=mock_chroma) as mock_chroma_class,
):
component = ChromaVectorStoreComponent().set(
collection_name="remote_collection",
persist_directory=None,
embedding=None,
chroma_server_host="chroma.example.com",
chroma_server_http_port=8100,
chroma_server_ssl_enabled=True,
ingest_data=[],
limit=None,
)
assert component.build_vector_store() is mock_chroma
mock_http_client.assert_called_once_with(host="chroma.example.com", port=8100, ssl=True)
mock_chroma_class.assert_called_once_with(
persist_directory=None,
client=mock_client,
embedding_function=None,
collection_name="remote_collection",
)
@pytest.mark.api_key_required
class TestChromaVectorStoreComponent(ComponentTestBaseWithoutClient):
@pytest.fixture

View File

@ -714,7 +714,7 @@ class TestPerformIngestionTask:
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client")
@patch("langflow.api.utils.kb_helpers.Chroma")
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper._build_embeddings", new_callable=AsyncMock)
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.build_embeddings", new_callable=AsyncMock)
@patch("langflow.api.utils.kb_helpers.KBAnalysisHelper.get_metadata")
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_directory_size")
@patch("langflow.api.utils.kb_helpers.KBAnalysisHelper.update_text_metrics")
@ -764,7 +764,7 @@ class TestPerformIngestionTask:
@patch("langflow.api.utils.kb_helpers.KBStorageHelper.get_fresh_chroma_client")
@patch("langflow.api.utils.kb_helpers.Chroma")
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper._build_embeddings", new_callable=AsyncMock)
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.build_embeddings", new_callable=AsyncMock)
@patch("langflow.api.utils.kb_helpers.KBIngestionHelper.cleanup_chroma_chunks_by_job", new_callable=AsyncMock)
async def test_perform_ingestion_rollback(
self, mock_cleanup, mock_build, mock_chroma, mock_fresh_client, mock_kb_path

View File

@ -0,0 +1,247 @@
"""Tests for migrate_orphaned_mcp_servers_config in langflow.services.utils.
Verifies that MCP server config files written under a previous default
superuser's UUID are picked up and migrated to the new default superuser
when the database is reset but the config directory is preserved
(typical of containerized deployments without a persisted DB volume).
"""
from __future__ import annotations
import json
import os
from typing import TYPE_CHECKING
from uuid import uuid4
import pytest
if TYPE_CHECKING:
from pathlib import Path
from langflow.services.database.models.file.model import File as UserFile
from langflow.services.deps import get_settings_service, session_scope
from langflow.services.utils import migrate_orphaned_mcp_servers_config
from sqlmodel import select
@pytest.fixture
async def initialized_services(monkeypatch, tmp_path):
"""Initialize DB + services with an isolated config dir."""
from langflow.services.utils import initialize_services, teardown_services
from lfx.services.manager import get_service_manager
db_path = tmp_path / "test.db"
config_dir = tmp_path / "config"
config_dir.mkdir(parents=True, exist_ok=True)
monkeypatch.setenv("LANGFLOW_DATABASE_URL", f"sqlite:///{db_path}")
monkeypatch.setenv("LANGFLOW_CONFIG_DIR", str(config_dir))
monkeypatch.setenv("LANGFLOW_AUTO_LOGIN", "true")
get_service_manager().factories.clear()
get_service_manager().services.clear()
await initialize_services()
yield config_dir
await teardown_services()
def _write_orphan(config_dir: Path, payload: dict, *, mtime_offset: float = 0.0) -> Path:
"""Create an orphaned _mcp_servers_{uuid}.json file in a UUID-named folder."""
orphan_id = uuid4()
orphan_dir = config_dir / str(orphan_id)
orphan_dir.mkdir(parents=True, exist_ok=True)
orphan_path = orphan_dir / f"_mcp_servers_{orphan_id}.json"
orphan_path.write_text(json.dumps(payload))
if mtime_offset:
stat = orphan_path.stat()
os.utime(orphan_path, (stat.st_atime + mtime_offset, stat.st_mtime + mtime_offset))
return orphan_path
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_migrate_orphaned_mcp_servers_config_recovers_previous_user_config(
initialized_services,
):
"""A single orphaned file is migrated to the current default superuser."""
config_dir: Path = initialized_services
expected_payload = {"mcpServers": {"my-server": {"command": "uvx", "args": ["mcp-proxy"]}}}
orphan_path = _write_orphan(config_dir, expected_payload)
settings = get_settings_service()
async with session_scope() as session:
from langflow.services.database.models.user.model import User
from lfx.services.settings.constants import DEFAULT_SUPERUSER
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
assert user is not None, "default superuser should exist after initialize_services"
# Simulate the fresh-DB scenario: the user has no MCP config row yet.
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
assert (await session.exec(stmt)).first() is None
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
assert migrated is True
# DB record should exist and point at the new user-specific path.
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
new_record = (await session.exec(stmt)).first()
assert new_record is not None
assert new_record.path == f"{user.id}/_mcp_servers_{user.id}.json"
# File should live under the new user's folder with the same contents.
target = config_dir / str(user.id) / f"_mcp_servers_{user.id}.json"
assert target.exists()
assert json.loads(target.read_text()) == expected_payload
# Orphan source is left intact (best-effort copy, not destructive move).
assert orphan_path.exists()
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_migrate_orphaned_mcp_servers_config_skips_when_multiple_orphans(
initialized_services,
):
"""With multiple orphan candidates we refuse to guess and leave state untouched.
MCP server entries can contain env/headers auth material, so importing an
unrelated user's config would be a security hazard. Operators must resolve
the ambiguity manually.
"""
config_dir: Path = initialized_services
old_payload = {"mcpServers": {"old": {}}}
new_payload = {"mcpServers": {"new": {}}}
old_path = _write_orphan(config_dir, old_payload, mtime_offset=-3600)
new_path = _write_orphan(config_dir, new_payload)
settings = get_settings_service()
async with session_scope() as session:
from langflow.services.database.models.user.model import User
from lfx.services.settings.constants import DEFAULT_SUPERUSER
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
assert migrated is False
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
assert (await session.exec(stmt)).first() is None
target = config_dir / str(user.id) / f"_mcp_servers_{user.id}.json"
assert not target.exists()
# Both orphans are preserved on disk for manual recovery.
assert json.loads(old_path.read_text()) == old_payload
assert json.loads(new_path.read_text()) == new_payload
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_migrate_orphaned_mcp_servers_config_no_orphans_is_noop(
initialized_services,
):
"""With no orphaned files the function returns False and does not touch the DB."""
config_dir: Path = initialized_services
# No UUID-named subdirectories should exist.
from uuid import UUID
def _is_uuid_dir(p: Path) -> bool:
if not p.is_dir():
return False
try:
UUID(p.name)
except ValueError:
return False
return True
assert not [p for p in config_dir.iterdir() if _is_uuid_dir(p)]
settings = get_settings_service()
async with session_scope() as session:
from langflow.services.database.models.user.model import User
from lfx.services.settings.constants import DEFAULT_SUPERUSER
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
assert migrated is False
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
assert (await session.exec(stmt)).first() is None
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_migrate_orphaned_mcp_servers_config_self_heals_missing_db_row(
initialized_services,
):
"""Re-register missing DB rows when the on-disk file already exists.
If the user's file is on disk but the DB row is missing, the migration
should recreate the row instead of leaving the config invisible. This
covers recovery from a previous migration attempt that wrote the file
but crashed before committing the DB row.
"""
config_dir: Path = initialized_services
settings = get_settings_service()
async with session_scope() as session:
from langflow.services.database.models.user.model import User
from lfx.services.settings.constants import DEFAULT_SUPERUSER
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
# Simulate a file-exists / DB-row-missing state.
existing_dir = config_dir / str(user.id)
existing_dir.mkdir(parents=True, exist_ok=True)
existing_payload = {"mcpServers": {"keep-me": {}}}
existing_path = existing_dir / f"_mcp_servers_{user.id}.json"
existing_path.write_text(json.dumps(existing_payload))
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
assert migrated is True
stmt = select(UserFile).where(UserFile.user_id == user.id).where(UserFile.name == f"_mcp_servers_{user.id}")
row = (await session.exec(stmt)).first()
assert row is not None
assert row.path == f"{user.id}/_mcp_servers_{user.id}.json"
assert row.size == existing_path.stat().st_size
# File contents are untouched.
assert json.loads(existing_path.read_text()) == existing_payload
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_migrate_orphaned_mcp_servers_config_skips_when_row_already_present(
initialized_services,
):
"""When the DB row already exists, do nothing — avoids double-registration."""
config_dir: Path = initialized_services
_write_orphan(config_dir, {"mcpServers": {"orphan": {}}})
settings = get_settings_service()
async with session_scope() as session:
from langflow.services.database.models.user.model import User
from lfx.services.settings.constants import DEFAULT_SUPERUSER
user = (await session.exec(select(User).where(User.username == DEFAULT_SUPERUSER))).first()
# Pre-register an MCP file row to simulate an already-migrated user.
session.add(
UserFile(
user_id=user.id,
name=f"_mcp_servers_{user.id}",
path=f"{user.id}/_mcp_servers_{user.id}.json",
size=0,
)
)
await session.commit()
migrated = await migrate_orphaned_mcp_servers_config(session, settings, user)
assert migrated is False

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@ -322,3 +322,15 @@ async def test_session_metadata_retrieval():
messages = await aget_messages(sender="User", session_id=session_id)
assert len(messages) == 1
assert messages[0].session_metadata == initial_metadata
@pytest.mark.usefixtures("client")
async def test_messageupdate_with_session_metadata(sample_session_metadata):
"""Test MessageUpdate schema with session_metadata."""
from langflow.services.database.models.message.model import MessageUpdate
message_update = MessageUpdate(
text="Updated text",
session_metadata=sample_session_metadata,
)
assert message_update.session_metadata == sample_session_metadata

View File

@ -19,10 +19,8 @@ export const useDeleteProviderAccount: useMutationFunctionType<
);
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["useDeleteProviderAccount"], fn, {
...options,
retry: false,
onSuccess: (...args) => {
queryClient.refetchQueries({ queryKey: ["useGetProviderAccounts"] });
options?.onSuccess?.(...args);

View File

@ -17,10 +17,8 @@ export const useDeleteDeployment: useMutationFunctionType<
await api.delete(`${getURL("DEPLOYMENTS")}/${deployment_id}`);
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["useDeleteDeployment"], fn, {
...options,
retry: false,
onSuccess: (...args) => {
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
options?.onSuccess?.(...args);

View File

@ -50,10 +50,8 @@ export const usePatchDeployment: useMutationFunctionType<
return res.data;
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["usePatchDeployment"], fn, {
...options,
retry: false,
onSuccess: (...args) => {
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
queryClient.removeQueries({

View File

@ -26,10 +26,8 @@ export const usePatchSnapshot: useMutationFunctionType<
return data;
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["usePatchSnapshot"], fn, {
...options,
retry: false,
onSuccess: (...args) => {
queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
options?.onSuccess?.(...args);

View File

@ -47,6 +47,5 @@ export const usePostDeploymentRun: useMutationFunctionType<
return res.data;
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["usePostDeploymentRun"], fn, { ...options, retry: false });
return mutate(["usePostDeploymentRun"], fn, options);
};

View File

@ -53,10 +53,8 @@ export const usePostDeployment: useMutationFunctionType<
return res.data;
};
// TODO: Add retries for transient server-side errors (5xx, timeouts).
return mutate(["usePostDeployment"], fn, {
...options,
retry: false,
onSuccess: () => {
return queryClient.refetchQueries({ queryKey: ["useGetDeployments"] });
},

View File

@ -0,0 +1,189 @@
// biome-ignore-all lint/suspicious/noExplicitAny: test mocks
const mockQueryClient = {
invalidateQueries: jest.fn(),
};
let mockCapturedQueryOptions: any = null;
let mockCapturedMutationOptions: any = null;
jest.mock("@tanstack/react-query", () => ({
useQueryClient: jest.fn(() => mockQueryClient),
useQuery: jest.fn((options: any) => {
mockCapturedQueryOptions = options;
return { data: undefined, isLoading: false };
}),
useMutation: jest.fn((options: any) => {
mockCapturedMutationOptions = options;
return { mutate: jest.fn(), mutateAsync: jest.fn() };
}),
}));
import { UseRequestProcessor } from "../request-processor";
// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
// `axios.isAxiosError` checks `error.isAxiosError === true`, so fixtures must
// set that flag to be classified as HTTP errors rather than transient.
const axiosError = (status: number) => ({
isAxiosError: true,
response: { status },
});
const axiosNetworkError = () => ({ isAxiosError: true, response: undefined });
const nonAxiosError = () => new Error("Network Error");
beforeEach(() => {
jest.clearAllMocks();
mockCapturedQueryOptions = null;
mockCapturedMutationOptions = null;
});
// ---------------------------------------------------------------------------
// Queries: retry policy
// ---------------------------------------------------------------------------
describe("UseRequestProcessor.query retry policy", () => {
const setup = (options: any = {}) => {
const { query } = UseRequestProcessor();
query(["k"], async () => ({}), options);
if (mockCapturedQueryOptions == null) {
throw new Error("query was not called by UseRequestProcessor");
}
return mockCapturedQueryOptions;
};
it("does not retry on any 4xx response", () => {
const { retry } = setup();
for (const status of [400, 401, 403, 404, 409, 422, 429, 499]) {
expect(retry(0, axiosError(status))).toBe(false);
}
});
it("retries up to 5 times on 5xx responses", () => {
const { retry } = setup();
for (const status of [500, 502, 503, 504]) {
expect(retry(0, axiosError(status))).toBe(true);
expect(retry(4, axiosError(status))).toBe(true);
expect(retry(5, axiosError(status))).toBe(false);
}
});
it("retries up to 5 times on axios errors with no response", () => {
const { retry } = setup();
expect(retry(0, axiosNetworkError())).toBe(true);
expect(retry(4, axiosNetworkError())).toBe(true);
expect(retry(5, axiosNetworkError())).toBe(false);
});
it("treats non-axios / unknown errors as transient (retries)", () => {
const { retry } = setup();
expect(retry(0, undefined)).toBe(true);
expect(retry(0, {})).toBe(true);
expect(retry(0, { response: {} })).toBe(true);
expect(retry(0, nonAxiosError())).toBe(true);
});
it("allows per-call options.retry to override the default", () => {
const { retry } = setup({ retry: false });
expect(retry).toBe(false);
});
it("allows per-call options.retry as a number to override the default", () => {
const { retry } = setup({ retry: 10 });
expect(retry).toBe(10);
});
it("allows per-call options.retryDelay to override the default", () => {
const customDelay = jest.fn(() => 42);
const { retryDelay } = setup({ retryDelay: customDelay });
expect(retryDelay).toBe(customDelay);
});
});
// ---------------------------------------------------------------------------
// Mutations: retry policy
// ---------------------------------------------------------------------------
describe("UseRequestProcessor.mutate retry policy", () => {
const setup = (options: any = {}) => {
const { mutate } = UseRequestProcessor();
mutate(["k"], async () => ({}), options);
if (mockCapturedMutationOptions == null) {
throw new Error("mutate was not called by UseRequestProcessor");
}
return mockCapturedMutationOptions;
};
it("does not retry on any 4xx response", () => {
const { retry } = setup();
for (const status of [400, 401, 403, 404, 409, 422, 429, 499]) {
expect(retry(0, axiosError(status))).toBe(false);
}
});
it("retries up to 3 times on 5xx responses", () => {
const { retry } = setup();
expect(retry(0, axiosError(500))).toBe(true);
expect(retry(2, axiosError(500))).toBe(true);
expect(retry(3, axiosError(500))).toBe(false);
});
it("retries up to 3 times on axios errors with no response", () => {
const { retry } = setup();
expect(retry(0, axiosNetworkError())).toBe(true);
expect(retry(2, axiosNetworkError())).toBe(true);
expect(retry(3, axiosNetworkError())).toBe(false);
});
it("respects options.retry === false", () => {
const { retry } = setup({ retry: false });
expect(retry).toBe(false);
});
it("respects options.retry === 0 (nullish coalescing, not falsy)", () => {
const { retry } = setup({ retry: 0 });
expect(retry).toBe(0);
});
it("respects a custom options.retry callback", () => {
const customRetry = jest.fn(() => true);
const { retry } = setup({ retry: customRetry });
expect(retry).toBe(customRetry);
});
it("respects a custom options.retryDelay", () => {
const customDelay = jest.fn(() => 42);
const { retryDelay } = setup({ retryDelay: customDelay });
expect(retryDelay).toBe(customDelay);
});
});
// ---------------------------------------------------------------------------
// retryDelay: exponential backoff capped at 30s
// ---------------------------------------------------------------------------
describe("UseRequestProcessor retryDelay", () => {
it("uses exponential backoff capped at 30s for queries", () => {
const { query } = UseRequestProcessor();
query(["k"], async () => ({}));
const { retryDelay } = mockCapturedQueryOptions;
expect(retryDelay(0)).toBe(1000);
expect(retryDelay(1)).toBe(2000);
expect(retryDelay(2)).toBe(4000);
expect(retryDelay(3)).toBe(8000);
expect(retryDelay(4)).toBe(16000);
expect(retryDelay(5)).toBe(30000);
expect(retryDelay(10)).toBe(30000);
});
it("uses exponential backoff capped at 30s for mutations", () => {
const { mutate } = UseRequestProcessor();
mutate(["k"], async () => ({}));
const { retryDelay } = mockCapturedMutationOptions;
expect(retryDelay(0)).toBe(1000);
expect(retryDelay(1)).toBe(2000);
expect(retryDelay(2)).toBe(4000);
expect(retryDelay(10)).toBe(30000);
});
});

View File

@ -6,11 +6,33 @@ import {
useQuery,
useQueryClient,
} from "@tanstack/react-query";
import axios from "axios";
import type {
MutationFunctionType,
QueryFunctionType,
} from "../../../types/api";
// 4xx responses are intentional client-side rejections (auth, validation,
// deployment guards, etc.) and won't change on retry.
function isClientError(error: unknown): boolean {
if (!axios.isAxiosError(error)) return false;
const status = error.response?.status;
return typeof status === "number" && status >= 400 && status < 500;
}
function makeRetry(maxRetries: number) {
return (failureCount: number, error: unknown) => {
if (isClientError(error)) return false;
return failureCount < maxRetries;
};
}
const queryRetry = makeRetry(5);
const mutationRetry = makeRetry(3);
const retryDelay = (attemptIndex: number) =>
Math.min(1000 * 2 ** attemptIndex, 30000);
export function UseRequestProcessor(): {
query: QueryFunctionType;
mutate: MutationFunctionType;
@ -26,8 +48,8 @@ export function UseRequestProcessor(): {
return useQuery({
queryKey,
queryFn,
retry: 5,
retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000),
retry: queryRetry,
retryDelay,
...options,
});
}
@ -45,8 +67,8 @@ export function UseRequestProcessor(): {
options.onSettled && options.onSettled(data, error, variables, context);
},
...options,
retry: options.retry ?? 3,
retryDelay: (attemptIndex) => Math.min(1000 * 2 ** attemptIndex, 30000),
retry: options.retry ?? mutationRetry,
retryDelay: options.retryDelay ?? retryDelay,
});
}

View File

@ -94,7 +94,7 @@ const InputWrapper: React.FC<InputWrapperProps> = ({
<div className="flex w-full flex-col-reverse">
<div
data-testid="input-wrapper"
className="flex w-full flex-col rounded-md border cursor-text border-input p-4 hover:border-muted-foreground focus:border-[1.75px] has-[:focus]:border-primary"
className="flex w-full flex-col rounded-md border cursor-text border-input bg-muted p-4 hover:border-muted-foreground focus:border-[1.75px] has-[:focus]:border-primary"
onClick={onClick}
onMouseDown={onMouseDown}
onKeyDown={onKeyDown}

View File

@ -34,7 +34,7 @@ const UploadFileButton = ({
/>
<Button
disabled={isBuilding}
className={`btn-playground-actions ${
className={`h-7 w-7 px-0 flex items-center justify-center ${
isBuilding
? "cursor-not-allowed"
: "text-muted-foreground hover:text-primary"
@ -42,7 +42,7 @@ const UploadFileButton = ({
onClick={handleClick}
unstyled
>
<ForwardedIconComponent name="File" />
<ForwardedIconComponent className="h-[18px] w-[18px]" name="File" />
</Button>
</div>
</ShadTooltip>

View File

@ -50,7 +50,7 @@ function ConfirmationModal({
const [flag, setFlag] = useState(false);
useEffect(() => {
if (open) setModalOpen(open);
if (open !== undefined) setModalOpen(open);
}, [open]);
useEffect(() => {
@ -81,7 +81,7 @@ function ConfirmationModal({
return (
<BaseModal
{...props}
open={open}
open={modalOpen}
setOpen={setModalOpen}
onOpenAutoFocus={onOpenAutoFocus}
>

View File

@ -218,6 +218,42 @@ describe("Create mode — canGoNext validation", () => {
expect(result.current.canGoNext).toBe(true);
});
it("blocks when trimmed name does not start with a letter", () => {
const { result } = renderCreateHook({
initialProvider: mockProvider,
initialInstance: mockInstance,
});
act(() => result.current.handleNext()); // → step 2
act(() => {
result.current.setDeploymentName(" 1 Agent");
result.current.setSelectedLlm("gpt-4");
});
expect(result.current.canGoNext).toBe(false);
expect(result.current.isDeploymentNameValid).toBe(false);
expect(result.current.hasDeploymentNameFormatError).toBe(true);
});
it("allows when trimmed name starts with a unicode letter", () => {
const { result } = renderCreateHook({
initialProvider: mockProvider,
initialInstance: mockInstance,
});
act(() => result.current.handleNext()); // → step 2
act(() => {
result.current.setDeploymentName(" Ágent");
result.current.setSelectedLlm("gpt-4");
});
expect(result.current.canGoNext).toBe(true);
expect(result.current.isDeploymentNameValid).toBe(true);
expect(result.current.hasDeploymentNameFormatError).toBe(false);
});
});
describe("Step 3 (Attach Flows)", () => {

View File

@ -180,4 +180,18 @@ describe("Custom tool naming", () => {
"Tool Beta",
);
});
it("rejects create payload when agent name does not start with a letter", () => {
const { result } = renderStepperHook();
act(() => {
result.current.setDeploymentName("1 Agent");
result.current.setSelectedLlm("test-model");
result.current.handleSelectVersion("flow-1", "ver-1", "v1");
});
expect(() => result.current.buildDeploymentPayload("provider-1")).toThrow(
"Deployment name must start with a letter",
);
});
});

View File

@ -0,0 +1,75 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import DeploymentStepperFooter from "../components/deployment-stepper-footer";
jest.mock(
"@/components/common/genericIconComponent",
() =>
function MockIcon({ name }: { name: string }) {
return <span data-testid={`icon-${name}`} />;
},
);
describe("DeploymentStepperFooter", () => {
it("shows done state controls when deployed", () => {
render(
<DeploymentStepperFooter
canGoNext={true}
currentStep={4}
isCreatingAccount={false}
isDeployed={true}
isDeploying={false}
isInDeployPhase={true}
isFinalStep={true}
minStep={1}
actionIcon="Rocket"
actionLabel="Deploy"
progressLabel="Deploying..."
onBack={jest.fn()}
onCancel={jest.fn()}
onClose={jest.fn()}
onPrimaryAction={jest.fn()}
/>,
);
expect(screen.getByRole("button", { name: "Done" })).toBeInTheDocument();
expect(
screen.queryByRole("button", { name: "Cancel" }),
).not.toBeInTheDocument();
expect(
screen.queryByRole("button", { name: "Back" }),
).not.toBeInTheDocument();
expect(
screen.queryByTestId("deployment-stepper-next"),
).not.toBeInTheDocument();
});
it("calls onClose from done button", async () => {
const user = userEvent.setup();
const onClose = jest.fn();
render(
<DeploymentStepperFooter
canGoNext={true}
currentStep={4}
isCreatingAccount={false}
isDeployed={true}
isDeploying={false}
isInDeployPhase={true}
isFinalStep={true}
minStep={1}
actionIcon="Rocket"
actionLabel="Deploy"
progressLabel="Deploying..."
onBack={jest.fn()}
onCancel={jest.fn()}
onClose={onClose}
onPrimaryAction={jest.fn()}
/>,
);
await user.click(screen.getByRole("button", { name: "Done" }));
expect(onClose).toHaveBeenCalled();
});
});

View File

@ -345,7 +345,7 @@ describe("Step labels", () => {
renderModal();
expect(screen.getByText("Provider")).toBeInTheDocument();
expect(screen.getByText("Type")).toBeInTheDocument();
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
expect(screen.getByText("Flows")).toBeInTheDocument();
expect(screen.getByText("Review")).toBeInTheDocument();
});
@ -355,7 +355,7 @@ describe("Step labels", () => {
renderModal({ editingDeployment: makeDeployment() });
expect(screen.queryByText("Provider")).not.toBeInTheDocument();
expect(screen.getByText("Type")).toBeInTheDocument();
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
expect(screen.getByText("Flows")).toBeInTheDocument();
expect(screen.getByText("Review")).toBeInTheDocument();
});
});

View File

@ -0,0 +1,86 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import DeploymentSuccessContent from "../components/deployment-success-content";
jest.mock(
"@/components/common/genericIconComponent",
() =>
function MockIcon({ name }: { name: string }) {
return <span data-testid={`icon-${name}`} />;
},
);
describe("DeploymentSuccessContent", () => {
it("renders success copy and provider link", () => {
render(
<DeploymentSuccessContent
deploymentName="My Agent"
providerName="watsonx Orchestrate"
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
showTestButton={true}
onTest={jest.fn()}
/>,
);
expect(screen.getByText("Deployment successful")).toBeInTheDocument();
expect(
screen.getByText("Deployed to watsonx Orchestrate as draft"),
).toBeInTheDocument();
expect(
screen.getByRole("link", { name: /watsonx Orchestrate/i }),
).toHaveAttribute(
"href",
"https://www.ibm.com/products/watsonx-orchestrate",
);
});
it("shows test button and calls handler", async () => {
const user = userEvent.setup();
const onTest = jest.fn();
render(
<DeploymentSuccessContent
deploymentName="My Agent"
providerName="watsonx Orchestrate"
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
showTestButton={true}
onTest={onTest}
/>,
);
await user.click(screen.getByRole("button", { name: "Test Deployment" }));
expect(onTest).toHaveBeenCalled();
});
it("hides test button without deployment name", () => {
render(
<DeploymentSuccessContent
providerName="watsonx Orchestrate"
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
showTestButton={true}
onTest={jest.fn()}
/>,
);
expect(
screen.queryByRole("button", { name: "Test Deployment" }),
).not.toBeInTheDocument();
});
it("hides test button when showTestButton is false", () => {
render(
<DeploymentSuccessContent
deploymentName="My Agent"
providerName="watsonx Orchestrate"
providerUrl="https://www.ibm.com/products/watsonx-orchestrate"
showTestButton={false}
onTest={jest.fn()}
/>,
);
expect(
screen.queryByRole("button", { name: "Test Deployment" }),
).not.toBeInTheDocument();
});
});

View File

@ -92,6 +92,31 @@ describe("Edit mode — basic state", () => {
expect(result.current.canGoNext).toBe(true);
});
it("blocks update flow when existing name does not start with a letter", () => {
const invalidDeployment = { ...mockDeployment, name: "1 Agent" };
const wrapper = ({ children }: { children: React.ReactNode }) => (
<DeploymentStepperProvider
initialState={{
editingDeployment: invalidDeployment,
selectedVersionByFlow: initialVersions,
initialLlm: "test-model",
initialToolNameByFlow: initialToolNames,
initialConnectionsByFlow: initialConnections,
}}
>
{children}
</DeploymentStepperProvider>
);
const { result } = renderHook(() => useDeploymentStepper(), { wrapper });
expect(result.current.canGoNext).toBe(false);
expect(result.current.isDeploymentNameValid).toBe(false);
expect(result.current.hasDeploymentNameFormatError).toBe(true);
expect(() => result.current.buildDeploymentUpdatePayload()).toThrow(
"Deployment name must start with a letter",
);
});
it("canGoNext on step 2 (Attach) allows proceeding in edit mode", () => {
const { result } = renderEditHook();
act(() => result.current.handleNext()); // step 2

View File

@ -230,14 +230,14 @@ beforeEach(() => {
// ---------------------------------------------------------------------------
describe("Three-panel layout rendering", () => {
it("renders the Attach Flows heading", () => {
it("renders the Flows heading", () => {
render(<StepAttachFlows />);
expect(screen.getByText("Attach Flows")).toBeInTheDocument();
expect(screen.getByText("Flows")).toBeInTheDocument();
});
it("renders the Available Flows panel header", () => {
it("renders the Available panel header", () => {
render(<StepAttachFlows />);
expect(screen.getByText("Available Flows")).toBeInTheDocument();
expect(screen.getByText("Available")).toBeInTheDocument();
});
it("renders flow items in the list", () => {

View File

@ -13,6 +13,8 @@ const mockSetSelectedLlm = jest.fn();
let mockIsEditMode = false;
let mockDeploymentType = "agent";
let mockDeploymentName = "";
let mockIsDeploymentNameValid = false;
let mockHasDeploymentNameFormatError = false;
let mockDeploymentDescription = "";
let mockSelectedLlm = "";
let mockSelectedInstance: { id: string } | null = { id: "inst-1" };
@ -26,6 +28,8 @@ jest.mock("../contexts/deployment-stepper-context", () => ({
setDeploymentType: mockSetDeploymentType,
deploymentName: mockDeploymentName,
setDeploymentName: mockSetDeploymentName,
isDeploymentNameValid: mockIsDeploymentNameValid,
hasDeploymentNameFormatError: mockHasDeploymentNameFormatError,
deploymentDescription: mockDeploymentDescription,
setDeploymentDescription: mockSetDeploymentDescription,
selectedLlm: mockSelectedLlm,
@ -61,6 +65,8 @@ beforeEach(() => {
mockIsEditMode = false;
mockDeploymentType = "agent";
mockDeploymentName = "";
mockIsDeploymentNameValid = false;
mockHasDeploymentNameFormatError = false;
mockDeploymentDescription = "";
mockSelectedLlm = "";
mockSelectedInstance = { id: "inst-1" };
@ -143,6 +149,30 @@ describe("Name input", () => {
screen.queryByText("Name cannot be changed after creation."),
).not.toBeInTheDocument();
});
it("shows validation error when name does not start with a letter", () => {
mockDeploymentName = "1 Agent";
mockHasDeploymentNameFormatError = true;
render(<StepType />);
expect(
screen.getByText("Agent name must start with a letter."),
).toBeInTheDocument();
expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute(
"aria-invalid",
"true",
);
});
it("does not show validation error for empty name", () => {
render(<StepType />);
expect(
screen.queryByText("Agent name must start with a letter."),
).not.toBeInTheDocument();
expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute(
"aria-invalid",
"false",
);
});
});
// ---------------------------------------------------------------------------

View File

@ -0,0 +1,89 @@
import ForwardedIconComponent from "@/components/common/genericIconComponent";
import { Button } from "@/components/ui/button";
interface DeploymentStepperFooterProps {
canGoNext: boolean;
currentStep: number;
isCreatingAccount: boolean;
isDeployed: boolean;
isDeploying: boolean;
isInDeployPhase: boolean;
isFinalStep: boolean;
minStep: number;
actionIcon: string;
actionLabel: string;
progressLabel: string;
onBack: () => void;
onCancel: () => void;
onClose: () => void;
onPrimaryAction: () => void;
}
export default function DeploymentStepperFooter({
canGoNext,
currentStep,
isCreatingAccount,
isDeployed,
isDeploying,
isInDeployPhase,
isFinalStep,
minStep,
actionIcon,
actionLabel,
progressLabel,
onBack,
onCancel,
onClose,
onPrimaryAction,
}: DeploymentStepperFooterProps) {
return (
<div className="flex items-center justify-between border-t border-border px-6 py-4">
{isDeployed ? (
<div />
) : (
<Button variant="ghost" onClick={onCancel}>
Cancel
</Button>
)}
<div className="flex items-center gap-3">
{!isDeployed && (
<Button
variant="outline"
onClick={onBack}
disabled={currentStep === minStep || isDeploying}
>
Back
</Button>
)}
{!isInDeployPhase && (
<Button
onClick={onPrimaryAction}
disabled={!canGoNext || isCreatingAccount}
data-testid="deployment-stepper-next"
>
{isFinalStep ? (
<>
<ForwardedIconComponent name={actionIcon} className="h-4 w-4" />
{actionLabel}
</>
) : isCreatingAccount ? (
"Connecting..."
) : (
"Next"
)}
</Button>
)}
{isDeploying && (
<Button disabled data-testid="deployment-stepper-next">
<ForwardedIconComponent
name={actionIcon}
className="h-4 w-4 animate-pulse"
/>
{progressLabel}
</Button>
)}
{isDeployed && <Button onClick={onClose}>Done</Button>}
</div>
</div>
);
}

View File

@ -1,7 +1,5 @@
import { useMemo, useState } from "react";
import { useParams } from "react-router-dom";
import ForwardedIconComponent from "@/components/common/genericIconComponent";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
@ -21,7 +19,9 @@ import {
} from "../contexts/deployment-stepper-context";
import { useErrorAlert } from "../hooks/use-error-alert";
import type { Deployment, DeploymentProvider, ProviderAccount } from "../types";
import DeploymentStepper from "./deployment-stepper";
import DeploymentStepper, { CREATE_DEPLOYED_STEPS } from "./deployment-stepper";
import DeploymentStepperFooter from "./deployment-stepper-footer";
import DeploymentSuccessContent from "./deployment-success-content";
import StepAttachFlows from "./step-attach-flows";
import StepDeployStatus from "./step-deploy-status";
import StepProvider from "./step-provider";
@ -206,6 +206,7 @@ function DeploymentStepperModalContent({
canGoNext,
handleNext,
handleBack,
selectedProvider,
selectedInstance,
setSelectedInstance,
needsProviderAccountCreation,
@ -225,6 +226,8 @@ function DeploymentStepperModalContent({
const isDeploying = deploymentPhase === "deploying";
const isDeployed = deploymentPhase === "deployed";
const isInDeployPhase = isDeploying || isDeployed;
const providerConsoleUrl = "https://www.ibm.com/products/watsonx-orchestrate";
const providerDisplayName = selectedProvider?.name ?? "watsonx Orchestrate";
// In edit mode, steps are shifted: 1=Type, 2=Attach, 3=Review.
const logicalStep = isEditMode ? currentStep + 1 : currentStep;
@ -316,19 +319,40 @@ function DeploymentStepperModalContent({
className="text-center text-2xl font-semibold"
data-testid="stepper-modal-title"
>
{isEditMode ? "Update Deployment" : "Create New Deployment"}
{isDeployed && !isEditMode
? "Deployed"
: isEditMode
? "Update Deployment"
: "Create New Deployment"}
</h2>
<DeploymentStepper />
<DeploymentStepper
steps={!isEditMode && isDeployed ? CREATE_DEPLOYED_STEPS : undefined}
currentStepOverride={!isEditMode && isDeployed ? 4 : undefined}
/>
</div>
{/* Content box: step content + footer */}
<div className="mx-4 mb-4 mt-4 flex flex-1 flex-col overflow-hidden rounded-lg border border-border bg-background">
<div className="flex min-h-0 flex-1 flex-col overflow-y-auto px-6 py-2">
{isInDeployPhase ? (
<StepDeployStatus
phase={isDeploying ? "deploying" : "deployed"}
deploymentName={createdDeployment?.name}
/>
isDeployed && !isEditMode ? (
<DeploymentSuccessContent
deploymentName={createdDeployment?.name}
providerName={providerDisplayName}
providerUrl={providerConsoleUrl}
showTestButton={
!!onTestDeployment &&
!!createdDeployment &&
!!selectedInstance?.id
}
onTest={handleTest}
/>
) : (
<StepDeployStatus
phase={isDeploying ? "deploying" : "deployed"}
deploymentName={createdDeployment?.name}
/>
)
) : (
<>
{logicalStep === 1 && <StepProvider />}
@ -339,61 +363,23 @@ function DeploymentStepperModalContent({
)}
</div>
{/* Footer */}
<div className="flex items-center justify-between border-t border-border px-6 py-4">
<Button variant="ghost" onClick={() => setOpen(false)}>
{isDeployed ? "Close" : "Cancel"}
</Button>
<div className="flex items-center gap-3">
{!isDeployed && (
<Button
variant="outline"
onClick={handleBack}
disabled={currentStep === minStep || isDeploying}
>
Back
</Button>
)}
{!isInDeployPhase && (
<Button
onClick={isFinalStep ? handleDeploy : handleStepNext}
disabled={!canGoNext || isCreatingAccount}
data-testid="deployment-stepper-next"
>
{isFinalStep ? (
<>
<ForwardedIconComponent
name={actionIcon}
className="h-4 w-4"
/>
{actionLabel}
</>
) : isCreatingAccount ? (
"Connecting..."
) : (
"Next"
)}
</Button>
)}
{isDeploying && (
<Button disabled data-testid="deployment-stepper-next">
<ForwardedIconComponent
name={actionIcon}
className="h-4 w-4 animate-pulse"
/>
{progressLabel}
</Button>
)}
{isDeployed && onTestDeployment && !isEditMode && (
<Button
data-testid="deployment-stepper-test"
onClick={handleTest}
>
Test
</Button>
)}
</div>
</div>
<DeploymentStepperFooter
canGoNext={canGoNext}
currentStep={currentStep}
isCreatingAccount={isCreatingAccount}
isDeployed={isDeployed}
isDeploying={isDeploying}
isInDeployPhase={isInDeployPhase}
isFinalStep={isFinalStep}
minStep={minStep}
actionIcon={actionIcon}
actionLabel={actionLabel}
progressLabel={progressLabel}
onBack={handleBack}
onCancel={() => setOpen(false)}
onClose={() => setOpen(false)}
onPrimaryAction={isFinalStep ? handleDeploy : handleStepNext}
/>
</div>
</>
);

View File

@ -1,25 +1,41 @@
import { cn } from "@/utils/utils";
import { useDeploymentStepper } from "../contexts/deployment-stepper-context";
const CREATE_STEPS = [
export const CREATE_STEPS = [
{ number: 1, label: "Provider" },
{ number: 2, label: "Type" },
{ number: 3, label: "Attach Flows" },
{ number: 3, label: "Flows" },
{ number: 4, label: "Review" },
] as const;
export const CREATE_DEPLOYED_STEPS = [
{ number: 1, label: "Provider" },
{ number: 2, label: "Type" },
{ number: 3, label: "Flows" },
{ number: 4, label: "Deployed" },
] as const;
const EDIT_STEPS = [
{ number: 1, label: "Type" },
{ number: 2, label: "Attach Flows" },
{ number: 2, label: "Flows" },
{ number: 3, label: "Review" },
] as const;
export const DEPLOYMENT_STEPS = CREATE_STEPS;
export default function DeploymentStepper() {
interface DeploymentStepperProps {
steps?: readonly { number: number; label: string }[];
currentStepOverride?: number;
}
export default function DeploymentStepper({
steps: stepsProp,
currentStepOverride,
}: DeploymentStepperProps) {
const { currentStep, isEditMode } = useDeploymentStepper();
const steps = isEditMode ? EDIT_STEPS : CREATE_STEPS;
const progressPercent = ((currentStep - 1) / (steps.length - 1)) * 100;
const steps = stepsProp ?? (isEditMode ? EDIT_STEPS : CREATE_STEPS);
const activeStep = currentStepOverride ?? currentStep;
const progressPercent = ((activeStep - 1) / (steps.length - 1)) * 100;
return (
<div className="relative mx-auto h-[52px] w-full max-w-[700px]">
@ -35,7 +51,7 @@ export default function DeploymentStepper() {
<div
className={cn(
"flex h-8 w-8 items-center justify-center rounded-full text-sm font-medium transition-colors",
currentStep >= step.number
activeStep >= step.number
? "bg-foreground text-background"
: "bg-muted text-muted-foreground",
)}
@ -45,7 +61,7 @@ export default function DeploymentStepper() {
<span
className={cn(
"whitespace-nowrap text-xs text-foreground",
currentStep >= step.number && "font-medium",
activeStep >= step.number && "font-medium",
)}
>
{step.label}

View File

@ -0,0 +1,63 @@
import ForwardedIconComponent from "@/components/common/genericIconComponent";
import { Button } from "@/components/ui/button";
interface DeploymentSuccessContentProps {
deploymentName?: string;
providerName: string;
providerUrl: string;
showTestButton: boolean;
onTest: () => void;
}
export default function DeploymentSuccessContent({
deploymentName,
providerName,
providerUrl,
showTestButton,
onTest,
}: DeploymentSuccessContentProps) {
return (
<div className="flex flex-1 flex-col items-center justify-center gap-8 py-10 text-center">
<div className="flex h-16 w-16 items-center justify-center rounded-full border-2 border-accent-emerald">
<ForwardedIconComponent
name="Check"
className="h-8 w-8 text-accent-emerald-foreground"
/>
</div>
<div className="space-y-2">
<h3 className="font-sans text-2xl font-semibold tracking-normal text-foreground">
Deployment successful
</h3>
<p className="font-sans text-base font-normal text-muted-foreground">
Deployed to {providerName} as draft
</p>
</div>
<div className="flex items-center gap-1 font-sans text-sm font-normal text-muted-foreground">
<span>Publish from Draft to Live in</span>
<a
href={providerUrl}
target="_blank"
rel="noopener noreferrer"
className="inline-flex items-center gap-1 font-semibold text-foreground hover:underline"
>
{providerName}
<ForwardedIconComponent name="ArrowUpRight" className="h-4 w-4" />
</a>
</div>
{showTestButton && deploymentName && (
<Button
variant="outline"
className="h-11 w-full max-w-lg gap-2 rounded-lg border-input font-sans text-base font-semibold text-foreground hover:bg-input"
data-testid="deployment-stepper-test"
onClick={onTest}
>
<ForwardedIconComponent name="Play" className="h-5 w-5" />
Test Deployment
</Button>
)}
</div>
);
}

View File

@ -29,7 +29,7 @@ export const FlowListPanel = memo(function FlowListPanel({
return (
<div className="flex w-[280px] flex-shrink-0 flex-col border-r border-border">
<div className="border-b border-border p-4 text-sm text-muted-foreground">
Available Flows
Available
</div>
<div className="flex-1 space-y-1 overflow-y-auto p-2">
{flows.map((flow) => {

View File

@ -259,7 +259,7 @@ export default function StepAttachFlows() {
return (
<div className="flex min-h-0 flex-1 flex-col gap-4 py-3">
<h2 className="text-lg font-semibold">Attach Flows</h2>
<h2 className="text-lg font-semibold">Flows</h2>
<div className="flex min-h-0 flex-1 overflow-hidden rounded-xl border border-border">
<FlowListPanel

View File

@ -32,6 +32,7 @@ export default function StepType() {
setDeploymentType,
deploymentName,
setDeploymentName,
hasDeploymentNameFormatError,
deploymentDescription,
setDeploymentDescription,
selectedLlm,
@ -142,11 +143,21 @@ export default function StepType() {
</span>
<Input
placeholder="e.g., Sales Bot"
className="bg-muted"
className={cn(
"bg-muted",
hasDeploymentNameFormatError &&
"border-destructive focus-visible:ring-0",
)}
value={deploymentName}
onChange={(e) => setDeploymentName(e.target.value)}
disabled={isEditMode}
aria-invalid={hasDeploymentNameFormatError}
/>
{hasDeploymentNameFormatError && (
<span className="mt-1 text-xs text-destructive">
Agent name must start with a letter.
</span>
)}
{isEditMode && (
<span className="mt-1 text-xs text-muted-foreground">
Name cannot be changed after creation.

View File

@ -73,6 +73,8 @@ interface DeploymentStepperContextType {
setDeploymentType: (type: DeploymentType) => void;
deploymentName: string;
setDeploymentName: (name: string) => void;
isDeploymentNameValid: boolean;
hasDeploymentNameFormatError: boolean;
deploymentDescription: string;
setDeploymentDescription: (description: string) => void;
selectedLlm: string;
@ -168,6 +170,11 @@ export function DeploymentStepperProvider({
>(initialState?.initialConnectionsByFlow ?? new Map());
const [hasToolNameErrors, setHasToolNameErrors] = useState(false);
const trimmedDeploymentName = deploymentName.trim();
const hasDeploymentNameFormatError =
trimmedDeploymentName !== "" && !/^\p{L}/u.test(trimmedDeploymentName);
const isDeploymentNameValid =
trimmedDeploymentName !== "" && !hasDeploymentNameFormatError;
// Edit mode: track which pre-existing flows the user wants to detach.
const [removedFlowIds, setRemovedFlowIds] = useState<Set<string>>(new Set());
@ -253,7 +260,7 @@ export function DeploymentStepperProvider({
);
}
if (logical === 2) {
return deploymentName.trim() !== "" && selectedLlm.trim() !== "";
return isDeploymentNameValid && selectedLlm.trim() !== "";
}
if (logical === 3) {
// In edit mode, user can proceed without new attachments (may just change desc/LLM).
@ -270,6 +277,7 @@ export function DeploymentStepperProvider({
selectedInstance,
hasValidCredentials,
deploymentName,
isDeploymentNameValid,
selectedLlm,
selectedVersionByFlow,
isEditMode,
@ -357,6 +365,9 @@ export function DeploymentStepperProvider({
const buildDeploymentPayload = useCallback(
(providerId: string): DeploymentCreateRequest => {
if (!isDeploymentNameValid) {
throw new Error("Deployment name must start with a letter");
}
const allConnectionIds = new Set<string>();
Array.from(attachedConnectionByFlow.values()).forEach((ids) => {
ids.forEach((id) => allConnectionIds.add(id));
@ -381,7 +392,7 @@ export function DeploymentStepperProvider({
...(initialState?.projectId
? { project_id: initialState.projectId }
: {}),
name: deploymentName,
name: trimmedDeploymentName,
description: deploymentDescription,
type: deploymentType,
provider_data: {
@ -396,10 +407,11 @@ export function DeploymentStepperProvider({
buildConnectionPayloads,
initialState?.projectId,
deploymentDescription,
deploymentName,
deploymentType,
isDeploymentNameValid,
selectedLlm,
selectedVersionByFlow,
trimmedDeploymentName,
toolNameByFlow,
],
);
@ -411,6 +423,9 @@ export function DeploymentStepperProvider({
"buildDeploymentUpdatePayload called outside edit mode",
);
}
if (!isDeploymentNameValid) {
throw new Error("Deployment name must start with a letter");
}
const result: DeploymentUpdateRequest = {
deployment_id: editingDeployment.id,
@ -510,6 +525,7 @@ export function DeploymentStepperProvider({
}, [
editingDeployment,
deploymentDescription,
isDeploymentNameValid,
selectedLlm,
initialVersionByFlow,
initialToolNameByFlow,
@ -541,6 +557,8 @@ export function DeploymentStepperProvider({
setDeploymentType,
deploymentName,
setDeploymentName,
isDeploymentNameValid,
hasDeploymentNameFormatError,
deploymentDescription,
setDeploymentDescription,
selectedLlm,
@ -581,6 +599,8 @@ export function DeploymentStepperProvider({
credentials,
deploymentType,
deploymentName,
isDeploymentNameValid,
hasDeploymentNameFormatError,
deploymentDescription,
selectedLlm,
connections,

View File

@ -150,12 +150,13 @@ async function selectProvider(page: Page) {
async function goToStepType(page: Page) {
await selectProvider(page);
await page.getByTestId("deployment-stepper-next").click();
// Wait for the Type step heading
await page.waitForSelector("text=Deployment Type");
await expect(
page.getByRole("heading", { name: /Deployment Type/i }),
).toBeVisible();
}
// ---------------------------------------------------------------------------
// Helper: navigate steps 1 → 2 → 3 (provider → type → attach flows)
// Helper: navigate steps 1 → 2 → 3 (provider → type → flows)
// ---------------------------------------------------------------------------
async function goToStepAttachFlows(page: Page) {
await goToStepType(page);
@ -165,13 +166,13 @@ async function goToStepAttachFlows(page: Page) {
// Select the LLM model
await page.getByRole("combobox").click();
await page.getByRole("option", { name: "ibm/granite-13b-chat" }).click();
// Advance to attach flows
// Advance to flows
await page.getByTestId("deployment-stepper-next").click();
await page.waitForSelector("text=Attach Flows");
await expect(page.getByRole("heading", { name: /^Flows$/i })).toBeVisible();
}
// ---------------------------------------------------------------------------
// Helper: navigate all steps through attach flows and select a flow+version
// Helper: navigate all steps through flows and select a flow+version
// ---------------------------------------------------------------------------
async function goToStepReview(page: Page) {
await goToStepAttachFlows(page);
@ -189,7 +190,9 @@ async function goToStepReview(page: Page) {
}
// Advance to review
await page.getByTestId("deployment-stepper-next").click();
await page.waitForSelector("text=Review & Confirm");
await expect(
page.getByRole("heading", { name: /Review & Confirm/i }),
).toBeVisible();
}
// ---------------------------------------------------------------------------
@ -197,7 +200,9 @@ async function goToStepReview(page: Page) {
// ---------------------------------------------------------------------------
test(
"deployment-create: opens stepper on New Deployment click",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -220,7 +225,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: step 1 provider - Next disabled without selection, enabled after selecting",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -245,7 +252,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: step 2 type - fill name and select type to enable Next",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -280,11 +289,13 @@ test(
);
// ---------------------------------------------------------------------------
// Test 4: Step 3 (Attach Flows) — select a flow and version, Next enables
// Test 4: Step 3 (Flows) — select a flow and version, Next enables
// ---------------------------------------------------------------------------
test(
"deployment-create: step 3 attach flows - select flow and version enables Next",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -331,7 +342,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: step 4 review - shows review content and Deploy button text",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -361,7 +374,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: clicking Deploy triggers POST and shows deploy status",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -414,7 +429,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: user can change tool name on review step",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -451,7 +468,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: review step shows error when tool name already exists in provider",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -476,7 +495,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: review step shows no error when tool name is unique",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -501,7 +522,9 @@ test(
// ---------------------------------------------------------------------------
test(
"deployment-create: editing tool name to unique value clears duplicate error",
{ tag: ["@deployment", "@workspace"] },
{
tag: ["@deployment", "@workspace"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",

View File

@ -83,9 +83,23 @@ async function openEditDialog(page: Parameters<typeof test>[2]["page"]) {
await page.waitForSelector('[data-testid="stepper-modal-title"]');
}
async function expectDeploymentTypeStep(
page: Parameters<typeof test>[2]["page"],
) {
await expect(
page.getByRole("heading", { name: /Deployment Type/i }),
).toBeVisible();
}
async function expectFlowsStep(page: Parameters<typeof test>[2]["page"]) {
await expect(page.getByRole("heading", { name: /^Flows$/i })).toBeVisible();
}
test(
"Opens edit stepper from actions menu",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -110,7 +124,9 @@ test(
test(
"Edit mode skips provider step — starts at Type",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -126,7 +142,7 @@ test(
await openEditDialog(page);
// Wait for stepper body to render (parallel fetches must complete)
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
await expect(page.getByText("Deployment Type")).toBeVisible();
@ -141,7 +157,9 @@ test(
test(
"Name field pre-populated in edit mode",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -157,7 +175,7 @@ test(
await openEditDialog(page);
// Wait for stepper body to render
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
const nameInput = page.getByPlaceholder("e.g., Sales Bot");
await expect(nameInput).toHaveValue("Test Deployment");
@ -166,7 +184,9 @@ test(
test(
"Submitting PATCH closes modal",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -182,13 +202,13 @@ test(
await openEditDialog(page);
// Wait for stepper body to render (parallel fetches must complete)
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
// Navigate through the stepper steps to reach Review
// Step: Type → click Next
await page.getByTestId("deployment-stepper-next").click();
// Step: Attach Flows → click Next
// Step: Flows → click Next
await page.getByTestId("deployment-stepper-next").click();
// Step: Review → click Update (final step)
@ -209,7 +229,9 @@ test(
test(
"Cancel during edit closes modal without calling PATCH",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -225,7 +247,7 @@ test(
await openEditDialog(page);
// Wait for stepper body to render
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
let patchCalled = false;
page.on("request", (req) => {
@ -356,7 +378,9 @@ async function setupRoutesWithConnections(
// ---------------------------------------------------------------------------
test(
"Edit mode includes new connections in PATCH request",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -399,11 +423,11 @@ test(
await page.waitForSelector('[data-testid="stepper-modal-title"]');
// Step 1 (Type) → Next
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
await page.getByTestId("deployment-stepper-next").click();
// Step 2 (Attach Flows) — flow "f1" should already be attached
await page.waitForSelector("text=Attach Flows");
// Step 2 (Flows) — flow "f1" should already be attached
await expectFlowsStep(page);
await page.waitForSelector('[data-testid="flow-item-f1"]');
// Click the pre-attached flow and its version to open the connection panel
@ -455,7 +479,9 @@ test(
// ---------------------------------------------------------------------------
test(
"Edit mode includes removed connections in PATCH request",
{ tag: ["@release", "@workspace", "@api"] },
{
tag: ["@release", "@workspace", "@api"],
},
async ({ page }) => {
test.skip(
process.env.LANGFLOW_FEATURE_WXO_DEPLOYMENTS !== "true",
@ -497,11 +523,11 @@ test(
await page.waitForSelector('[data-testid="stepper-modal-title"]');
// Step 1 (Type) → Next
await page.waitForSelector('h2:has-text("Deployment Type")');
await expectDeploymentTypeStep(page);
await page.getByTestId("deployment-stepper-next").click();
// Step 2 (Attach Flows)
await page.waitForSelector("text=Attach Flows");
// Step 2 (Flows)
await expectFlowsStep(page);
await page.waitForSelector('[data-testid="flow-item-f1"]');
await page.getByTestId("flow-item-f1").click();
await page.waitForSelector('[data-testid="version-item-fv1"]');

View File

@ -0,0 +1,103 @@
import { expect, test } from "../../fixtures";
import { adjustScreenView } from "../../utils/adjust-screen-view";
import { awaitBootstrapTest } from "../../utils/await-bootstrap-test";
import {
closeAdvancedOptions,
openAdvancedOptions,
} from "../../utils/open-advanced-options";
/**
* Covers the user-facing contract of the "Default Agent Tools" feature:
* see CZL/MANUAL_TEST_DEFAULT_AGENT_TOOLS.md scenarios S1, S3, S5.
*
* Backing unit tests (pytest) already cover runtime behaviour; these tests
* validate the UI wiring so a regression in the inputs list or default
* prompt value is caught in the release gate.
*/
async function dragAgentOntoCanvas(page: import("@playwright/test").Page) {
await awaitBootstrapTest(page);
await page.waitForSelector('[data-testid="blank-flow"]', { timeout: 30000 });
await page.getByTestId("blank-flow").click();
await page.waitForSelector('[data-testid="sidebar-search-input"]', {
timeout: 30000,
});
await page.getByTestId("sidebar-search-input").click();
await page.getByTestId("sidebar-search-input").fill("agent");
await page.waitForSelector('[data-testid="models_and_agentsAgent"]', {
timeout: 30000,
});
await page
.getByTestId("models_and_agentsAgent")
.dragTo(page.locator('//*[@id="react-flow-id"]'));
await adjustScreenView(page);
}
test(
"Agent ships with Calculator and Current Date toggles enabled by default (S1/S3)",
{ tag: ["@release", "@workspace", "@components"] },
async ({ page }) => {
await dragAgentOntoCanvas(page);
// Focus the Agent node so its advanced-field drawer is reachable.
await page.getByTestId("div-generic-node").click();
await openAdvancedOptions(page);
// Both advanced toggles exist as show-on-canvas checkboxes.
// Their default `value=True` is validated by the pytest suite
// (`test_should_have_placeholders_in_default_system_prompt` covers the
// default contract of the inputs list).
await expect(
page.locator('//*[@id="showadd_current_date_tool"]'),
).toBeVisible({ timeout: 10000 });
await expect(
page.locator('//*[@id="showadd_calculator_tool"]'),
).toBeVisible({ timeout: 10000 });
// Flip the Calculator field visible on canvas so we can assert the toggle
// is active and can be switched off and on (S3).
await page.locator('//*[@id="showadd_calculator_tool"]').click();
await closeAdvancedOptions(page);
await adjustScreenView(page);
const calculatorToggle = page.getByTestId(
"toggle_bool_add_calculator_tool",
);
await expect(calculatorToggle).toBeVisible({ timeout: 10000 });
expect(await calculatorToggle.isChecked()).toBeTruthy();
// S3 — user disables the toggle.
await calculatorToggle.click();
expect(await calculatorToggle.isChecked()).toBeFalsy();
// Re-enable to confirm the control is bi-directional.
await calculatorToggle.click();
expect(await calculatorToggle.isChecked()).toBeTruthy();
},
);
test(
"Agent default system prompt contains {current_date} and {model_name} placeholders (S5)",
{ tag: ["@release", "@workspace", "@components"] },
async ({ page }) => {
await dragAgentOntoCanvas(page);
// The placeholders must be present inside the Agent Instructions textarea
// so the dynamic injection has a discoverable effect out of the box.
const instructionsTextarea = page.getByTestId("textarea_str_system_prompt");
await expect(instructionsTextarea).toBeVisible({ timeout: 10000 });
// <textarea> stores content in its `value`, not textContent.
const promptText = await instructionsTextarea.inputValue();
expect(promptText).toContain("{current_date}");
expect(promptText).toContain("{model_name}");
},
);

View File

@ -41,13 +41,23 @@ test(
const cleanCode = await extractAndCleanCode(page);
// Replace the multiline string in the code
// Sanity-check: the original Ollama source has many lines. If we read it
// back as a single concatenated line, the textarea-value path is the
// problem and Ace surgery downstream cannot save us.
const cleanCodeNewlines = (cleanCode.match(/\n/g) || []).length;
if (cleanCodeNewlines < 50) {
throw new Error(
`extractAndCleanCode returned code with only ${cleanCodeNewlines} newlines (length ${cleanCode.length}); expected the multi-line Ollama source. The hidden #codeValue textarea may be returning a stripped value on this platform.`,
);
}
// Replace the multiline string in the code.
// Use a regex so the match is resilient to line-ending differences
// (LF on macOS/Linux vs CRLF on Windows after git checkout).
const originalSliderBlockRegex =
/name="temperature",\s+display_name="Temperature",\s+value=0\.1,\s+range_spec=RangeSpec\(min=0, max=1, step=0\.01\),\s+advanced=True,/;
const newCode = cleanCode.replace(
`name="temperature",
display_name="Temperature",
value=0.1,
range_spec=RangeSpec(min=0, max=1, step=0.01),
advanced=True,`,
originalSliderBlockRegex,
`name="temperature",
display_name="Temperature",
value=0.2,
@ -63,9 +73,7 @@ test(
);
// make sure codes are different
expect(cleanCode).not.toEqual(newCode);
await page.locator("textarea").last().press(`ControlOrMeta+a`);
await page.keyboard.press("Backspace");
await page.locator("textarea").last().fill(newCode);
await setAceEditorValue(page, newCode);
await page.locator('//*[@id="checkAndSaveBtn"]').click();
await adjustScreenView(page);
@ -104,13 +112,24 @@ test(
);
async function extractAndCleanCode(page: Page): Promise<string> {
// The hidden `#codeValue` mirror is rendered with `<Input value={code} />`,
// i.e. a single-line `<input>` rather than a textarea. Browsers strip
// newlines from `.value` of single-line inputs by HTML spec, so reading
// `(el as HTMLInputElement).value` returns the entire source concatenated
// onto one line. The backend then rejects it as
// `invalid decimal literal (<unknown>, line 1)`.
//
// Read the rendered HTML attribute instead — React serializes the prop as
// `value="..."` with newlines encoded (`&#10;`, `&#13;`, or literal LF
// inside attribute text), and the regex below recovers them faithfully.
// This is the same strategy queryInputComponent.spec.ts uses.
const outerHTML = await page
.locator('//*[@id="codeValue"]')
.evaluate((el) => el.outerHTML);
const valueMatch = outerHTML.match(/value="([\s\S]*?)"/);
if (!valueMatch) {
throw new Error("Could not find value attribute in the HTML");
throw new Error("Could not find value attribute in the #codeValue HTML");
}
const codeContent = valueMatch[1]
@ -119,9 +138,45 @@ async function extractAndCleanCode(page: Page): Promise<string> {
.replace(/&lt;/g, "<")
.replace(/&gt;/g, ">")
.replace(/&#x27;/g, "'")
.replace(/&#x2F;/g, "/");
.replace(/&#x2F;/g, "/")
.replace(/&#10;/g, "\n")
.replace(/&#13;/g, "\r")
.replace(/&#xA;/gi, "\n")
.replace(/&#xD;/gi, "\r");
return codeContent;
// Normalize line endings so downstream string operations are OS-agnostic
// (Windows git checkouts of .py files can end up with CRLF).
return codeContent.replace(/\r\n/g, "\n").replace(/\r/g, "\n");
}
// Set the Ace editor's content using the exact same pattern that
// queryInputComponent.spec.ts uses successfully on Windows CI:
// `textarea.last().fill(newCode)` against Ace's hidden text-input textarea.
// Textareas (unlike single-line `<input>`s) preserve `\n` in `.value`, so
// fill() reliably round-trips the multi-line source. Ace's text-input
// listener picks up the resulting `input` event, applies it to the buffer,
// and fires the `change` event that react-ace listens to — which is what
// updates the React `code` state that gets POSTed on save.
async function setAceEditorValue(page: Page, newCode: string): Promise<void> {
const expectedNewlines = (newCode.match(/\n/g) || []).length;
if (expectedNewlines < 10) {
throw new Error(
`setAceEditorValue: newCode has only ${expectedNewlines} newlines (length ${newCode.length}); upstream extractAndCleanCode likely lost newlines.`,
);
}
await page.locator("textarea").last().press("ControlOrMeta+a");
await page.keyboard.press("Backspace");
await page.locator("textarea").last().fill(newCode);
// Wait for the change to propagate into the controlled React state. The
// `#codeValue` mirror is rendered by `<Input>` (single-line), so browsers
// strip newlines from its `.value` — matching the substring is enough to
// confirm the new code landed.
await expect(page.locator("#codeValue")).toHaveValue(
/range_spec=RangeSpec\(min=3, max=30, step=1\)/,
{ timeout: 10000 },
);
}
async function mutualValidation(page: Page) {

File diff suppressed because one or more lines are too long

View File

@ -1,7 +1,6 @@
from copy import deepcopy
from typing import TYPE_CHECKING
from chromadb.config import Settings
from langchain_chroma import Chroma
from typing_extensions import override
@ -92,23 +91,22 @@ class ChromaVectorStoreComponent(LCVectorStoreComponent):
def build_vector_store(self) -> Chroma:
"""Builds the Chroma object."""
try:
from chromadb import Client
from langchain_chroma import Chroma
except ImportError as e:
msg = "Could not import Chroma integration package. Please install it with `pip install langchain-chroma`."
raise ImportError(msg) from e
# Chroma settings
chroma_settings = None
client = None
if self.chroma_server_host:
chroma_settings = Settings(
chroma_server_cors_allow_origins=self.chroma_server_cors_allow_origins or [],
chroma_server_host=self.chroma_server_host,
chroma_server_http_port=self.chroma_server_http_port or None,
chroma_server_grpc_port=self.chroma_server_grpc_port or None,
chroma_server_ssl_enabled=self.chroma_server_ssl_enabled,
try:
from chromadb import HttpClient
except ImportError as e:
msg = "Could not import chromadb. Please install it with `pip install chromadb`."
raise ImportError(msg) from e
client = HttpClient(
host=self.chroma_server_host,
port=self.chroma_server_http_port or 8000,
ssl=bool(self.chroma_server_ssl_enabled),
)
client = Client(settings=chroma_settings)
# Check persist_directory and expand it if it is a relative path
persist_directory = self.resolve_path(self.persist_directory) if self.persist_directory is not None else None

View File

@ -2,6 +2,7 @@ from __future__ import annotations
import json
import re
from datetime import datetime, timezone
from typing import TYPE_CHECKING
from pydantic import ValidationError
@ -20,7 +21,7 @@ from lfx.base.models.unified_models import (
)
from lfx.base.models.watsonx_constants import IBM_WATSONX_URLS
from lfx.components.agentics.helpers.model_config import validate_model_selection
from lfx.components.helpers import CurrentDateComponent
from lfx.components.helpers import CalculatorComponent, CurrentDateComponent
from lfx.components.langchain_utilities.tool_calling import ToolCallingAgentComponent
from lfx.custom.custom_component.component import get_component_toolkit
from lfx.field_typing.range_spec import RangeSpec
@ -84,8 +85,14 @@ class AgentComponent(ToolCallingAgentComponent):
MultilineInput(
name="system_prompt",
display_name="Agent Instructions",
info="System Prompt: Initial instructions and context provided to guide the agent's behavior.",
value="You are a helpful assistant that can use tools to answer questions and perform tasks.",
info=(
"System Prompt: Initial instructions and context provided to guide the agent's behavior. "
"Supports dynamic placeholders: {current_date}, {model_name}."
),
value=(
"You are a helpful assistant that can use tools to answer questions and perform tasks. "
"Today is {current_date}. You are powered by {model_name}."
),
advanced=False,
),
MessageTextInput(
@ -181,6 +188,16 @@ class AgentComponent(ToolCallingAgentComponent):
info="If true, will add a tool to the agent that returns the current date.",
value=True,
),
BoolInput(
name="add_calculator_tool",
display_name="Calculator",
advanced=True,
info=(
"If true, adds a zero-config arithmetic calculator tool to the agent "
"(safe: only +, -, *, /, ** operators via AST)."
),
value=True,
),
]
outputs = [
Output(name="response", display_name="Response", method="message_response"),
@ -274,11 +291,60 @@ class AgentComponent(ToolCallingAgentComponent):
raise TypeError(msg)
self.tools.append(current_date_tool)
# Add calculator tool if enabled (zero-config arithmetic)
if getattr(self, "add_calculator_tool", False):
if not isinstance(self.tools, list): # type: ignore[has-type]
self.tools = []
calculator_tool = (await CalculatorComponent(**self.get_base_args()).to_toolkit()).pop(0)
if not isinstance(calculator_tool, StructuredTool):
msg = "CalculatorComponent must be converted to a StructuredTool"
raise TypeError(msg)
self.tools.append(calculator_tool)
# Set shared callbacks for tracing the tools used by the agent
self.set_tools_callbacks(self.tools, self._get_shared_callbacks())
return llm_model, self.chat_history, self.tools
def _get_resolved_model_name(self) -> str:
"""Best-effort human-readable model name for {model_name} injection."""
try:
from langchain_core.language_models import BaseLanguageModel
if isinstance(self.model, BaseLanguageModel):
return type(self.model).__name__
except ImportError:
pass
if isinstance(self.model, list) and self.model:
first = self.model[0]
if isinstance(first, dict):
name = first.get("name")
if isinstance(name, str) and name:
return name
legacy_model_name = getattr(self, "model_name", None)
if isinstance(legacy_model_name, str) and legacy_model_name:
return legacy_model_name
return ""
def _inject_dynamic_prompt_values(self, prompt: str | None) -> str | None:
"""Replace {current_date} / {model_name} placeholders in the system prompt.
Uses str.replace (not str.format) so user prompts containing literal braces
such as JSON examples ({"key": 1}) never break the agent.
"""
if not prompt:
return prompt
replacements = {
"{current_date}": datetime.now(tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC"),
"{model_name}": self._get_resolved_model_name(),
}
for placeholder, value in replacements.items():
prompt = prompt.replace(placeholder, value)
return prompt
async def message_response(self) -> Message:
try:
llm_model, self.chat_history, self.tools = await self.get_agent_requirements()
@ -288,7 +354,7 @@ class AgentComponent(ToolCallingAgentComponent):
tools=self.tools or [],
chat_history=self.chat_history,
input_value=self.input_value,
system_prompt=self.system_prompt,
system_prompt=self._inject_dynamic_prompt_values(self.system_prompt),
)
agent = self.create_agent_runnable()
result = await self.run_agent(agent)
@ -392,8 +458,10 @@ class AgentComponent(ToolCallingAgentComponent):
try:
system_components = []
# 1. Agent Instructions (system_prompt)
agent_instructions = getattr(self, "system_prompt", "") or ""
# 1. Agent Instructions (system_prompt).
# Inject dynamic placeholders HERE so user-authored format_instructions
# and schema descriptions appended later keep their literal {...} tokens.
agent_instructions = self._inject_dynamic_prompt_values(getattr(self, "system_prompt", "") or "") or ""
if agent_instructions:
system_components.append(f"{agent_instructions}")
@ -423,7 +491,9 @@ class AgentComponent(ToolCallingAgentComponent):
except (ValidationError, ValueError, TypeError, KeyError) as e:
await logger.aerror(f"Could not build schema for prompt: {e}", exc_info=True)
# Combine all components
# Combine all components. Injection already applied on agent_instructions
# above; do NOT re-run it here so literal {...} tokens in format
# instructions / schema descriptions stay intact.
combined_instructions = "\n\n".join(system_components) if system_components else ""
llm_model, self.chat_history, self.tools = await self.get_agent_requirements()
self.set(
@ -542,6 +612,7 @@ class AgentComponent(ToolCallingAgentComponent):
"tools",
"input_value",
"add_current_date_tool",
"add_calculator_tool",
"system_prompt",
"agent_description",
"max_iterations",

View File

@ -193,14 +193,7 @@ class YouTubeCommentsComponent(Component):
else:
request = None
# Convert to DataFrame
comments_df = pd.DataFrame(comments_data)
# Add video metadata
comments_df["video_id"] = video_id
comments_df["video_url"] = self.video_url
# Sort columns for better organization
# Define column order
column_order = [
"video_id",
"video_url",
@ -217,7 +210,20 @@ class YouTubeCommentsComponent(Component):
if self.include_metrics:
column_order.extend(["like_count", "reply_count"])
comments_df = comments_df[column_order]
# Handle empty comments case
if not comments_data:
# Create empty DataFrame with proper columns
comments_df = pd.DataFrame(columns=column_order)
else:
# Convert to DataFrame
comments_df = pd.DataFrame(comments_data)
# Add video metadata
comments_df["video_id"] = video_id
comments_df["video_url"] = self.video_url
# Reorder columns
comments_df = comments_df[column_order]
return DataFrame(comments_df)

View File

@ -272,8 +272,9 @@ def create_class(code, class_name):
module = ast.parse(code)
exec_globals = prepare_global_scope(module)
future_imports = [n for n in module.body if isinstance(n, ast.ImportFrom) and n.module == "__future__"]
class_code = extract_class_code(module, class_name)
compiled_class = compile_class_code(class_code)
compiled_class = compile_class_code(class_code, future_imports)
return build_class_constructor(compiled_class, exec_globals, class_name)
@ -394,11 +395,15 @@ def prepare_global_scope(module):
exec_globals = globals().copy()
imports = []
import_froms = []
future_imports = []
definitions = []
for node in module.body:
if isinstance(node, ast.Import):
imports.append(node)
elif isinstance(node, ast.ImportFrom) and node.module == "__future__":
# __future__ imports are compiler directives — collect separately
future_imports.append(node)
elif isinstance(node, ast.ImportFrom) and node.module is not None:
import_froms.append(node)
elif isinstance(node, ast.ClassDef | ast.FunctionDef | ast.Assign | ast.AnnAssign):
@ -468,7 +473,8 @@ def prepare_global_scope(module):
raise ModuleNotFoundError(msg)
if definitions:
combined_module = ast.Module(body=definitions, type_ignores=[])
# Prepend __future__ imports so compiler directives (e.g. PEP 563 annotations) take effect
combined_module = ast.Module(body=future_imports + definitions, type_ignores=[])
compiled_code = compile(combined_module, "<string>", "exec")
exec(compiled_code, exec_globals)
@ -491,16 +497,18 @@ def extract_class_code(module, class_name):
return class_code
def compile_class_code(class_code):
def compile_class_code(class_code, future_imports=None):
"""Compiles the AST node of a class into a code object.
Args:
class_code: AST node of the class
future_imports: Optional list of __future__ ImportFrom nodes to prepend as compiler directives
Returns:
Compiled code object of the class
"""
return compile(ast.Module(body=[class_code], type_ignores=[]), "<string>", "exec")
body = (future_imports or []) + [class_code]
return compile(ast.Module(body=body, type_ignores=[]), "<string>", "exec")
def build_class_constructor(compiled_class, exec_globals, class_name):

View File

@ -41,6 +41,9 @@ class ResultData(BaseModel):
if message is None:
continue
if not isinstance(message, dict):
continue
if "stream_url" in message and "type" in message:
stream_url = StreamURL(location=message["stream_url"])
values["outputs"].update({key: OutputValue(message=stream_url, type=message["type"])})

View File

@ -111,7 +111,7 @@ def build_output_logs(vertex, result) -> dict:
type_ = get_type(output_result)
match type_:
case LogType.STREAM if "stream_url" in message:
case LogType.STREAM if isinstance(message, dict) and "stream_url" in message:
message = StreamURL(location=message["stream_url"])
case LogType.STREAM:
@ -124,9 +124,13 @@ def build_output_logs(vertex, result) -> dict:
message = ""
case LogType.ARRAY:
if isinstance(message, DataFrame):
if message is None:
message = []
elif isinstance(message, DataFrame):
message = message.to_dict(orient="records")
message = [serialize(item) for item in message]
message = [serialize(item) for item in message]
else:
message = [serialize(item) for item in message]
name = output.get("name", f"output_{index}")
outputs |= {name: OutputValue(message=message, type=type_).model_dump()}

View File

@ -319,6 +319,7 @@ class Settings(BaseSettings):
max_flow_version_entries_per_flow: int = 50
"""Max version history entries per flow. Oldest entries pruned on next snapshot.
If retroactively lowered below the current count for a flow,
the oldest entries are deleted only when the next entry is created.
"""
@ -336,6 +337,7 @@ class Settings(BaseSettings):
max_items_length: int = MAX_ITEMS_LENGTH
"""Maximum number of items to store and display in the UI. Lists longer than this
will be truncated when displayed in the UI. Does not affect data passed between components nor outputs."""
max_ingestion_timeout_secs: int = 600
# MCP Server
mcp_server_enabled: bool = True

View File

@ -39,6 +39,34 @@ class TestLoggingComponent(Component):
assert result.__name__ == "TestLoggingComponent"
def test_create_class_future_annotations_with_type_checking():
"""Regression test for issue #12776.
`from __future__ import annotations` must act as a compiler directive so that TYPE_CHECKING-only
imports don't raise NameError at classdefinition time.
"""
code = dedent("""
from __future__ import annotations
from typing import TYPE_CHECKING
from langflow.custom import Component
if TYPE_CHECKING:
from typing import List
class TypeCheckingComponent(Component):
display_name = "Test"
def build(self, value: List[str]) -> str:
return str(value)
""")
result = create_class(code, "TypeCheckingComponent")
assert result.__name__ == "TypeCheckingComponent"
# With PEP 563 active, annotations should be stored as strings rather than evaluated
hints = result.build.__annotations__
assert hints.get("value") == "List[str]"
assert hints.get("return") == "str"
def test_execute_function_supports_aliased_dotted_imports():
code = dedent("""
import urllib.request as request

View File

@ -0,0 +1,13 @@
import threading
from lfx.graph.schema import ResultData
def test_result_data_ignores_non_dict_artifact_values():
"""Vector DB artifacts may include non-iterable objects such as locks."""
lock = threading.Lock()
result = ResultData(artifacts={"vector_db": lock})
assert result.artifacts == {"vector_db": lock}
assert result.outputs == {}