From 0d98c9c10e5f307f39645cc86e89daafc18f188d Mon Sep 17 00:00:00 2001 From: omercengiz Date: Thu, 20 Aug 2026 21:12:27 +0300 Subject: [PATCH 01/13] Add MCP server and DAG analysis tools --- pyproject.toml | 3 +- src/flowsense/cli/main.py | 66 +++++----------------- src/flowsense/engine/analyzer.py | 65 ++++++++++++++++++++++ src/flowsense/mcp/__init__.py | 0 src/flowsense/mcp/server.py | 53 ++++++++++++++++++ src/flowsense/models/__init__.py | 6 +- src/flowsense/models/dag_analysis.py | 17 ++++++ tests/test_analyzer.py | 50 +++++++++++++++++ tests/test_mcp_server.py | 82 ++++++++++++++++++++++++++++ uv.lock | 19 ++++++- 10 files changed, 306 insertions(+), 55 deletions(-) create mode 100644 src/flowsense/engine/analyzer.py create mode 100644 src/flowsense/mcp/__init__.py create mode 100644 src/flowsense/mcp/server.py create mode 100644 src/flowsense/models/dag_analysis.py create mode 100644 tests/test_analyzer.py create mode 100644 tests/test_mcp_server.py diff --git a/pyproject.toml b/pyproject.toml index d50a021..f966611 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,11 +21,12 @@ dev = [ ] mcp = [ - "mcp", + "mcp[cli]>=2.0", ] [project.scripts] flowsense = "flowsense.cli.main:app" +flowsense-mcp = "flowsense.mcp.server:main" [build-system] requires = ["setuptools>=75"] diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index 429e0c5..ee01015 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -4,10 +4,7 @@ from rich.console import Console from rich.table import Table -from flowsense.collector.airflow_client import AirflowClient -from flowsense.engine.drift import calculate_drift -from flowsense.engine.history import build_duration_history -from flowsense.engine.propagation import analyze_propagation +from flowsense.engine.analyzer import analyze_dag app = typer.Typer( name="flowsense", @@ -27,39 +24,12 @@ def main() -> None: def analyze( dag_id: str = typer.Argument( ..., - help="Airflow DAG ID to analyze.", + help="Airflow DAG id to analyze.", ), ) -> None: - client = AirflowClient() + analysis = analyze_dag(dag_id) - task_runs = client.collect_task_runs(dag_id) - - if not task_runs: - console.print(f"[red]No successful task runs found for DAG: {dag_id}[/red]") - raise typer.Exit(code=1) - - history = build_duration_history(task_runs) - - drift_results = {} - - for task_id, durations in history.items(): - result = calculate_drift( - task_id, - durations, - ) - - drift_results[task_id] = result - - dependencies = client.get_dag_dependencies(dag_id) - - propagation_results = analyze_propagation( - drift_results, - dependencies, - ) - - console.print() - console.print(f"[bold]FlowSense Analysis[/bold] — {dag_id}") - console.print() + console.print(f"\n[bold]FlowSense Analysis — {analysis.dag_id}[/bold]\n") table = Table() @@ -70,9 +40,9 @@ def analyze( table.add_column("Z-Score") table.add_column("Severity") - for result in drift_results.values(): + for task_id, result in analysis.drift_results.items(): table.add_row( - result.task_id, + task_id, f"{result.baseline:.2f}s", f"{result.current:.2f}s", f"{result.deviation_percent:+.1f}%", @@ -82,21 +52,15 @@ def analyze( console.print(table) - console.print() - console.print("[bold]Propagation Analysis[/bold]") - - if not propagation_results: - console.print("No propagation detected.") - return - - for propagation in propagation_results: - console.print() - console.print(f"Origin: [bold]{propagation.origin_task}[/bold]") - - console.print("Path: " + " -> ".join(propagation.path)) + console.print(f"\nOverall Severity: [bold]{analysis.overall_severity}[/bold]") - console.print(f"Propagation Score: {propagation.propagation_score:.2f}") + if analysis.primary_origin: + console.print(f"Primary Origin: [bold]{analysis.primary_origin}[/bold]") + if analysis.propagation_results: + console.print("\n[bold]Propagation Analysis[/bold]\n") -if __name__ == "__main__": - app() + for result in analysis.propagation_results: + console.print(f"Origin: {result.origin_task}") + console.print(f"Path: {' -> '.join(result.path)}") + console.print(f"Propagation Score: {result.propagation_score:.2f}") diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py new file mode 100644 index 0000000..4ed9903 --- /dev/null +++ b/src/flowsense/engine/analyzer.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from flowsense.collector.airflow_client import AirflowClient +from flowsense.engine.drift import calculate_drift +from flowsense.engine.history import build_duration_history +from flowsense.engine.propagation import analyze_propagation +from flowsense.models import DAGAnalysis + + +def analyze_dag(dag_id: str) -> DAGAnalysis: + client = AirflowClient() + + task_runs = client.collect_task_runs(dag_id) + duration_history = build_duration_history(task_runs) + + drift_results = {} + + for task_id, durations in duration_history.items(): + try: + drift_results[task_id] = calculate_drift( + task_id=task_id, + durations=durations, + ) + except ValueError: + continue + + dependencies = client.get_dag_dependencies(dag_id) + + propagation_results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + severity_order = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, + } + + overall_severity = "NORMAL" + + if drift_results: + overall_severity = max( + drift_results.values(), + key=lambda result: severity_order[result.severity], + ).severity + + primary_origin = None + + if propagation_results: + primary_origin = max( + propagation_results, + key=lambda result: result.propagation_score, + ).origin_task + + return DAGAnalysis( + dag_id=dag_id, + runs_analyzed=len({run.dag_run_id for run in task_runs}), + overall_severity=overall_severity, + primary_origin=primary_origin, + drift_results=drift_results, + propagation_results=propagation_results, + dependencies=dependencies, + ) diff --git a/src/flowsense/mcp/__init__.py b/src/flowsense/mcp/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py new file mode 100644 index 0000000..43b3729 --- /dev/null +++ b/src/flowsense/mcp/server.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from mcp.server import MCPServer + +from flowsense.engine.analyzer import analyze_dag +from flowsense.models import DAGAnalysis + +mcp = MCPServer("FlowSense Engine") + + +def serialize_analysis(analysis: DAGAnalysis) -> dict: + return { + "dag_id": analysis.dag_id, + "runs_analyzed": analysis.runs_analyzed, + "overall_severity": analysis.overall_severity, + "primary_origin": analysis.primary_origin, + "drift_results": { + task_id: { + "baseline": result.baseline, + "current": result.current, + "mad": result.mad, + "robust_z_score": result.robust_z_score, + "deviation_percent": result.deviation_percent, + "severity": result.severity, + } + for task_id, result in analysis.drift_results.items() + }, + "propagation_results": [ + { + "origin_task": result.origin_task, + "affected_tasks": result.affected_tasks, + "path": result.path, + "propagation_score": result.propagation_score, + } + for result in analysis.propagation_results + ], + "dependencies": analysis.dependencies, + } + + +@mcp.tool() +def analyze_airflow_dag(dag_id: str) -> dict: + """Analyze an Apache Airflow DAG for temporal drift and propagation.""" + analysis = analyze_dag(dag_id) + return serialize_analysis(analysis) + + +def main() -> None: + mcp.run() + + +if __name__ == "__main__": + main() diff --git a/src/flowsense/models/__init__.py b/src/flowsense/models/__init__.py index 7de82b2..133820d 100644 --- a/src/flowsense/models/__init__.py +++ b/src/flowsense/models/__init__.py @@ -1,3 +1,7 @@ +from flowsense.models.dag_analysis import DAGAnalysis from flowsense.models.task_run import TaskRun -__all__ = ["TaskRun"] +__all__ = [ + "DAGAnalysis", + "TaskRun", +] diff --git a/src/flowsense/models/dag_analysis.py b/src/flowsense/models/dag_analysis.py new file mode 100644 index 0000000..b735d81 --- /dev/null +++ b/src/flowsense/models/dag_analysis.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult +from flowsense.engine.propagation import PropagationResult + + +@dataclass +class DAGAnalysis: + dag_id: str + runs_analyzed: int + overall_severity: str + primary_origin: str | None + drift_results: dict[str, DriftResult] + propagation_results: list[PropagationResult] + dependencies: dict[str, list[str]] diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py new file mode 100644 index 0000000..a462261 --- /dev/null +++ b/tests/test_analyzer.py @@ -0,0 +1,50 @@ +from unittest.mock import MagicMock, patch + +from flowsense.engine.analyzer import analyze_dag +from flowsense.models import TaskRun + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_identifies_primary_origin( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="transform", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [3.0, 3.1, 2.9, 3.0, 9.5], + start=1, + ) + ] + [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="load", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [1.0, 1.1, 0.9, 1.0, 2.0], + start=1, + ) + ] + + client.get_dag_dependencies.return_value = { + "transform": ["load"], + "load": [], + } + + analysis = analyze_dag("demo") + + assert analysis.overall_severity == "CRITICAL" + assert analysis.primary_origin == "transform" + assert len(analysis.propagation_results) == 1 + assert analysis.propagation_results[0].origin_task == "transform" + assert analysis.propagation_results[0].affected_tasks == ["load"] diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py new file mode 100644 index 0000000..6d101da --- /dev/null +++ b/tests/test_mcp_server.py @@ -0,0 +1,82 @@ +import pytest +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + +from flowsense.engine.drift import DriftResult +from flowsense.mcp.server import serialize_analysis +from flowsense.models import DAGAnalysis + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +@pytest.mark.anyio +async def test_mcp_server_exposes_analyze_tool() -> None: + server_params = StdioServerParameters( + command="uv", + args=[ + "run", + "python", + "-m", + "flowsense.mcp.server", + ], + ) + + async with ( + stdio_client(server_params) as (read, write), + ClientSession(read, write) as session, + ): + await session.initialize() + + response = await session.list_tools() + + tool_names = [tool.name for tool in response.tools] + + assert "analyze_airflow_dag" in tool_names + + +def test_serialize_analysis() -> None: + analysis = DAGAnalysis( + dag_id="demo", + runs_analyzed=5, + overall_severity="CRITICAL", + primary_origin="transform", + drift_results={ + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=9.5, + mad=0.1, + robust_z_score=43.84, + deviation_percent=216.67, + severity="CRITICAL", + ) + }, + propagation_results=[], + dependencies={ + "transform": [], + }, + ) + + result = serialize_analysis(analysis) + + assert result["dag_id"] == "demo" + assert result["runs_analyzed"] == 5 + assert result["overall_severity"] == "CRITICAL" + assert result["primary_origin"] == "transform" + + transform = result["drift_results"]["transform"] + + assert transform["baseline"] == 3.0 + assert transform["current"] == 9.5 + assert transform["mad"] == 0.1 + assert transform["robust_z_score"] == 43.84 + assert transform["deviation_percent"] == 216.67 + assert transform["severity"] == "CRITICAL" + + assert result["propagation_results"] == [] + assert result["dependencies"] == { + "transform": [], + } diff --git a/uv.lock b/uv.lock index e0a7539..64a7413 100644 --- a/uv.lock +++ b/uv.lock @@ -235,13 +235,13 @@ dev = [ { name = "ruff" }, ] mcp = [ - { name = "mcp" }, + { name = "mcp", extra = ["cli"] }, ] [package.metadata] requires-dist = [ { name = "httpx", specifier = ">=0.27" }, - { name = "mcp", marker = "extra == 'mcp'" }, + { name = "mcp", extras = ["cli"], marker = "extra == 'mcp'", specifier = ">=2.0" }, { name = "numpy", specifier = ">=2.0" }, { name = "pandas", specifier = ">=2.2" }, { name = "pydantic", specifier = ">=2.8" }, @@ -411,6 +411,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/67/72/7d7897418912c1d12e87556630dfb7bf0eac71160e9bef8b447960804ee3/mcp-2.0.0-py3-none-any.whl", hash = "sha256:1cb4c75d2d2c7b8c1d756355e5d82a39f2822cc7f13e22a2051d7ca3592349d6", size = 349980, upload-time = "2026-07-28T13:45:28.853Z" }, ] +[package.optional-dependencies] +cli = [ + { name = "python-dotenv" }, + { name = "typer" }, +] + [[package]] name = "mcp-types" version = "2.0.0" @@ -732,6 +738,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, ] +[[package]] +name = "python-dotenv" +version = "1.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/53/ed9d74092561d4b01a2ef1349d52cdbc135e526c245f366b089cfca6de49/python_dotenv-1.2.3.tar.gz", hash = "sha256:a20a594dabeaa385725aa239d5244871c143ecb356add8a20fcf23773a6c3a35", size = 58945, upload-time = "2026-08-16T16:54:54.067Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/17/c5c6b53ddc18f297992099b3d9ec16c855c0ccc83263a21fe4d1c625ec6c/python_dotenv-1.2.3-py3-none-any.whl", hash = "sha256:904552145e8bfed22162c09dab1c2b9b54fefa7b23ba780f4f26ca0316b0f0d9", size = 22780, upload-time = "2026-08-16T16:54:52.473Z" }, +] + [[package]] name = "python-multipart" version = "0.0.32" From a552e3d87d74dc36c3442f5968000ad4a05316bc Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:32:24 +0300 Subject: [PATCH 02/13] feat: support multi-hop propagation analysis --- src/flowsense/engine/propagation.py | 148 ++++++++++++++++++---- tests/test_propagation.py | 186 ++++++++++++++++++++++++++++ 2 files changed, 309 insertions(+), 25 deletions(-) diff --git a/src/flowsense/engine/propagation.py b/src/flowsense/engine/propagation.py index c5382af..4b292a6 100644 --- a/src/flowsense/engine/propagation.py +++ b/src/flowsense/engine/propagation.py @@ -20,49 +20,147 @@ class PropagationResult: propagation_score: float +def _build_reverse_dependencies( + dependencies: dict[str, list[str]], +) -> dict[str, list[str]]: + reverse_dependencies: dict[str, list[str]] = { + task_id: [] for task_id in dependencies + } + + for upstream_task, downstream_tasks in dependencies.items(): + for downstream_task in downstream_tasks: + reverse_dependencies.setdefault(downstream_task, []).append(upstream_task) + + return reverse_dependencies + + +def _has_anomalous_upstream( + task_id: str, + drift_results: dict[str, DriftResult], + reverse_dependencies: dict[str, list[str]], + visited: set[str] | None = None, +) -> bool: + if visited is None: + visited = set() + + if task_id in visited: + return False + + visited = {*visited, task_id} + + for upstream_task in reverse_dependencies.get(task_id, []): + upstream_drift = drift_results.get(upstream_task) + + if upstream_drift is not None and upstream_drift.severity in { + "HIGH", + "CRITICAL", + }: + return True + + if _has_anomalous_upstream( + task_id=upstream_task, + drift_results=drift_results, + reverse_dependencies=reverse_dependencies, + visited=visited, + ): + return True + + return False + + +def _find_propagation_paths( + origin_task: str, + current_task: str, + drift_results: dict[str, DriftResult], + dependencies: dict[str, list[str]], + path: list[str], + visited: set[str], +) -> list[list[str]]: + paths: list[list[str]] = [] + + for downstream_task in dependencies.get(current_task, []): + if downstream_task in visited: + continue + + downstream_drift = drift_results.get(downstream_task) + + if downstream_drift is None: + continue + + if downstream_drift.severity == "NORMAL": + continue + + next_path = [*path, downstream_task] + next_visited = {*visited, downstream_task} + + child_paths = _find_propagation_paths( + origin_task=origin_task, + current_task=downstream_task, + drift_results=drift_results, + dependencies=dependencies, + path=next_path, + visited=next_visited, + ) + + if child_paths: + paths.extend(child_paths) + else: + paths.append(next_path) + + return paths + + def analyze_propagation( drift_results: dict[str, DriftResult], dependencies: dict[str, list[str]], ) -> list[PropagationResult]: results: list[PropagationResult] = [] + reverse_dependencies = _build_reverse_dependencies(dependencies) + for task_id, drift in drift_results.items(): if drift.severity not in {"HIGH", "CRITICAL"}: continue - downstream = dependencies.get(task_id, []) + if _has_anomalous_upstream( + task_id=task_id, + drift_results=drift_results, + reverse_dependencies=reverse_dependencies, + ): + continue - affected: list[str] = [] + paths = _find_propagation_paths( + origin_task=task_id, + current_task=task_id, + drift_results=drift_results, + dependencies=dependencies, + path=[task_id], + visited={task_id}, + ) - for downstream_task in downstream: - downstream_drift = drift_results.get(downstream_task) + for path in paths: + affected_tasks = path[1:] - if downstream_drift is None: + if not affected_tasks: continue - if downstream_drift.severity != "NORMAL": - affected.append(downstream_task) - - if not affected: - continue + origin_score = SEVERITY_SCORE[drift.severity] - origin_score = SEVERITY_SCORE[drift.severity] + downstream_scores = [ + SEVERITY_SCORE[drift_results[task].severity] for task in affected_tasks + ] - downstream_scores = [ - SEVERITY_SCORE[drift_results[task].severity] for task in affected - ] - - propagation_score = sum(downstream_scores) / ( - len(downstream_scores) * origin_score - ) + propagation_score = sum(downstream_scores) / ( + len(downstream_scores) * origin_score + ) - results.append( - PropagationResult( - origin_task=task_id, - affected_tasks=affected, - path=[task_id, *affected], - propagation_score=propagation_score, + results.append( + PropagationResult( + origin_task=task_id, + affected_tasks=affected_tasks, + path=path, + propagation_score=propagation_score, + ) ) - ) return results diff --git a/tests/test_propagation.py b/tests/test_propagation.py index 30bddc5..94dbfcb 100644 --- a/tests/test_propagation.py +++ b/tests/test_propagation.py @@ -122,3 +122,189 @@ def test_medium_task_is_not_treated_as_origin() -> None: ) assert results == [] + + +def test_detects_multi_hop_propagation() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=6.0, + mad=0.5, + robust_z_score=4.0, + deviation_percent=100.0, + severity="HIGH", + ), + "load": DriftResult( + task_id="load", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform"], + "transform": ["load"], + "load": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + extract_result = next( + result for result in results if result.origin_task == "extract" + ) + + assert extract_result.path == [ + "extract", + "transform", + "load", + ] + assert extract_result.affected_tasks == [ + "transform", + "load", + ] + assert extract_result.propagation_score > 0 + + +def test_does_not_report_downstream_task_as_duplicate_origin() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform": DriftResult( + task_id="transform", + baseline=3.0, + current=6.0, + mad=0.5, + robust_z_score=4.0, + deviation_percent=100.0, + severity="HIGH", + ), + "load": DriftResult( + task_id="load", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform"], + "transform": ["load"], + "load": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + origins = [result.origin_task for result in results] + + assert origins == ["extract"] + + +def test_detects_branching_propagation_paths() -> None: + drift_results = { + "extract": DriftResult( + task_id="extract", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "transform_a": DriftResult( + task_id="transform_a", + baseline=2.0, + current=5.0, + mad=0.3, + robust_z_score=4.5, + deviation_percent=150.0, + severity="HIGH", + ), + "load_a": DriftResult( + task_id="load_a", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + "transform_b": DriftResult( + task_id="transform_b", + baseline=2.0, + current=4.0, + mad=0.3, + robust_z_score=3.8, + deviation_percent=100.0, + severity="HIGH", + ), + "load_b": DriftResult( + task_id="load_b", + baseline=1.0, + current=1.6, + mad=0.2, + robust_z_score=2.2, + deviation_percent=60.0, + severity="MEDIUM", + ), + } + + dependencies = { + "extract": ["transform_a", "transform_b"], + "transform_a": ["load_a"], + "load_a": [], + "transform_b": ["load_b"], + "load_b": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + extract_results = [result for result in results if result.origin_task == "extract"] + + paths = [result.path for result in extract_results] + + assert len(extract_results) == 2 + + assert [ + "extract", + "transform_a", + "load_a", + ] in paths + + assert [ + "extract", + "transform_b", + "load_b", + ] in paths From ba6575565c4fee980e698c9e87a6010c29724460 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:33:55 +0300 Subject: [PATCH 03/13] feat: add task handoff timing analysis --- src/flowsense/engine/timing.py | 92 +++++++++++++++++++ tests/test_timing.py | 162 +++++++++++++++++++++++++++++++++ 2 files changed, 254 insertions(+) create mode 100644 src/flowsense/engine/timing.py create mode 100644 tests/test_timing.py diff --git a/src/flowsense/engine/timing.py b/src/flowsense/engine/timing.py new file mode 100644 index 0000000..7c43cac --- /dev/null +++ b/src/flowsense/engine/timing.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult, calculate_drift +from flowsense.models import TaskRun + + +@dataclass +class HandoffTiming: + upstream_task: str + downstream_task: str + dag_run_id: str + handoff_delay: float + + +def calculate_handoff_delay( + upstream_run: TaskRun, + downstream_run: TaskRun, +) -> HandoffTiming: + if upstream_run.end_date is None: + raise ValueError("Upstream task end_date is required.") + + if downstream_run.start_date is None: + raise ValueError("Downstream task start_date is required.") + + if upstream_run.dag_run_id != downstream_run.dag_run_id: + raise ValueError("Task runs must belong to the same DAG run.") + + handoff_delay = (downstream_run.start_date - upstream_run.end_date).total_seconds() + + return HandoffTiming( + upstream_task=upstream_run.task_id, + downstream_task=downstream_run.task_id, + dag_run_id=upstream_run.dag_run_id, + handoff_delay=handoff_delay, + ) + + +def build_handoff_history( + task_runs: list[TaskRun], + dependencies: dict[str, list[str]], +) -> dict[tuple[str, str], list[float]]: + runs_by_id: dict[str, dict[str, TaskRun]] = {} + + for task_run in task_runs: + runs_by_id.setdefault(task_run.dag_run_id, {})[task_run.task_id] = task_run + + history: dict[tuple[str, str], list[float]] = {} + + for tasks_in_run in runs_by_id.values(): + for upstream_task, downstream_tasks in dependencies.items(): + upstream_run = tasks_in_run.get(upstream_task) + + if upstream_run is None: + continue + + for downstream_task in downstream_tasks: + downstream_run = tasks_in_run.get(downstream_task) + + if downstream_run is None: + continue + + try: + timing = calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + except ValueError: + continue + + edge = ( + upstream_task, + downstream_task, + ) + + history.setdefault(edge, []).append(timing.handoff_delay) + + return history + + +def calculate_handoff_drift( + upstream_task: str, + downstream_task: str, + handoff_delays: list[float], +) -> DriftResult: + edge_id = f"{upstream_task}->{downstream_task}" + + return calculate_drift( + task_id=edge_id, + durations=handoff_delays, + ) diff --git a/tests/test_timing.py b/tests/test_timing.py new file mode 100644 index 0000000..c694854 --- /dev/null +++ b/tests/test_timing.py @@ -0,0 +1,162 @@ +from datetime import UTC, datetime, timedelta + +import pytest + +from flowsense.engine.timing import ( + build_handoff_history, + calculate_handoff_delay, + calculate_handoff_drift, +) +from flowsense.models import TaskRun + + +def test_calculate_handoff_delay() -> None: + upstream_end = datetime(2026, 8, 20, 10, 0, 5, tzinfo=UTC) + + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + state="success", + end_date=upstream_end, + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + state="success", + start_date=upstream_end + timedelta(seconds=7), + ) + + result = calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + assert result.upstream_task == "extract" + assert result.downstream_task == "transform" + assert result.dag_run_id == "run_1" + assert result.handoff_delay == 7.0 + + +def test_calculate_handoff_delay_requires_upstream_end_date() -> None: + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="transform", + start_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ) + + with pytest.raises(ValueError): + calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + +def test_calculate_handoff_delay_requires_same_dag_run() -> None: + upstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_1", + task_id="extract", + end_date=datetime(2026, 8, 20, 10, 0, tzinfo=UTC), + ) + + downstream_run = TaskRun( + dag_id="demo", + dag_run_id="run_2", + task_id="transform", + start_date=datetime(2026, 8, 20, 10, 0, 5, tzinfo=UTC), + ) + + with pytest.raises(ValueError): + calculate_handoff_delay( + upstream_run=upstream_run, + downstream_run=downstream_run, + ) + + +def test_build_handoff_history() -> None: + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + task_runs = [] + + for index, delay in enumerate( + [2.0, 3.0, 4.0], + start=1, + ): + dag_run_id = f"run_{index}" + + upstream_end = base_time + timedelta( + minutes=index, + ) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + end_date=upstream_end, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=upstream_end + timedelta(seconds=delay), + ), + ] + ) + + dependencies = { + "extract": ["transform"], + "transform": [], + } + + history = build_handoff_history( + task_runs=task_runs, + dependencies=dependencies, + ) + + assert history == { + ("extract", "transform"): [ + 2.0, + 3.0, + 4.0, + ] + } + + +def test_calculate_handoff_drift_detects_anomaly() -> None: + result = calculate_handoff_drift( + upstream_task="extract", + downstream_task="transform", + handoff_delays=[ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ], + ) + + assert result.task_id == "extract->transform" + assert result.severity == "CRITICAL" + assert result.current == 8.0 + assert result.baseline == 2.0 From e917241c6202102feca21044168fc479bb892ee7 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:34:02 +0300 Subject: [PATCH 04/13] feat: add task impact classification --- src/flowsense/engine/impact.py | 67 +++++++++++++++ tests/test_impact.py | 149 +++++++++++++++++++++++++++++++++ 2 files changed, 216 insertions(+) create mode 100644 src/flowsense/engine/impact.py create mode 100644 tests/test_impact.py diff --git a/src/flowsense/engine/impact.py b/src/flowsense/engine/impact.py new file mode 100644 index 0000000..fd69ca0 --- /dev/null +++ b/src/flowsense/engine/impact.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult + + +@dataclass +class TaskImpact: + task_id: str + classification: str + task_severity: str + upstream_handoff_severity: str | None + + +def classify_task_impact( + task_id: str, + task_drift: DriftResult, + upstream_handoff_drifts: list[DriftResult], +) -> TaskImpact: + anomalous_handoffs = [ + drift + for drift in upstream_handoff_drifts + if drift.severity + in { + "MEDIUM", + "HIGH", + "CRITICAL", + } + ] + + task_is_anomalous = task_drift.severity in { + "MEDIUM", + "HIGH", + "CRITICAL", + } + + if task_is_anomalous and anomalous_handoffs: + classification = "COMBINED" + elif task_is_anomalous: + classification = "OWN_DRIFT" + elif anomalous_handoffs: + classification = "INHERITED_DELAY" + else: + classification = "NORMAL" + + upstream_handoff_severity = None + + if anomalous_handoffs: + severity_order = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, + } + + upstream_handoff_severity = max( + anomalous_handoffs, + key=lambda result: severity_order[result.severity], + ).severity + + return TaskImpact( + task_id=task_id, + classification=classification, + task_severity=task_drift.severity, + upstream_handoff_severity=upstream_handoff_severity, + ) diff --git a/tests/test_impact.py b/tests/test_impact.py new file mode 100644 index 0000000..e3b1b36 --- /dev/null +++ b/tests/test_impact.py @@ -0,0 +1,149 @@ +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import classify_task_impact + + +def test_classifies_inherited_delay() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=3.1, + mad=0.2, + robust_z_score=0.3, + deviation_percent=3.3, + severity="NORMAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=8.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.task_id == "transform" + assert result.classification == "INHERITED_DELAY" + assert result.task_severity == "NORMAL" + assert result.upstream_handoff_severity == "CRITICAL" + + +def test_classifies_own_drift() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=9.0, + mad=0.3, + robust_z_score=6.5, + deviation_percent=200.0, + severity="CRITICAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=2.1, + mad=0.2, + robust_z_score=0.3, + deviation_percent=5.0, + severity="NORMAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.classification == "OWN_DRIFT" + assert result.task_severity == "CRITICAL" + assert result.upstream_handoff_severity is None + + +def test_classifies_normal_task() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=3.1, + mad=0.2, + robust_z_score=0.2, + deviation_percent=3.3, + severity="NORMAL", + ) + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=[], + ) + + assert result.classification == "NORMAL" + assert result.task_severity == "NORMAL" + assert result.upstream_handoff_severity is None + + +def test_classifies_combined_impact() -> None: + task_drift = DriftResult( + task_id="transform", + baseline=3.0, + current=8.0, + mad=0.3, + robust_z_score=5.5, + deviation_percent=166.7, + severity="CRITICAL", + ) + + upstream_handoff_drifts = [ + DriftResult( + task_id="extract->transform", + baseline=2.0, + current=7.0, + mad=0.2, + robust_z_score=5.0, + deviation_percent=250.0, + severity="CRITICAL", + ) + ] + + result = classify_task_impact( + task_id="transform", + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + + assert result.classification == "COMBINED" + assert result.task_severity == "CRITICAL" + assert result.upstream_handoff_severity == "CRITICAL" + + +def test_classifies_medium_task_as_own_drift() -> None: + task_drift = DriftResult( + task_id="load", + baseline=1.4, + current=2.1, + mad=0.2, + robust_z_score=3.17, + deviation_percent=50.0, + severity="MEDIUM", + ) + + result = classify_task_impact( + task_id="load", + task_drift=task_drift, + upstream_handoff_drifts=[], + ) + + assert result.classification == "OWN_DRIFT" + assert result.task_severity == "MEDIUM" + assert result.upstream_handoff_severity is None From 860c89ceea276e6e2b6152ac93e8be3b6130ef44 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:34:06 +0300 Subject: [PATCH 05/13] feat: add primary root cause selection --- src/flowsense/engine/root_cause.py | 152 +++++++++++++++++++ tests/test_root_cause.py | 232 +++++++++++++++++++++++++++++ 2 files changed, 384 insertions(+) create mode 100644 src/flowsense/engine/root_cause.py create mode 100644 tests/test_root_cause.py diff --git a/src/flowsense/engine/root_cause.py b/src/flowsense/engine/root_cause.py new file mode 100644 index 0000000..4e2e03f --- /dev/null +++ b/src/flowsense/engine/root_cause.py @@ -0,0 +1,152 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.propagation import PropagationResult + + +@dataclass +class RootCauseResult: + task_id: str + classification: str + severity: str + propagation_score: float + + +SEVERITY_SCORE = { + "NORMAL": 0, + "MEDIUM": 1, + "HIGH": 2, + "CRITICAL": 3, +} + + +def _build_reverse_dependencies( + dependencies: dict[str, list[str]], +) -> dict[str, list[str]]: + reverse_dependencies: dict[str, list[str]] = {} + + for upstream_task, downstream_tasks in dependencies.items(): + reverse_dependencies.setdefault(upstream_task, []) + + for downstream_task in downstream_tasks: + reverse_dependencies.setdefault( + downstream_task, + [], + ).append(upstream_task) + + return reverse_dependencies + + +def _has_candidate_upstream( + task_id: str, + candidate_tasks: set[str], + reverse_dependencies: dict[str, list[str]], + visited: set[str] | None = None, +) -> bool: + if visited is None: + visited = set() + + if task_id in visited: + return False + + visited.add(task_id) + + for upstream_task in reverse_dependencies.get(task_id, []): + if upstream_task in candidate_tasks: + return True + + if _has_candidate_upstream( + task_id=upstream_task, + candidate_tasks=candidate_tasks, + reverse_dependencies=reverse_dependencies, + visited=visited, + ): + return True + + return False + + +def select_primary_origin( + drift_results: dict[str, DriftResult], + task_impacts: dict[str, TaskImpact], + dependencies: dict[str, list[str]], + propagation_results: list[PropagationResult], +) -> RootCauseResult | None: + candidate_tasks = { + task_id + for task_id, impact in task_impacts.items() + if impact.classification + in { + "OWN_DRIFT", + "COMBINED", + } + } + + if not candidate_tasks: + return None + + reverse_dependencies = _build_reverse_dependencies( + dependencies, + ) + + root_candidates = { + task_id + for task_id in candidate_tasks + if not _has_candidate_upstream( + task_id=task_id, + candidate_tasks=candidate_tasks, + reverse_dependencies=reverse_dependencies, + ) + } + + if not root_candidates: + root_candidates = candidate_tasks + + propagation_scores: dict[str, float] = {} + + for result in propagation_results: + current_score = propagation_scores.get( + result.origin_task, + 0.0, + ) + + propagation_scores[result.origin_task] = max( + current_score, + result.propagation_score, + ) + + def ranking(task_id: str) -> tuple[float, int]: + propagation_score = propagation_scores.get( + task_id, + 0.0, + ) + + drift = drift_results.get(task_id) + + severity_score = SEVERITY_SCORE[drift.severity] if drift else 0 + + return ( + propagation_score, + severity_score, + ) + + primary_task = max( + root_candidates, + key=ranking, + ) + + impact = task_impacts[primary_task] + drift = drift_results[primary_task] + + return RootCauseResult( + task_id=primary_task, + classification=impact.classification, + severity=drift.severity, + propagation_score=propagation_scores.get( + primary_task, + 0.0, + ), + ) diff --git a/tests/test_root_cause.py b/tests/test_root_cause.py new file mode 100644 index 0000000..7b2faa7 --- /dev/null +++ b/tests/test_root_cause.py @@ -0,0 +1,232 @@ +from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.propagation import PropagationResult +from flowsense.engine.root_cause import select_primary_origin + + +def _drift( + task_id: str, + severity: str, +) -> DriftResult: + return DriftResult( + task_id=task_id, + baseline=1.0, + current=2.0, + mad=0.1, + robust_z_score=3.0, + deviation_percent=100.0, + severity=severity, + ) + + +def _impact( + task_id: str, + classification: str, + severity: str, +) -> TaskImpact: + return TaskImpact( + task_id=task_id, + classification=classification, + task_severity=severity, + upstream_handoff_severity=None, + ) + + +def test_selects_single_own_drift_as_primary_origin() -> None: + drift_results = { + "transform": _drift( + "transform", + "CRITICAL", + ), + } + + task_impacts = { + "transform": _impact( + "transform", + "OWN_DRIFT", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "transform": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "transform" + + +def test_inherited_delay_is_not_primary_origin() -> None: + drift_results = { + "transform": _drift( + "transform", + "NORMAL", + ), + } + + task_impacts = { + "transform": _impact( + "transform", + "INHERITED_DELAY", + "NORMAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "transform": [], + }, + propagation_results=[], + ) + + assert result is None + + +def test_prefers_upstream_root_over_downstream_combined_task() -> None: + drift_results = { + "extract": _drift( + "extract", + "CRITICAL", + ), + "transform": _drift( + "transform", + "CRITICAL", + ), + } + + task_impacts = { + "extract": _impact( + "extract", + "OWN_DRIFT", + "CRITICAL", + ), + "transform": _impact( + "transform", + "COMBINED", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract": ["transform"], + "transform": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "extract" + + +def test_selects_more_severe_independent_root() -> None: + drift_results = { + "extract_a": _drift( + "extract_a", + "MEDIUM", + ), + "extract_b": _drift( + "extract_b", + "CRITICAL", + ), + } + + task_impacts = { + "extract_a": _impact( + "extract_a", + "OWN_DRIFT", + "MEDIUM", + ), + "extract_b": _impact( + "extract_b", + "OWN_DRIFT", + "CRITICAL", + ), + } + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract_a": [], + "extract_b": [], + }, + propagation_results=[], + ) + + assert result is not None + assert result.task_id == "extract_b" + + +def test_prefers_higher_propagation_score_between_independent_roots() -> None: + drift_results = { + "extract_a": _drift( + "extract_a", + "CRITICAL", + ), + "extract_b": _drift( + "extract_b", + "CRITICAL", + ), + } + + task_impacts = { + "extract_a": _impact( + "extract_a", + "OWN_DRIFT", + "CRITICAL", + ), + "extract_b": _impact( + "extract_b", + "OWN_DRIFT", + "CRITICAL", + ), + } + + propagation_results = [ + PropagationResult( + origin_task="extract_a", + affected_tasks=["load_a"], + path=["extract_a", "load_a"], + propagation_score=0.4, + ), + PropagationResult( + origin_task="extract_b", + affected_tasks=[ + "transform_b", + "load_b", + ], + path=[ + "extract_b", + "transform_b", + "load_b", + ], + propagation_score=0.8, + ), + ] + + result = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies={ + "extract_a": ["load_a"], + "load_a": [], + "extract_b": ["transform_b"], + "transform_b": ["load_b"], + "load_b": [], + }, + propagation_results=propagation_results, + ) + + assert result is not None + assert result.task_id == "extract_b" + assert result.propagation_score == 0.8 From 3c12c7db7439a1ce038e3c5190e8d6862b527419 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:34:11 +0300 Subject: [PATCH 06/13] feat: integrate timing and root cause analysis --- src/flowsense/engine/analyzer.py | 64 ++++++++++- src/flowsense/models/dag_analysis.py | 6 +- tests/test_analyzer.py | 163 ++++++++++++++++++++++++++- 3 files changed, 225 insertions(+), 8 deletions(-) diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index 4ed9903..61d83da 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -3,7 +3,13 @@ from flowsense.collector.airflow_client import AirflowClient from flowsense.engine.drift import calculate_drift from flowsense.engine.history import build_duration_history +from flowsense.engine.impact import classify_task_impact from flowsense.engine.propagation import analyze_propagation +from flowsense.engine.root_cause import select_primary_origin +from flowsense.engine.timing import ( + build_handoff_history, + calculate_handoff_drift, +) from flowsense.models import DAGAnalysis @@ -26,6 +32,43 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: dependencies = client.get_dag_dependencies(dag_id) + handoff_history = build_handoff_history( + task_runs=task_runs, + dependencies=dependencies, + ) + + handoff_drift_results = {} + + for edge, delays in handoff_history.items(): + _upstream_task, downstream_task = edge + + try: + handoff_drift_results[edge] = calculate_handoff_drift( + upstream_task=_upstream_task, + downstream_task=downstream_task, + handoff_delays=delays, + ) + except ValueError: + continue + + task_impacts = {} + + for task_id, task_drift in drift_results.items(): + upstream_handoff_drifts = [ + drift + for ( + _upstream_task, + downstream_task, + ), drift in handoff_drift_results.items() + if downstream_task == task_id + ] + + task_impacts[task_id] = classify_task_impact( + task_id=task_id, + task_drift=task_drift, + upstream_handoff_drifts=upstream_handoff_drifts, + ) + propagation_results = analyze_propagation( drift_results=drift_results, dependencies=dependencies, @@ -40,19 +83,26 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: overall_severity = "NORMAL" - if drift_results: + all_drift_results = [ + *drift_results.values(), + *handoff_drift_results.values(), + ] + + if all_drift_results: overall_severity = max( - drift_results.values(), + all_drift_results, key=lambda result: severity_order[result.severity], ).severity primary_origin = None if propagation_results: - primary_origin = max( - propagation_results, - key=lambda result: result.propagation_score, - ).origin_task + primary_origin = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies=dependencies, + propagation_results=propagation_results, + ) return DAGAnalysis( dag_id=dag_id, @@ -60,6 +110,8 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: overall_severity=overall_severity, primary_origin=primary_origin, drift_results=drift_results, + handoff_drift_results=handoff_drift_results, + task_impacts=task_impacts, propagation_results=propagation_results, dependencies=dependencies, ) diff --git a/src/flowsense/models/dag_analysis.py b/src/flowsense/models/dag_analysis.py index b735d81..e963ffb 100644 --- a/src/flowsense/models/dag_analysis.py +++ b/src/flowsense/models/dag_analysis.py @@ -3,7 +3,9 @@ from dataclasses import dataclass from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact from flowsense.engine.propagation import PropagationResult +from flowsense.engine.root_cause import RootCauseResult @dataclass @@ -11,7 +13,9 @@ class DAGAnalysis: dag_id: str runs_analyzed: int overall_severity: str - primary_origin: str | None + primary_origin: RootCauseResult | None drift_results: dict[str, DriftResult] + handoff_drift_results: dict[tuple[str, str], DriftResult] + task_impacts: dict[str, TaskImpact] propagation_results: list[PropagationResult] dependencies: dict[str, list[str]] diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index a462261..a11e122 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -1,3 +1,4 @@ +from datetime import UTC, datetime, timedelta from unittest.mock import MagicMock, patch from flowsense.engine.analyzer import analyze_dag @@ -44,7 +45,167 @@ def test_analyze_dag_identifies_primary_origin( analysis = analyze_dag("demo") assert analysis.overall_severity == "CRITICAL" - assert analysis.primary_origin == "transform" + assert analysis.primary_origin is not None + assert analysis.primary_origin.task_id == "transform" assert len(analysis.propagation_results) == 1 assert analysis.propagation_results[0].origin_task == "transform" assert analysis.propagation_results[0].affected_tasks == ["load"] + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_calculates_handoff_drift( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + handoff_delays = [ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ] + + task_runs = [] + + for index, delay in enumerate( + handoff_delays, + start=1, + ): + dag_run_id = f"run_{index}" + + extract_start = base_time + timedelta(minutes=index) + extract_end = extract_start + timedelta(seconds=1) + + transform_start = extract_end + timedelta(seconds=delay) + transform_end = transform_start + timedelta(seconds=3) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + start_date=extract_start, + end_date=extract_end, + duration=1.0, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=transform_start, + end_date=transform_end, + duration=3.0, + ), + ] + ) + + client.collect_task_runs.return_value = task_runs + + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + edge = ("extract", "transform") + + assert edge in analysis.handoff_drift_results + + handoff_drift = analysis.handoff_drift_results[edge] + + assert handoff_drift.current == 8.0 + assert handoff_drift.baseline == 2.0 + assert handoff_drift.severity == "CRITICAL" + + +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_uses_handoff_for_overall_severity( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + base_time = datetime( + 2026, + 8, + 20, + 10, + 0, + tzinfo=UTC, + ) + + handoff_delays = [ + 2.0, + 2.1, + 1.9, + 2.0, + 8.0, + ] + + task_runs = [] + + for index, delay in enumerate( + handoff_delays, + start=1, + ): + dag_run_id = f"run_{index}" + + extract_start = base_time + timedelta(minutes=index) + extract_end = extract_start + timedelta(seconds=1) + + transform_start = extract_end + timedelta(seconds=delay) + transform_end = transform_start + timedelta(seconds=3) + + task_runs.extend( + [ + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="extract", + state="success", + start_date=extract_start, + end_date=extract_end, + duration=1.0, + ), + TaskRun( + dag_id="demo", + dag_run_id=dag_run_id, + task_id="transform", + state="success", + start_date=transform_start, + end_date=transform_end, + duration=3.0, + ), + ] + ) + + client.collect_task_runs.return_value = task_runs + + client.get_dag_dependencies.return_value = { + "extract": ["transform"], + "transform": [], + } + + analysis = analyze_dag("demo") + + assert analysis.drift_results["extract"].severity == "NORMAL" + assert analysis.drift_results["transform"].severity == "NORMAL" + + assert ( + analysis.handoff_drift_results[("extract", "transform")].severity == "CRITICAL" + ) + + assert analysis.overall_severity == "CRITICAL" From f5cb87e6549fab86425b5ff069b76cd2b47db7b6 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:34:15 +0300 Subject: [PATCH 07/13] feat: display task impact and root cause details in CLI --- src/flowsense/cli/main.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/flowsense/cli/main.py b/src/flowsense/cli/main.py index ee01015..33d5604 100644 --- a/src/flowsense/cli/main.py +++ b/src/flowsense/cli/main.py @@ -39,8 +39,12 @@ def analyze( table.add_column("Deviation") table.add_column("Z-Score") table.add_column("Severity") + table.add_column("Impact") for task_id, result in analysis.drift_results.items(): + impact = analysis.task_impacts.get(task_id) + impact_label = impact.classification if impact else "-" + table.add_row( task_id, f"{result.baseline:.2f}s", @@ -48,6 +52,7 @@ def analyze( f"{result.deviation_percent:+.1f}%", f"{result.robust_z_score:.2f}", result.severity, + impact_label, ) console.print(table) @@ -55,7 +60,12 @@ def analyze( console.print(f"\nOverall Severity: [bold]{analysis.overall_severity}[/bold]") if analysis.primary_origin: - console.print(f"Primary Origin: [bold]{analysis.primary_origin}[/bold]") + console.print(f"Primary Origin: [bold]{analysis.primary_origin.task_id}[/bold]") + console.print(f"Reason: [bold]{analysis.primary_origin.classification}[/bold]") + console.print(f"Severity: [bold]{analysis.primary_origin.severity}[/bold]") + console.print( + f"Propagation Score: {analysis.primary_origin.propagation_score:.2f}" + ) if analysis.propagation_results: console.print("\n[bold]Propagation Analysis[/bold]\n") From 069ed2edf72774732370077dc8ea9a9c71b4e2ff Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:34:18 +0300 Subject: [PATCH 08/13] feat: expose impact and root cause analysis through MCP --- src/flowsense/mcp/server.py | 33 +++++++++++++++++++- tests/test_mcp_integration.py | 55 ++++++++++++++++++++++++++++++++ tests/test_mcp_server.py | 59 +++++++++++++++++++++++++++++++++-- 3 files changed, 144 insertions(+), 3 deletions(-) create mode 100644 tests/test_mcp_integration.py diff --git a/src/flowsense/mcp/server.py b/src/flowsense/mcp/server.py index 43b3729..af8d483 100644 --- a/src/flowsense/mcp/server.py +++ b/src/flowsense/mcp/server.py @@ -13,7 +13,16 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: "dag_id": analysis.dag_id, "runs_analyzed": analysis.runs_analyzed, "overall_severity": analysis.overall_severity, - "primary_origin": analysis.primary_origin, + "primary_origin": ( + { + "task_id": analysis.primary_origin.task_id, + "classification": analysis.primary_origin.classification, + "severity": analysis.primary_origin.severity, + "propagation_score": analysis.primary_origin.propagation_score, + } + if analysis.primary_origin + else None + ), "drift_results": { task_id: { "baseline": result.baseline, @@ -25,6 +34,28 @@ def serialize_analysis(analysis: DAGAnalysis) -> dict: } for task_id, result in analysis.drift_results.items() }, + "handoff_drift_results": { + f"{upstream}->{downstream}": { + "baseline": result.baseline, + "current": result.current, + "mad": result.mad, + "robust_z_score": result.robust_z_score, + "deviation_percent": result.deviation_percent, + "severity": result.severity, + } + for ( + upstream, + downstream, + ), result in analysis.handoff_drift_results.items() + }, + "task_impacts": { + task_id: { + "classification": impact.classification, + "task_severity": impact.task_severity, + "upstream_handoff_severity": impact.upstream_handoff_severity, + } + for task_id, impact in analysis.task_impacts.items() + }, "propagation_results": [ { "origin_task": result.origin_task, diff --git a/tests/test_mcp_integration.py b/tests/test_mcp_integration.py new file mode 100644 index 0000000..979d6c4 --- /dev/null +++ b/tests/test_mcp_integration.py @@ -0,0 +1,55 @@ +import os + +import pytest +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +@pytest.mark.integration +@pytest.mark.anyio +async def test_mcp_analyze_airflow_dag_end_to_end() -> None: + dag_id = os.getenv( + "FLOWSENSE_TEST_DAG_ID", + "flowsense_demo", + ) + + env = os.environ.copy() + + existing_pythonpath = env.get("PYTHONPATH") + + if existing_pythonpath: + env["PYTHONPATH"] = f"src:{existing_pythonpath}" + else: + env["PYTHONPATH"] = "src" + + server_params = StdioServerParameters( + command="uv", + args=[ + "run", + "python", + "-m", + "flowsense.mcp.server", + ], + env=env, + ) + + async with ( + stdio_client(server_params) as (read, write), + ClientSession(read, write) as session, + ): + await session.initialize() + + result = await session.call_tool( + "analyze_airflow_dag", + { + "dag_id": dag_id, + }, + ) + + assert result.is_error is False + assert result.content diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 6d101da..4ab3ccd 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,8 +1,12 @@ +import os + import pytest from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from flowsense.engine.drift import DriftResult +from flowsense.engine.impact import TaskImpact +from flowsense.engine.root_cause import RootCauseResult from flowsense.mcp.server import serialize_analysis from flowsense.models import DAGAnalysis @@ -14,6 +18,15 @@ def anyio_backend() -> str: @pytest.mark.anyio async def test_mcp_server_exposes_analyze_tool() -> None: + env = os.environ.copy() + + existing_pythonpath = env.get("PYTHONPATH") + + if existing_pythonpath: + env["PYTHONPATH"] = f"src:{existing_pythonpath}" + else: + env["PYTHONPATH"] = "src" + server_params = StdioServerParameters( command="uv", args=[ @@ -22,6 +35,7 @@ async def test_mcp_server_exposes_analyze_tool() -> None: "-m", "flowsense.mcp.server", ], + env=env, ) async with ( @@ -42,7 +56,12 @@ def test_serialize_analysis() -> None: dag_id="demo", runs_analyzed=5, overall_severity="CRITICAL", - primary_origin="transform", + primary_origin=RootCauseResult( + task_id="transform", + classification="COMBINED", + severity="CRITICAL", + propagation_score=0.8, + ), drift_results={ "transform": DriftResult( task_id="transform", @@ -54,6 +73,25 @@ def test_serialize_analysis() -> None: severity="CRITICAL", ) }, + handoff_drift_results={ + ("extract", "transform"): DriftResult( + task_id="extract->transform", + baseline=2.0, + current=8.0, + mad=0.2, + robust_z_score=20.24, + deviation_percent=300.0, + severity="CRITICAL", + ) + }, + task_impacts={ + "transform": TaskImpact( + task_id="transform", + classification="COMBINED", + task_severity="CRITICAL", + upstream_handoff_severity="CRITICAL", + ) + }, propagation_results=[], dependencies={ "transform": [], @@ -65,8 +103,13 @@ def test_serialize_analysis() -> None: assert result["dag_id"] == "demo" assert result["runs_analyzed"] == 5 assert result["overall_severity"] == "CRITICAL" - assert result["primary_origin"] == "transform" + primary_origin = result["primary_origin"] + + assert primary_origin["task_id"] == "transform" + assert primary_origin["classification"] == "COMBINED" + assert primary_origin["severity"] == "CRITICAL" + assert primary_origin["propagation_score"] == 0.8 transform = result["drift_results"]["transform"] assert transform["baseline"] == 3.0 @@ -76,6 +119,18 @@ def test_serialize_analysis() -> None: assert transform["deviation_percent"] == 216.67 assert transform["severity"] == "CRITICAL" + handoff = result["handoff_drift_results"]["extract->transform"] + + assert handoff["baseline"] == 2.0 + assert handoff["current"] == 8.0 + assert handoff["severity"] == "CRITICAL" + + impact = result["task_impacts"]["transform"] + + assert impact["classification"] == "COMBINED" + assert impact["task_severity"] == "CRITICAL" + assert impact["upstream_handoff_severity"] == "CRITICAL" + assert result["propagation_results"] == [] assert result["dependencies"] == { "transform": [], From 6e3d375d96e5db7f9f4a0f70c4705596d75543bd Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:35:33 +0300 Subject: [PATCH 09/13] ci: install MCP dependencies for tests --- .github/workflows/ci.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 282a8f0..6a0bdca 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -23,7 +23,7 @@ jobs: uses: astral-sh/setup-uv@v6 - name: Install dependencies - run: uv pip install --system -e ".[dev]" + run: uv pip install --system -e ".[dev,mcp]" - name: Run Ruff run: ruff check . From 79f6baa7f36727be3246c25fefe3a3895e78b6ce Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:39:29 +0300 Subject: [PATCH 10/13] fix: use current Python environment for MCP tests --- tests/test_mcp_integration.py | 5 ++--- tests/test_mcp_server.py | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/test_mcp_integration.py b/tests/test_mcp_integration.py index 979d6c4..f045671 100644 --- a/tests/test_mcp_integration.py +++ b/tests/test_mcp_integration.py @@ -1,4 +1,5 @@ import os +import sys import pytest from mcp import ClientSession, StdioServerParameters @@ -28,10 +29,8 @@ async def test_mcp_analyze_airflow_dag_end_to_end() -> None: env["PYTHONPATH"] = "src" server_params = StdioServerParameters( - command="uv", + command=sys.executable, args=[ - "run", - "python", "-m", "flowsense.mcp.server", ], diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 4ab3ccd..bbee8f1 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -1,4 +1,5 @@ import os +import sys import pytest from mcp import ClientSession, StdioServerParameters @@ -28,10 +29,8 @@ async def test_mcp_server_exposes_analyze_tool() -> None: env["PYTHONPATH"] = "src" server_params = StdioServerParameters( - command="uv", + command=sys.executable, args=[ - "run", - "python", "-m", "flowsense.mcp.server", ], From b3a56d97427f810c59cda393af0de34460359895 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 25 Aug 2026 23:41:12 +0300 Subject: [PATCH 11/13] docs: update README with analysis and MCP features --- README.md | 59 +++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index f623a2f..b769207 100644 --- a/README.md +++ b/README.md @@ -26,8 +26,12 @@ FlowSense is designed to answer questions such as: - Median-based historical baselines - MAD-based robust Z-score drift detection - Severity classification -- Dependency-aware propagation analysis +- Task handoff delay analysis +- Task impact classification (`OWN_DRIFT`, `INHERITED_DELAY`, and `COMBINED`) +- Multi-hop and branching propagation analysis +- Primary root-cause selection - CLI-based DAG analysis +- MCP server integration ## Example @@ -40,13 +44,19 @@ Example output: ```text FlowSense Analysis — flowsense_demo -┏━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┓ -┃ Task ┃ Baseline ┃ Current ┃ Deviation ┃ Z-Score ┃ Severity ┃ -┡━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━┩ -│ extract │ 1.56s │ 1.61s │ +3.4% │ 0.17 │ NORMAL │ -│ transform │ 3.34s │ 9.61s │ +187.6% │ 7.61 │ CRITICAL │ -│ load │ 1.40s │ 2.11s │ +50.2% │ 3.17 │ MEDIUM │ -└───────────┴──────────┴─────────┴───────────┴─────────┴──────────┘ +┏━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━┓ +┃ Task ┃ Baseline ┃ Current ┃ Deviation ┃ Z-Score ┃ Severity ┃ Impact ┃ +┡━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━┩ +│ extract │ 1.56s │ 1.61s │ +3.4% │ 0.17 │ NORMAL │ NORMAL │ +│ transform │ 3.34s │ 9.61s │ +187.6% │ 7.61 │ CRITICAL │ OWN_DRIFT │ +│ load │ 1.40s │ 2.11s │ +50.2% │ 3.17 │ MEDIUM │ COMBINED │ +└───────────┴──────────┴─────────┴───────────┴─────────┴──────────┴───────────┘ + +Overall Severity: CRITICAL +Primary Origin: transform +Reason: OWN_DRIFT +Severity: CRITICAL +Propagation Score: 0.33 Propagation Analysis @@ -64,17 +74,16 @@ Apache Airflow Collector │ ▼ -Task Run History +Task Run and Handoff History │ ▼ - Drift Engine +Drift and Impact Analysis │ ▼ -Propagation Analysis +Propagation and Root-Cause Analysis │ ├── CLI - │ - └── MCP Server (planned) + └── MCP Server ``` ## Installation @@ -93,7 +102,7 @@ Create a virtual environment and install the project: ```bash uv venv --python 3.12 source .venv/bin/activate -uv pip install -e ".[dev]" +uv pip install -e ".[dev,mcp]" ``` ## Configuration @@ -126,6 +135,18 @@ Then run: flowsense analyze ``` +## MCP Server + +Start the FlowSense MCP server over stdio: + +```bash +flowsense-mcp +``` + +The server exposes the `analyze_airflow_dag` tool, which returns task drift, +handoff drift, impact classification, propagation paths, and primary root-cause +information for a DAG. + ## Development Run unit tests: @@ -167,7 +188,11 @@ src/flowsense/ ├── engine/ │ ├── drift.py │ ├── history.py -│ └── propagation.py +│ ├── impact.py +│ ├── propagation.py +│ ├── root_cause.py +│ └── timing.py +├── mcp/ └── models/ ``` @@ -194,13 +219,9 @@ Planned areas include: - DAG-level analysis models - configurable historical baseline windows - improved propagation scoring -- temporal delay analysis -- upstream/downstream impact separation -- root-cause analysis - change-point detection - trend detection - richer CLI reporting -- MCP server integration - broader Airflow compatibility testing ## License From 112741e7762cac185f7f6062201a80f75ab53dc4 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:36:47 +0300 Subject: [PATCH 12/13] fix: detect primary origin for isolated anomalies --- src/flowsense/engine/analyzer.py | 15 ++++++--------- tests/test_analyzer.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 9 deletions(-) diff --git a/src/flowsense/engine/analyzer.py b/src/flowsense/engine/analyzer.py index 61d83da..aa45c3c 100644 --- a/src/flowsense/engine/analyzer.py +++ b/src/flowsense/engine/analyzer.py @@ -94,15 +94,12 @@ def analyze_dag(dag_id: str) -> DAGAnalysis: key=lambda result: severity_order[result.severity], ).severity - primary_origin = None - - if propagation_results: - primary_origin = select_primary_origin( - drift_results=drift_results, - task_impacts=task_impacts, - dependencies=dependencies, - propagation_results=propagation_results, - ) + primary_origin = select_primary_origin( + drift_results=drift_results, + task_impacts=task_impacts, + dependencies=dependencies, + propagation_results=propagation_results, + ) return DAGAnalysis( dag_id=dag_id, diff --git a/tests/test_analyzer.py b/tests/test_analyzer.py index a11e122..f76d9b8 100644 --- a/tests/test_analyzer.py +++ b/tests/test_analyzer.py @@ -52,6 +52,37 @@ def test_analyze_dag_identifies_primary_origin( assert analysis.propagation_results[0].affected_tasks == ["load"] +@patch("flowsense.engine.analyzer.AirflowClient") +def test_analyze_dag_identifies_isolated_primary_origin( + mock_client_class: MagicMock, +) -> None: + client = mock_client_class.return_value + + client.collect_task_runs.return_value = [ + TaskRun( + dag_id="demo", + dag_run_id=f"run_{index}", + task_id="transform", + state="success", + duration=duration, + ) + for index, duration in enumerate( + [3.0, 3.1, 2.9, 3.0, 9.5], + start=1, + ) + ] + + client.get_dag_dependencies.return_value = { + "transform": [], + } + + analysis = analyze_dag("demo") + + assert analysis.propagation_results == [] + assert analysis.primary_origin is not None + assert analysis.primary_origin.task_id == "transform" + + @patch("flowsense.engine.analyzer.AirflowClient") def test_analyze_dag_calculates_handoff_drift( mock_client_class: MagicMock, From 3b38a37893bf909a91cddb359790017e81aae6f9 Mon Sep 17 00:00:00 2001 From: omercengiz Date: Tue, 1 Sep 2026 22:36:58 +0300 Subject: [PATCH 13/13] fix: preserve anomaly origins across normal dependency gaps --- src/flowsense/engine/propagation.py | 5 ++- tests/test_propagation.py | 57 +++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/src/flowsense/engine/propagation.py b/src/flowsense/engine/propagation.py index 4b292a6..507cdbc 100644 --- a/src/flowsense/engine/propagation.py +++ b/src/flowsense/engine/propagation.py @@ -51,7 +51,10 @@ def _has_anomalous_upstream( for upstream_task in reverse_dependencies.get(task_id, []): upstream_drift = drift_results.get(upstream_task) - if upstream_drift is not None and upstream_drift.severity in { + if upstream_drift is None or upstream_drift.severity == "NORMAL": + continue + + if upstream_drift.severity in { "HIGH", "CRITICAL", }: diff --git a/tests/test_propagation.py b/tests/test_propagation.py index 94dbfcb..c44ac98 100644 --- a/tests/test_propagation.py +++ b/tests/test_propagation.py @@ -308,3 +308,60 @@ def test_detects_branching_propagation_paths() -> None: "transform_b", "load_b", ] in paths + + +def test_preserves_origin_after_normal_dependency_gap() -> None: + drift_results = { + "upstream": DriftResult( + task_id="upstream", + baseline=1.0, + current=4.0, + mad=0.2, + robust_z_score=6.0, + deviation_percent=300.0, + severity="CRITICAL", + ), + "normal_bridge": DriftResult( + task_id="normal_bridge", + baseline=2.0, + current=2.1, + mad=0.3, + robust_z_score=0.2, + deviation_percent=5.0, + severity="NORMAL", + ), + "independent_origin": DriftResult( + task_id="independent_origin", + baseline=3.0, + current=7.0, + mad=0.5, + robust_z_score=5.4, + deviation_percent=133.3, + severity="CRITICAL", + ), + "downstream": DriftResult( + task_id="downstream", + baseline=1.0, + current=1.8, + mad=0.2, + robust_z_score=2.5, + deviation_percent=80.0, + severity="MEDIUM", + ), + } + + dependencies = { + "upstream": ["normal_bridge"], + "normal_bridge": ["independent_origin"], + "independent_origin": ["downstream"], + "downstream": [], + } + + results = analyze_propagation( + drift_results=drift_results, + dependencies=dependencies, + ) + + assert len(results) == 1 + assert results[0].origin_task == "independent_origin" + assert results[0].path == ["independent_origin", "downstream"]