From fc59834aeaf6790ae8145765c5dcd4e646312b98 Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Wed, 16 Sep 2026 11:38:03 +0100 Subject: [PATCH 1/3] Add generic agentic bundle component passthrough Co-authored-by: omnigent --- src/flowx/bundler/dab_writer.py | 9 ++ src/flowx/ir_serde.py | 7 ++ src/flowx/models/ir.py | 21 +++++ .../activity_preparers/agentic_component.py | 28 ++++++ src/flowx/preparer/workflow_preparer.py | 3 + src/flowx/validate/bundle_invariants.py | 24 ++++- tests/unit/test_agentic_component.py | 89 +++++++++++++++++++ tests/unit/test_bundle_invariants.py | 46 ++++++++++ 8 files changed, 226 insertions(+), 1 deletion(-) create mode 100644 src/flowx/preparer/activity_preparers/agentic_component.py create mode 100644 tests/unit/test_agentic_component.py diff --git a/src/flowx/bundler/dab_writer.py b/src/flowx/bundler/dab_writer.py index 8cdf20b..f02423c 100644 --- a/src/flowx/bundler/dab_writer.py +++ b/src/flowx/bundler/dab_writer.py @@ -29,6 +29,7 @@ from flowx.models.dab import DabNotebook from flowx.models.ir import ( Activity, + AgenticComponentActivity, AppendVariableActivity, ControlEdge, CopyActivity, @@ -2279,6 +2280,14 @@ def _reconstruct_ir(task_ir: dict[str, Any]) -> Activity: consolidate_metadata_driven=bool(task_ir.get("consolidate_metadata_driven", False)), lookup_values=list(task_ir.get("lookup_values") or []), ) + if task_type == "AgenticComponentActivity": + return AgenticComponentActivity( + **base, + files=list(task_ir.get("files") or []), + resources=list(task_ir.get("resources") or []), + task=dict(task_ir.get("task") or {}), + raw_definition=task_ir.get("raw_definition"), + ) if task_type == "UnsupportedActivity": return UnsupportedActivity( **base, diff --git a/src/flowx/ir_serde.py b/src/flowx/ir_serde.py index 72b00a6..d865880 100644 --- a/src/flowx/ir_serde.py +++ b/src/flowx/ir_serde.py @@ -21,6 +21,7 @@ from flowx.models.ir import ( Activity, + AgenticComponentActivity, AppendVariableActivity, CopyActivity, DataAsset, @@ -242,6 +243,12 @@ def activity_extra_fields(activity: Activity) -> dict[str, Any]: extra: dict[str, Any] = {} match activity: + case AgenticComponentActivity(): + extra["files"] = activity.files + extra["resources"] = activity.resources + extra["task"] = activity.task + if activity.raw_definition is not None: + extra["raw_definition"] = activity.raw_definition case NotebookActivity(): extra["notebook_path"] = activity.notebook_path if activity.base_parameters: diff --git a/src/flowx/models/ir.py b/src/flowx/models/ir.py index a1b097a..1b70937 100644 --- a/src/flowx/models/ir.py +++ b/src/flowx/models/ir.py @@ -685,6 +685,27 @@ class DbtFactoryActivity(Activity): nodes: list[dict[str, Any]] = field(default_factory=list) +@dataclass(slots=True, kw_only=True) +class AgenticComponentActivity(Activity): + """Bundle components authored for a source activity the typed engine cannot express. + + Attributes: + files: Files to write below the bundle's ``src`` directory. Each entry + carries ``path`` and either UTF-8 ``content`` or base64-encoded + ``binary_content``. + resources: Pipeline resources in the existing ``resource_key`` plus + raw ``definition`` shape used by the bundle writer. + task: Raw Databricks task fragment containing either ``pipeline_task`` + or ``notebook_task`` wiring to an authored resource or file. + raw_definition: Original source definition retained for auditing. + """ + + files: list[dict[str, Any]] = field(default_factory=list) + resources: list[dict[str, Any]] = field(default_factory=list) + task: dict[str, Any] = field(default_factory=dict) + raw_definition: dict[str, Any] | None = None + + @dataclass(slots=True, kw_only=True) class UnsupportedActivity(Activity): """Sentinel for activities that could not be translated. diff --git a/src/flowx/preparer/activity_preparers/agentic_component.py b/src/flowx/preparer/activity_preparers/agentic_component.py new file mode 100644 index 0000000..738348e --- /dev/null +++ b/src/flowx/preparer/activity_preparers/agentic_component.py @@ -0,0 +1,28 @@ +"""Lower agent-authored bundle components through the existing output channels.""" + +from __future__ import annotations + +import base64 + +from flowx.models.dab import DabNotebook +from flowx.models.ir import AgenticComponentActivity +from flowx.preparer.workflow_preparer import PreparedActivity, build_common_task_fields + + +def prepare(activity: AgenticComponentActivity, *, scope: str = "") -> PreparedActivity: + """Pass authored files, resources, and task wiring to the bundle writer.""" + del scope + notebooks = [ + DabNotebook( + relative_path=str(file["path"]), + content=str(file.get("content", "")), + binary_content=(base64.b64decode(str(file["binary_content"])) if "binary_content" in file else None), + ) + for file in activity.files + ] + task = {**build_common_task_fields(activity), **activity.task, "task_key": activity.task_key} + return PreparedActivity( + task=task, + notebooks=notebooks, + pipeline_resources=list(activity.resources), + ) diff --git a/src/flowx/preparer/workflow_preparer.py b/src/flowx/preparer/workflow_preparer.py index 3ab1df1..e7d664d 100644 --- a/src/flowx/preparer/workflow_preparer.py +++ b/src/flowx/preparer/workflow_preparer.py @@ -9,6 +9,7 @@ from flowx.models.dab import DabNotebook, ParameterApproximation, SecretInstruction, SetupTask from flowx.models.ir import ( Activity, + AgenticComponentActivity, AppendVariableActivity, CopyActivity, DbtFactoryActivity, @@ -142,6 +143,7 @@ def prepare_activity( ) -> PreparedActivity: """Dispatches to the appropriate activity preparer based on activity type.""" from flowx.preparer.activity_preparers import ( + agentic_component, append_variable, copy, databricks_job, @@ -164,6 +166,7 @@ def prepare_activity( ) dispatch: dict[type, Any] = { + AgenticComponentActivity: agentic_component.prepare, NotebookActivity: notebook.prepare, SparkJarActivity: spark_jar.prepare, SparkPythonActivity: spark_python.prepare, diff --git a/src/flowx/validate/bundle_invariants.py b/src/flowx/validate/bundle_invariants.py index 2ed275f..153b30c 100644 --- a/src/flowx/validate/bundle_invariants.py +++ b/src/flowx/validate/bundle_invariants.py @@ -25,6 +25,7 @@ _ANCHOR_RE = re.compile(r"[&*]id\d+\b") _JOB_PARAM_REF_RE = re.compile(r"\{\{\s*job\.parameters\.([A-Za-z0-9_]+)\s*\}\}") _JOB_RESOURCE_ID_RE = re.compile(r"\$\{resources\.jobs\.([^.}]+)\.id\}") +_PIPELINE_RESOURCE_ID_RE = re.compile(r"\$\{resources\.pipelines\.([^.}]+)\.id\}") _PYDABS_JOB_RE = re.compile(r"resources\.add_job\(\s*['\"]([^'\"]+)['\"]") @@ -246,10 +247,15 @@ def check_bundle_dir(bundle_dir: Path) -> BundleInvariantResult: documents.append((path, document)) known_jobs: set[str] = set() + known_pipelines: set[str] = set() for _path, document in documents: - jobs = (document.get("resources") or {}).get("jobs") or {} + document_resources = document.get("resources") or {} + jobs = document_resources.get("jobs") or {} if isinstance(jobs, dict): known_jobs.update(str(job_key) for job_key in jobs) + pipelines = document_resources.get("pipelines") or {} + if isinstance(pipelines, dict): + known_pipelines.update(str(pipeline_key) for pipeline_key in pipelines) python_resources = (document.get("python") or {}).get("resources") or [] for resource in python_resources: if not isinstance(resource, str): @@ -268,6 +274,22 @@ def check_bundle_dir(bundle_dir: Path) -> BundleInvariantResult: if not isinstance(job, dict): continue for task in _iter_tasks(job.get("tasks") or []): + pipeline_task = task.get("pipeline_task") or {} + pipeline_id = pipeline_task.get("pipeline_id") if isinstance(pipeline_task, dict) else None + pipeline_match = ( + _PIPELINE_RESOURCE_ID_RE.fullmatch(pipeline_id) if isinstance(pipeline_id, str) else None + ) + if pipeline_match is not None and pipeline_match.group(1) not in known_pipelines: + findings.append( + BundleFinding( + code="dangling_pipeline_reference", + location=f"{path.name}, job '{job_key}', task '{task.get('task_key', '')}'", + message=( + f"pipeline_task references bundle pipeline '{pipeline_match.group(1)}', which is not " + "declared in static resource YAML." + ), + ) + ) run_job = task.get("run_job_task") or {} job_id = run_job.get("job_id") if isinstance(run_job, dict) else None match = _JOB_RESOURCE_ID_RE.fullmatch(job_id) if isinstance(job_id, str) else None diff --git a/tests/unit/test_agentic_component.py b/tests/unit/test_agentic_component.py new file mode 100644 index 0000000..2b66b39 --- /dev/null +++ b/tests/unit/test_agentic_component.py @@ -0,0 +1,89 @@ +"""Tests for generic agent-authored bundle components.""" + +from __future__ import annotations + +import base64 +import json + +import yaml + +from flowx.bundler.dab_writer import pipeline_dict_to_ir, write_bundle +from flowx.ir_serde import activity_to_dict, pipeline_to_dict +from flowx.models.ir import AgenticComponentActivity, Pipeline +from flowx.preparer.workflow_preparer import prepare_workflow +from flowx.validate.bundle_invariants import check_bundle_dir + +SOURCE_DEFINITION = {"type": "ExecuteDataFlow", "typeProperties": {"dataflow": "orders"}} +FILES = [ + {"path": "pipelines/orders.py", "content": "from pyspark import pipelines as dp\n"}, + {"path": "libraries/orders.whl", "binary_content": "UEsDBAoAAAAA"}, +] +PIPELINE_DEFINITION = { + "name": "orders_ingestion", + "catalog": "${var.catalog}", + "target": "${var.schema}", + "ingestion_definition": { + "connection_name": "flowx_orders_connection", + "objects": [ + { + "table": { + "source_catalog": "sales", + "source_schema": "dbo", + "source_table": "orders", + "destination_catalog": "${var.catalog}", + "destination_schema": "${var.schema}", + "destination_table": "orders", + } + } + ], + }, +} +RESOURCES = [{"resource_key": "orders_ingestion", "definition": PIPELINE_DEFINITION}] +TASK = {"pipeline_task": {"pipeline_id": "${resources.pipelines.orders_ingestion.id}"}} + + +def _activity() -> AgenticComponentActivity: + return AgenticComponentActivity( + name="Ingest orders", + task_key="ingest_orders", + files=FILES, + resources=RESOURCES, + task=TASK, + raw_definition=SOURCE_DEFINITION, + ) + + +def test_agentic_component_round_trips_through_serialized_ir(): + serialized = json.loads(json.dumps(activity_to_dict(_activity()))) + pipeline, _ = pipeline_dict_to_ir({"name": "orders", "tasks": [serialized]}) + restored = pipeline.tasks[0] + + assert type(restored).__name__ == "AgenticComponentActivity" + assert activity_to_dict(restored) == serialized + assert serialized == { + "name": "Ingest orders", + "task_key": "ingest_orders", + "type": "AgenticComponentActivity", + "files": FILES, + "resources": RESOURCES, + "task": TASK, + "raw_definition": SOURCE_DEFINITION, + } + + +def test_agentic_component_packages_authored_files_resource_and_pipeline_task(tmp_path): + serialized_pipeline = json.loads(json.dumps(pipeline_to_dict(Pipeline(name="orders", tasks=[_activity()])))) + restored_pipeline, _ = pipeline_dict_to_ir(serialized_pipeline) + + write_bundle(prepare_workflow(restored_pipeline), tmp_path) + + assert (tmp_path / "src" / "pipelines" / "orders.py").read_text(encoding="utf-8") == FILES[0]["content"] + assert (tmp_path / "src" / "libraries" / "orders.whl").read_bytes() == base64.b64decode(FILES[1]["binary_content"]) + pipeline_resource = yaml.safe_load((tmp_path / "resources" / "orders_ingestion.yml").read_text(encoding="utf-8")) + assert pipeline_resource == {"resources": {"pipelines": {"orders_ingestion": PIPELINE_DEFINITION}}} + + job_resource = yaml.safe_load((tmp_path / "resources" / "orders.yml").read_text(encoding="utf-8")) + assert job_resource["resources"]["jobs"]["orders"]["tasks"] == [ + {"task_key": "ingest_orders", "pipeline_task": TASK["pipeline_task"]} + ] + assert check_bundle_dir(tmp_path).ok diff --git a/tests/unit/test_bundle_invariants.py b/tests/unit/test_bundle_invariants.py index 8b52e55..5e4ead4 100644 --- a/tests/unit/test_bundle_invariants.py +++ b/tests/unit/test_bundle_invariants.py @@ -130,3 +130,49 @@ def test_bundle_job_reference_to_unknown_resource_is_flagged(tmp_path): assert "parent.yml" in finding.location assert "call_missing" in finding.location assert "dangling_run_job_reference" in format_result(result) + + +def test_bundle_pipeline_reference_can_target_pipeline_in_another_resource_file(tmp_path): + resources = tmp_path / "resources" + resources.mkdir() + (resources / "job.yml").write_text( + "resources:\n" + " jobs:\n" + " parent:\n" + " tasks:\n" + " - task_key: run_pipeline\n" + " pipeline_task:\n" + " pipeline_id: ${resources.pipelines.ingestion.id}\n", + encoding="utf-8", + ) + (resources / "pipeline.yml").write_text( + "resources:\n pipelines:\n ingestion:\n name: ingestion\n", + encoding="utf-8", + ) + + result = check_bundle_dir(tmp_path) + + assert "dangling_pipeline_reference" not in _codes(result.findings) + + +def test_bundle_pipeline_reference_to_unknown_resource_is_flagged(tmp_path): + resources = tmp_path / "resources" + resources.mkdir() + (resources / "job.yml").write_text( + "resources:\n" + " jobs:\n" + " parent:\n" + " tasks:\n" + " - task_key: run_pipeline\n" + " pipeline_task:\n" + " pipeline_id: ${resources.pipelines.missing.id}\n", + encoding="utf-8", + ) + + result = check_bundle_dir(tmp_path) + finding = next(finding for finding in result.findings if finding.code == "dangling_pipeline_reference") + + assert finding.severity == "violation" + assert "job.yml" in finding.location + assert "run_pipeline" in finding.location + assert "dangling_pipeline_reference" in format_result(result) From cc47ead43810accf18a94eeeee3582f3134e6ada Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Wed, 16 Sep 2026 11:42:52 +0100 Subject: [PATCH 2/3] Keep agentic files within bundle source tree Co-authored-by: omnigent --- .../activity_preparers/agentic_component.py | 13 ++++++- tests/unit/test_agentic_component.py | 37 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/src/flowx/preparer/activity_preparers/agentic_component.py b/src/flowx/preparer/activity_preparers/agentic_component.py index 738348e..fcbec40 100644 --- a/src/flowx/preparer/activity_preparers/agentic_component.py +++ b/src/flowx/preparer/activity_preparers/agentic_component.py @@ -3,18 +3,29 @@ from __future__ import annotations import base64 +from pathlib import PurePosixPath from flowx.models.dab import DabNotebook from flowx.models.ir import AgenticComponentActivity from flowx.preparer.workflow_preparer import PreparedActivity, build_common_task_fields +def _source_relative_path(raw_path: object) -> str: + """Return a safe path that the shared writer always treats as relative to ``src``.""" + relative_path = PurePosixPath(str(raw_path)) + if relative_path.is_absolute() or relative_path.as_posix() == "." or ".." in relative_path.parts: + raise ValueError(f"Agentic component file path {raw_path!r} must be relative to the bundle src directory") + # The bundle writer reserves resources/* and pyproject.toml for PyDABs root artifacts. + # A leading ./ preserves the authored path while keeping agentic files on its src channel. + return f"./{relative_path.as_posix()}" + + def prepare(activity: AgenticComponentActivity, *, scope: str = "") -> PreparedActivity: """Pass authored files, resources, and task wiring to the bundle writer.""" del scope notebooks = [ DabNotebook( - relative_path=str(file["path"]), + relative_path=_source_relative_path(file["path"]), content=str(file.get("content", "")), binary_content=(base64.b64decode(str(file["binary_content"])) if "binary_content" in file else None), ) diff --git a/tests/unit/test_agentic_component.py b/tests/unit/test_agentic_component.py index 2b66b39..46c9f21 100644 --- a/tests/unit/test_agentic_component.py +++ b/tests/unit/test_agentic_component.py @@ -5,6 +5,7 @@ import base64 import json +import pytest import yaml from flowx.bundler.dab_writer import pipeline_dict_to_ir, write_bundle @@ -87,3 +88,39 @@ def test_agentic_component_packages_authored_files_resource_and_pipeline_task(tm {"task_key": "ingest_orders", "pipeline_task": TASK["pipeline_task"]} ] assert check_bundle_dir(tmp_path).ok + + +def test_agentic_component_notebook_files_with_reserved_paths_stay_under_src(tmp_path): + activity = AgenticComponentActivity( + name="Custom notebook", + task_key="custom_notebook", + files=[ + {"path": "resources/custom.py", "content": "print('custom')\n"}, + {"path": "pyproject.toml", "content": "[project]\nname = 'custom'\n"}, + ], + task={"notebook_task": {"notebook_path": "../src/resources/custom.py"}}, + ) + + write_bundle(prepare_workflow(Pipeline(name="custom", tasks=[activity])), tmp_path) + + assert (tmp_path / "src" / "resources" / "custom.py").read_text(encoding="utf-8") == "print('custom')\n" + assert (tmp_path / "src" / "pyproject.toml").read_text(encoding="utf-8") == "[project]\nname = 'custom'\n" + assert not (tmp_path / "resources" / "custom.py").exists() + assert not (tmp_path / "pyproject.toml").exists() + job_resource = yaml.safe_load((tmp_path / "resources" / "custom.yml").read_text(encoding="utf-8")) + assert job_resource["resources"]["jobs"]["custom"]["tasks"][0]["notebook_task"] == { + "notebook_path": "../src/resources/custom.py" + } + + +@pytest.mark.parametrize("path", ["../outside.py", "/tmp/outside.py"]) +def test_agentic_component_rejects_file_paths_that_escape_src(path): + activity = AgenticComponentActivity( + name="Unsafe file", + task_key="unsafe_file", + files=[{"path": path, "content": "unsafe\n"}], + task={"notebook_task": {"notebook_path": "../src/safe.py"}}, + ) + + with pytest.raises(ValueError, match="relative to the bundle src directory"): + prepare_workflow(Pipeline(name="unsafe", tasks=[activity])) From 6e63191a580ad5efd99b99a250c63be387c30c8e Mon Sep 17 00:00:00 2001 From: Matthew Moorcroft Date: Wed, 16 Sep 2026 11:48:44 +0100 Subject: [PATCH 3/3] Preserve agentic paths through shared bundle packaging Co-authored-by: omnigent --- src/flowx/bundler/dab_writer.py | 15 +++--- src/flowx/models/dab.py | 4 ++ .../activity_preparers/agentic_component.py | 18 ++++--- tests/unit/test_agentic_component.py | 47 +++++++++++++++++-- 4 files changed, 67 insertions(+), 17 deletions(-) diff --git a/src/flowx/bundler/dab_writer.py b/src/flowx/bundler/dab_writer.py index f02423c..93b8933 100644 --- a/src/flowx/bundler/dab_writer.py +++ b/src/flowx/bundler/dab_writer.py @@ -239,11 +239,7 @@ def write_bundle( src_dir = output_dir / "src" def _write_generated(notebooks: list[DabNotebook]) -> None: - root_artifacts = [ - notebook - for notebook in notebooks - if notebook.relative_path.startswith("resources/") or notebook.relative_path == "pyproject.toml" - ] + root_artifacts = [notebook for notebook in notebooks if _is_bundle_root_artifact(notebook)] rest = [notebook for notebook in notebooks if notebook not in root_artifacts] if rest: created_files.extend(write_notebooks(rest, src_dir)) @@ -661,6 +657,13 @@ def _known_bundle_job_keys(workflow: PreparedWorkflow, resource_key: str) -> set return keys +def _is_bundle_root_artifact(notebook: DabNotebook) -> bool: + """Return whether a generated file belongs at the bundle root instead of below ``src``.""" + if notebook.write_to_bundle_root is not None: + return notebook.write_to_bundle_root + return notebook.relative_path.startswith("resources/") or notebook.relative_path == "pyproject.toml" + + def _namespace_workflow_assets(workflow: PreparedWorkflow) -> PreparedWorkflow: """Namespaces generated source files by DAG while preserving workspace paths.""" cloned = copy.deepcopy(workflow) @@ -697,7 +700,7 @@ def _namespace_workflow_assets(workflow: PreparedWorkflow) -> PreparedWorkflow: notebook.content = notebook.content.replace("src/dbt_project", f"src/{prefix}/dbt_project") notebook.content = notebook.content.replace("src/dbt_profiles", f"src/{prefix}/dbt_profiles") continue - if original_path.startswith("resources/") or original_path == "pyproject.toml": + if _is_bundle_root_artifact(notebook): continue notebook.relative_path = f"{prefix}/{original_path}" replacements[f"../src/{original_path}"] = f"../src/{notebook.relative_path}" diff --git a/src/flowx/models/dab.py b/src/flowx/models/dab.py index ec0b85c..334d045 100644 --- a/src/flowx/models/dab.py +++ b/src/flowx/models/dab.py @@ -20,12 +20,16 @@ class DabNotebook: language: Notebook language (``"python"``, ``"sql"``, ``"scala"``, ``"r"``). binary_content: Raw bytes for binary files (e.g. JARs). When set, the notebook writer writes these bytes instead of ``content``. + write_to_bundle_root: Controls whether the bundle writer places this + file at the bundle root instead of below ``src``. ``None`` keeps + the legacy path-based routing for PyDABs artifacts. """ relative_path: str content: str = "" language: str = "python" binary_content: bytes | None = None + write_to_bundle_root: bool | None = None # --------------------------------------------------------------------------- diff --git a/src/flowx/preparer/activity_preparers/agentic_component.py b/src/flowx/preparer/activity_preparers/agentic_component.py index fcbec40..9cfd91b 100644 --- a/src/flowx/preparer/activity_preparers/agentic_component.py +++ b/src/flowx/preparer/activity_preparers/agentic_component.py @@ -3,7 +3,7 @@ from __future__ import annotations import base64 -from pathlib import PurePosixPath +from pathlib import PurePosixPath, PureWindowsPath from flowx.models.dab import DabNotebook from flowx.models.ir import AgenticComponentActivity @@ -11,13 +11,18 @@ def _source_relative_path(raw_path: object) -> str: - """Return a safe path that the shared writer always treats as relative to ``src``.""" + """Return a safe path relative to the bundle's ``src`` directory.""" relative_path = PurePosixPath(str(raw_path)) - if relative_path.is_absolute() or relative_path.as_posix() == "." or ".." in relative_path.parts: + windows_path = PureWindowsPath(str(raw_path)) + if ( + relative_path.is_absolute() + or windows_path.is_absolute() + or relative_path.as_posix() == "." + or ".." in relative_path.parts + or ".." in windows_path.parts + ): raise ValueError(f"Agentic component file path {raw_path!r} must be relative to the bundle src directory") - # The bundle writer reserves resources/* and pyproject.toml for PyDABs root artifacts. - # A leading ./ preserves the authored path while keeping agentic files on its src channel. - return f"./{relative_path.as_posix()}" + return relative_path.as_posix() def prepare(activity: AgenticComponentActivity, *, scope: str = "") -> PreparedActivity: @@ -28,6 +33,7 @@ def prepare(activity: AgenticComponentActivity, *, scope: str = "") -> PreparedA relative_path=_source_relative_path(file["path"]), content=str(file.get("content", "")), binary_content=(base64.b64decode(str(file["binary_content"])) if "binary_content" in file else None), + write_to_bundle_root=False, ) for file in activity.files ] diff --git a/tests/unit/test_agentic_component.py b/tests/unit/test_agentic_component.py index 46c9f21..8186c81 100644 --- a/tests/unit/test_agentic_component.py +++ b/tests/unit/test_agentic_component.py @@ -8,7 +8,7 @@ import pytest import yaml -from flowx.bundler.dab_writer import pipeline_dict_to_ir, write_bundle +from flowx.bundler.dab_writer import _combine_airflow_workflows, pipeline_dict_to_ir, write_bundle from flowx.ir_serde import activity_to_dict, pipeline_to_dict from flowx.models.ir import AgenticComponentActivity, Pipeline from flowx.preparer.workflow_preparer import prepare_workflow @@ -108,12 +108,49 @@ def test_agentic_component_notebook_files_with_reserved_paths_stay_under_src(tmp assert not (tmp_path / "resources" / "custom.py").exists() assert not (tmp_path / "pyproject.toml").exists() job_resource = yaml.safe_load((tmp_path / "resources" / "custom.yml").read_text(encoding="utf-8")) - assert job_resource["resources"]["jobs"]["custom"]["tasks"][0]["notebook_task"] == { - "notebook_path": "../src/resources/custom.py" - } + assert job_resource["resources"]["jobs"]["custom"]["tasks"][0]["notebook_task"]["notebook_path"] == ( + "../src/resources/custom.py" + ) + + +def test_agentic_component_notebook_path_tracks_shared_workflow_namespacing(tmp_path): + custom = AgenticComponentActivity( + name="Custom notebook", + task_key="custom_notebook", + files=[{"path": "resources/custom.py", "content": "print('custom')\n"}], + task={"notebook_task": {"notebook_path": "../src/resources/custom.py"}}, + ) + other = AgenticComponentActivity( + name="Other notebook", + task_key="other_notebook", + files=[{"path": "notebooks/other.py", "content": "print('other')\n"}], + task={"notebook_task": {"notebook_path": "../src/notebooks/other.py"}}, + ) + combined = _combine_airflow_workflows( + [ + prepare_workflow(Pipeline(name="custom", tasks=[custom])), + prepare_workflow(Pipeline(name="other", tasks=[other])), + ] + ) + + write_bundle(combined, tmp_path) + + assert (tmp_path / "src" / "custom" / "resources" / "custom.py").exists() + custom_job = yaml.safe_load((tmp_path / "resources" / "custom.yml").read_text(encoding="utf-8")) + assert custom_job["resources"]["jobs"]["custom"]["tasks"][0]["notebook_task"]["notebook_path"] == ( + "../src/custom/resources/custom.py" + ) + assert (tmp_path / "src" / "other" / "notebooks" / "other.py").exists() + other_job = yaml.safe_load((tmp_path / "resources" / "other.yml").read_text(encoding="utf-8")) + assert other_job["resources"]["jobs"]["other"]["tasks"][0]["notebook_task"]["notebook_path"] == ( + "../src/other/notebooks/other.py" + ) -@pytest.mark.parametrize("path", ["../outside.py", "/tmp/outside.py"]) +@pytest.mark.parametrize( + "path", + ["../outside.py", "/tmp/outside.py", "..\\outside.py", "C:\\tmp\\outside.py"], +) def test_agentic_component_rejects_file_paths_that_escape_src(path): activity = AgenticComponentActivity( name="Unsafe file",