Aegis: DataHub OnCall Agent

Aegis is an AI on-call engineer that monitors DataHub 24/7, investigates issues via real lineage and live data, proposes a fix, waits for approval, then applies it and logs the outcome.

Comment

Inspiration

The scariest data incidents are never the loud ones. A pipeline that throws an error gets fixed in minutes. The dangerous failure is silent — a dashboard quietly showing zero, or stale numbers, for hours or days before a human happens to notice and asks "wait, is this right?" By the time someone catches it, they're starting the investigation from zero: which table, which hop in the pipeline, who owns it, what changed.

We wanted an on-call engineer that doesn't wait to be asked. Something that watches the data the way a human on-call engineer would if they never needed sleep — but that still respects the one thing that actually matters in an incident response system: nothing consequential happens without a human saying yes.

That second half mattered as much as the first. It would have been easy to build something that "does everything" autonomously. We deliberately didn't. The interesting design problem wasn't "can an agent fix things" — it was "where exactly should the human stay in the loop, and how do you enforce that in code instead of just a prompt."

What it does

Aegis is an AI on-call engineer for data pipelines, built on top of DataHub's Analytics Agent framework.

  1. Detects on its own — a deterministic rule engine (zero LLM calls) continuously checks for anomalies: zero revenue, sudden drops against a trailing average, abnormal order volume.
  2. Investigates with real evidence — using DataHub's Agent Context Kit, it walks the lineage graph upstream from the symptom, checks freshness and schema at every hop, and cross-references it against live SQL queries, not just metadata.
  3. Proposes one fix from a fixed catalog — root causes are classified into a small, auditable set of categories, each mapped to a specific, non-improvised remediation.
  4. Waits for a human — nothing executes without explicit approval. A non-admin's approval only queues the finding for review; only an admin clicking approve in the dashboard authorizes real action.
  5. Actually executes the fix — a real data backfill and pipeline rebuild, visible live on the dashboard.
  6. Writes the outcome back to DataHub — but only after real action has occurred, not at proposal time, so the catalog reflects what actually happened.

A second, general-purpose chat assistant ("Copilot") sits alongside it for open-ended questions about the DataHub catalog — using the same underlying agent, but deliberately barred from writing to the incident queue itself, so every tracked finding goes through the same consistent pipeline no matter how it originated.

How we built it

The foundation is DataHub's open-source Analytics Agent — a LangGraph-based chat agent wired to DataHub's Agent Context Kit (datahub-agent-context), giving it real tools for search, lineage traversal, schema inspection, and entity lookup.

On top of that, we added:

  • remediation_tools.py — two new LangChain tools, get_remediation_options (a fixed lookup table mapping anomaly categories to remediations) and record_incident_outcome (the tool that queues findings for review or, for an authorized admin, actually executes the fix and writes back to DataHub).
  • autonomous_monitor.py — a standalone script running a small deterministic rule engine on a timer, invoking the same LangGraph agent headlessly whenever a rule fires — no human types anything.
  • A role-based approval queue — a MySQL table plus dashboard UI where developer-triggered or autonomous findings sit as PENDING until an admin approves or rejects them.
  • A live operational dashboard (FastAPI + a real MySQL pipeline mirroring the DataHub lineage graph) — live KPIs, a Sentinel health-check button, the Copilot chat panel, and an activity feed with full incident reports viewable and exportable as PDF.

Everything reasons through Claude; everything that actually does something — classifying an anomaly type, authorizing execution, running the backfill — is deterministic Python the model cannot talk its way around.

Challenges we ran into

  • Two real bugs in the underlying framework. A custom-patched tool (list_schema_fields) was failing with "No DataHubClient in context" because it skipped the context-wrapping the base framework applies to every other tool. Separately, three built-in tools (list_schema_fields, get_lineage, get_entities) were blowing past DataHub's 15,000-token GraphQL query limit on our non-cloud GMS instance, because their built-in queries pull a much heavier payload than the limit allows. We wrote lightweight, targeted GraphQL queries to replace all three.
  • A nested-event-loop bug that surfaced as a bizarre Invalid URL 'datahub/config' error — calling the framework's synchronous graph-building function from inside an already-running async FastAPI request caused its internal event-loop usage to collide with the outer one. Fixed by offloading it to a thread.
  • A concurrency race on approval. Nothing initially stopped two overlapping "Approve" requests from both running the backfill logic at once — since that logic does DROP TABLE / CREATE TABLE AS SELECT, a second overlapping run could wipe out data the first run had just written. Fixed with an atomic claim (UPDATE ... WHERE status = 'PENDING') so only one request can ever proceed.
  • A subtle escaping bug in the "View Report" feature — encodeURIComponent() doesn't escape apostrophes, and RCA reports are full of them ("today's revenue," "doesn't"), which silently broke the button's onclick handler on almost any real report.
  • Getting the human-in-the-loop boundary actually correct, not just prompted. At one point we tried instructing the autonomous monitor's prompt to self-approve findings, framing it as "safe" since a non-admin role can't trigger real execution anyway. The model correctly refused, flagging the instruction as an unverifiable claim asking it to bypass an approval gate — which was the right call, and told us the queuing decision belonged in deterministic code, not a prompt, all along.

Accomplishments that we're proud of

  • A fully verified, end-to-end loop — autonomous detection through real execution through DataHub write-back — tested live, not just described.
  • A genuinely hybrid architecture: deterministic code owns every consequential decision (is this an anomaly, is execution authorized), the LLM owns reasoning and explanation, and that boundary is enforced in code the model cannot override.
  • Two real, verifiable contributions back to the open-source framework we built on, not just usage of it.
  • A demo that shows a dashboard going from broken to fixed, live, with zero staging or fake data.

What we learned

The most valuable lesson was where not to trust the model. Every time we were tempted to let the LLM decide something consequential — whether to self-approve, how to classify risk — we ended up moving that decision into plain, auditable code instead, and the system got both safer and easier to reason about. Building an agent that's genuinely useful for infrastructure work looks less like "let the model do everything" and more like carefully deciding which five percent of decisions actually need judgment, and keeping the other ninety-five percent boring and deterministic.

We also learned a lot about the practical edges of DataHub's Agent Context Kit on a self-hosted, non-cloud GMS instance - token limits and schema differences that don't show up until you're running against a real instance, not just reading the docs.

What's next for Aegis

  • Real orchestrator integration — trigger an actual Airflow DAG run via its REST API instead of a direct data backfill, so "re-run the pipeline" means exactly what it says against real infrastructure.
  • A learning loop — once a human validates a diagnosis enough times, compile that pattern into a deterministic check that catches it instantly next time, reducing reliance on the LLM for repeat incidents.
  • Real authentication in place of the current UI-level role selector.
  • A broader anomaly rule library and Slack/PagerDuty notifications so approval doesn't require someone to be looking at the dashboard.
  • Real dbt-based lineage — replacing the hand-authored DataHub metadata with lineage derived from actual transformation code, ingested automatically.

Built With

Share this project:

Updates