coSTAR: DatabricksでAIエージェントを迅速かつ安全にリリースする方法

You'd never let a coding assistant refactor your codebase without a test suite. Without tests, the assistant flies blind. It might fix one function and silently break three others. The tests are what close the loop: run them, observe failures, fix the code, run them again. No tests, no confidence.

At Databricks we continuously develop and deploy agents that cover a wide range of functionality, from new features in the Databricks platform (e.g., the data-engineering, trace analysis, and machine learning capabilities in Genie Code), to OSS projects (e.g., the MLflow assistant), to internal engineering workflows (e.g., on-call support or automated code reviewers). These agents can perform long-running tasks, generate thousands of lines of code, and create new data and AI assets among other things. While we had some basic checks in place early on, we lacked the kind of comprehensive, automated test suite that would let us iterate with confidence. This post describes how we closed that gap using MLflow, and the best-practices coSTAR (coupled Scenario, Trace, Assess, Refine) methodology we built around it. coSTAR runs two coupled loops: one that aligns judges with human expert judgment so they can be trusted, and one that uses those trusted judges to automatically refine the agent until it passes all test scenarios.

Figure: The coSTAR framework runs two mirrored STAR loops (Scenario → Trace → Assess → Refine) . The agent loop (blue) uses judges to auto-score traces and refines the agent to align with judges. The judge loop (orange) uses human experts to score traces and refines the judges to align with their assessments. Both loops share the same scenarios and traces.

The Problem: Coding Without Tests

Early on, our development loop looked like this: run the agent, manually review its output, spot a flaw, tell a coding assistant to fix it. Repeat.

If this reminds you of writing code without tests and manually QA-ing every change, that's exactly what it was. And it failed in exactly the way you'd predict. The obvious reaction is "so write tests." But agent testing is structurally different from testing a deterministic function, and several challenges compound at once:

• Non-determinism. The same implementation, the same input, can produce different outputs on different runs. Tests need to evaluate properties of the output rather than assert exact outputs.
• Slow feedback loops. A single agent execution can take tens of minutes. There's no iterating the way a sub-second test suite allows. Every evaluation cycle is expensive.
• Cascading errors. A bad decision at step 3 causes a failure at step 7. By the time the symptom surfaces, the root cause is buried several steps back in the agent's execution.
• Subjective quality. For many testing dimensions (is this feature engineering code any good? is this data cleaning approach appropriate?) there's no ground truth. Judging these dimensions depends on domain expertise.

These constraints shaped every design decision that follows. They're also what makes this problem interesting: we're not just building a test runner, we're building a automated optimization methodology for stochastic, long-running, multi-step processes where "correct" is a judgment call.

The Analogy That Guides Our Approach

If you squint, agent development maps cleanly onto the dev loop that every engineer already knows:

This analogy isn't just illustrative. It's the literal architecture of our system, which we call coSTAR: two coupled loops that use Scenario definitions as test fixtures, Trace capture as the test harness, Assess with judges as the test suite, and Refine as the red-green loop. Let's walk through each piece.

S - Scenario Definitions

In traditional testing, a test fixture sets up the preconditions: create a database, seed it with data, configure the environment. Our equivalent is a scenario definition: a structured description of the initial state, the user prompt, and the expected outcomes.

Here's a simplified scenario for testing a Data Analyst agent against a messy dataset:

Each scenario bundles the setup, the input, and the success criteria in one place, just like a test fixture. We maintain a suite of these across different agents, covering common cases, edge cases, and known past failures. The suite grows over time as we discover new failure modes: every bug we find in production becomes a new scenario, the same way every production bug should become a regression test.

Why bother with this structure? Because agent runs are expensive. A single scenario takes minutes to execute. We need to be deliberate about what we test, and we need the scenario definitions to be portable: the same scenario can run against different agent implementations or different versions of the same agent.

T - Trace Capture

To run our test suite, we use a harness that sends each scenario's prompt to the agent under test (AUT). Each execution is captured as a MLflow trace: a structured log of every tool call, every intermediate output, and every artifact the agent produces. Think of it as a flight recorder: it captures everything the agent did, in order, so we can inspect any part of the execution after the fact.

A key architectural decision: we decouple execution from scoring. The test harness produces traces; the judges (which we'll introduce next) score them. These are separate steps. By persisting traces, we can iterate on judges without re-running scenarios. Adjust a threshold? Re-score the recorded traces in seconds. Add a new judge? Run it against every trace you've ever collected. Suspect a judge is wrong? Compare its verdicts against the recordings and debug it offline. One expensive agent run produces data that gets reused many times, including as candidates for the Golden Set we'll use to align judges later.

A - Assess with Judges

Judges operate on traces and reason about properties of the execution: did the agent produce valid code? Did the output meet a quality threshold? Did the agent follow the right process? As mentioned earlier, this evaluation is different from traditional unit tests: agent output is non-deterministic and rich, and so asserting exact outputs is essentially useless.

The standard approach to implementing these judges is "LLM-as-a-Judge": feed the full trace to a model and ask for a score and equally importantly a rationale for that score. However, that's like writing a test that dumps the entire program state into an assertion. It's expensive, fragile, and hard to debug. For our agents, a single trace can be thousands of lines long. Stuffing it into a judge's context window degrades judgment quality.

Instead, we use MLflow’s agentic judges: judges that are themselves agents, equipped with tools to explore the trace selectively. Just like a well-written test calls a specific function and checks a specific return value, an agentic judge calls a specific tool on the trace and checks a specific property.

Here are some example judges that we have used across our agents:

Skill invocation judge explores the trace and identifies whether the agent invoked skills that are targeted by the scenario (if not, then the skill’s purpose is not clear to the AUT):

Best-practices judge explores whether the output follows best practices according to Databricks official documentation:

Outcome Judge inspects the trace for output assets and asserts certain properties. Going back to the Data Analyst example, identify the part of the trace where engineering code was authored and evaluate whether the code is appropriate for the task at hand:

This judge is interesting because it tackles the subjective quality problem head-on: what counts as good feature engineering depends on domain expertise. An LLM judge can't get this right out of the box. It's tempting to try writing out the complete criteria in the judge's prompt: "prefer median imputation over mean for skewed distributions, always scale features before distance-based models, ..." But encoding a domain expert's full judgment into a prompt is laborious and brittle. It's much easier for humans to look at an example and say "this is good" or "this is bad" than to write out the complete spec. This is exactly why alignment works, as we'll cover shortly.

In general, our test suite for a single agent includes judges across several categories:

Deterministic checks, things we can verify mechanically, no LLM needed:

• Syntax/linting on generated code
• Output schema validation (do expected tables exist? are column types correct?)
• Tool sequence linting (did the agent read the error logs before trying to fix the issue, or did it skip straight to editing code?)

LLM-based checks, judgment calls that require understanding context:

• Code diff guidelines (did the agent change unrelated lines? did it introduce deprecated APIs?)
• Best practice adherence (is the generated code following the conventions for this domain?)

Operational metrics, signals that don't pass/fail individually but track health over time:

• Token usage (high token counts often signal the agent is struggling, retrying, backtracking, or going in circles)
• Tool call counts and failure ratios (a spike in failed tool calls indicates something is wrong)
• Latency (wall-clock time for the agent to complete the task)

The operational metrics deserve a note. They don't gate a release the way pass/fail judges do, but they're critical for cost management and early warning. If token usage doubles after a change, something went wrong even if all judges still pass; the agent is probably doing more work than it should. We track these over time and alert on anomalies.

Growing the test suite over time

Test suites don't get authored in one sitting. They evolve over time. They start with the simplest checks that give a signal: does the output exist? Does it parse? Then structural checks follow: does the output have the right schema, the right columns, the right types? Only later come end-to-end data validation judges: does the output actually produce correct results when you run it?

This mirrors how test suites mature in traditional software. Exhaustive integration tests don't come on day one. It starts with smoke tests, then unit tests as failure modes emerge, building toward end-to-end coverage over time. The key is that the infrastructure supports adding new judges cheaply, so the test suite grows alongside the agent.

Testing the Tests: Judge Alignment

Here's a problem every engineer knows: a flaky or wrong test suite that greenlights bad code ships bugs with confidence. Similarly, judges who approve poor outcomes give a false sense of security. This is where the second loop of the coSTAR framework comes in: the same scenarios and traces that drive agent refinement also drive judge refinement, with human expert scores as the ground truth. This matters because, unlike traditional testing where test correctness can be verified by inspection, LLM judges are stochastic and can drift in how they interpret natural-language criteria. So we need a way to verify them and keep them aligned with human experts.

To do this alignment, we first curate a Golden Set of typically dozens of examples of agent outputs that our engineers have manually assessed. This is the ground truth the judges must agree with. Then we leverage MLflow's alignment capabilities (powered by techniques like GEPA and MemAlign) to automatically refine the judge against the Golden Set. Notice this is structurally the same STAR loop we use to refine the AUT itself, but the assess step is performed by human experts and the refine step applies to the judge.

R - Refine

With judges that the judge loop has aligned against human expert judgment, we can now trust the agent loop. A coding assistant treats the agent as its codebase and the judges as its test suite. It reads failures, diagnoses root causes, patches the agent, and re-runs everything. The engineer is still the reviewer and final arbiter of the proposed changes to the agent, but this automated iteration saves considerable human effort in analyzing and improving the agent.

Here's what one iteration looked like for the Data Analyst agent:

Red. We ran the initial version of the agent against our scenario suite. The best-practices judge flagged a discrepancy: our agent was generating code for logical views that was different from our official recommendations/documentation. While this discrepancy would not affect correctness, it had implications on the maintenance and deployment of the generated code. This is an example of an insidious regression that would be hard to catch by manual investigation.

Green. The coding assistant analyzed the judge feedback and identified the gap: the agent was using a skill that was not prescriptive about the type of views that should be created (temporary vs permanent). After adding the relevant guidance to the skill, the tests passed successfully and the change was verified to not introduce other regression (based on other test scenarios).

Regression Tests for Infrastructure, Not Just the Agent

So far we've described judges as tests for the agent, catching regressions when the agent implementation changes. But in practice, the agent itself isn't the only thing that changes. The agent depends on external tools and infrastructure, and those change too.

Our agents call MCP tools, standardized interfaces for data access, code execution, environment setup, and more. These tools have their own development teams and release cycles. When a tool changes its implementation (say, a code execution tool starts returning stderr in a different format, or a data access tool changes how it handles null values) the agent hasn't changed at all, but the agent's behavior can break.

Because we run our judges on every nightly build, they act as regression tests against the full stack, not just the agent’s current implementation. When a tool team ships a change that causes an agent to start failing its judges then we catch the error immediately, before it reaches customers. More importantly, the judge's failure tells us what broke (the specific quality dimension that regressed), which makes it far easier to triage whether the root cause is in the agent or in a tool the agent depends on.

This is the same value that integration tests provide in traditional software: they guard the contract between the code and its dependencies. The only difference is that here, the "code" is an agent and the "dependencies" are MCP tools.

From Eval to Production Monitoring

There's one more extension of the testing analogy that turned out to be surprisingly valuable: running the same judges on production traffic.

In traditional software, testing doesn't stop at CI. Production gets monitored too: error rates, latency percentiles, business metrics on live traffic. The same test logic that validates code in dev often reappears as health checks and alerts in prod.

We do the same thing. The judges we built for eval are designed to score any agent conversation, not just eval scenarios. So we run them (or a sampled subset) on real production conversations. This gives us:

• Early warning on drift. If judge's pass rate drops on production conversations, something changed. Maybe a model upgrade degraded quality, maybe user prompts shifted in a way the agent handles poorly. We see it in the judge scores before we see it in user complaints.
• Real-world signal for the test suite. Production conversations that judges flag as failures become candidates for new eval scenarios. This is how the test suite grows organically: real failures feed back into eval, closing the loop between production and development.
• Cost monitoring at the agent level. We track token usage and tool call counts on production conversations. A quality-neutral change that triples cost is still a regression.

重要なのは、同じ評価インフラストラクチャ（評価者、メトリクス、記録されたトレース）が二重の役割を果たすということです。評価のために一度構築すれば、本番環境の監視は副次的な効果として得られます。

現状

私たちは、Databricksプラットフォームでリリースしたいくつかのエージェント（例：Genieのデータエンジニアリング、機械学習、トレース分析機能）、開発者生産性向上のための社内エージェント、およびその他の顧客向けエージェント（例：AI Dev Kit、OSS MLflow Assistant）全体でこの方法論を採用してきました。全体として、具体的なメリットを実感しています。

• 手動評価と比較して、自動テストスイートにより、変更の検証にかかる時間が2週間から数時間に短縮されました。これにより、チームはより迅速に改善をリリースできるようになりました。
• いくつかのテストスイートは、エージェントごとに数百のテストシナリオにまで拡張され、リグレッションの検出に対する信頼性が向上しました。
• 統合テストは、依存インフラストラクチャの変更をフラグ付けし、本番環境でのリグレッションを防ぐことができました。これらの変更の例としては、基盤となるモデルのTODO管理動作、レイテンシに影響を与える変更、またはモデルの変更などが挙げられます。

MLflowは、GenAIテストプラットフォームとしても不可欠であり、エンジニアが方法論を標準化し、テスト開発を加速し、チーム間でベストプラクティスを共有するのに役立っています。

まだ機能しないこと

ここでもテストのアナロジーが役立ちます。私たちの限界は、よく知られたテストの問題にマッピングされます。

シナリオ生成は手動です（テストケースの作成はコストがかかります）。スコアリング、アライメント、最適化は自動化しましたが、シナリオ自体の生成はまだ人間のタスクです。各シナリオでは、現実的な初期状態、意味のあるプロンプト、および正しい期待値を慎重に作成する必要があります。これがテストスイートのサイズを制限するボトルネックであり、狭いテストスイートは次の問題に直接つながります。シナリオ生成の自動化（本番トラフィックパターンまたはエージェントの仕様から、多様で現実的なテストケースを合成すること）は、私たちにとって活発な研究分野です。

コーディングアシスタントは過学習する可能性があります（テストスイートが狭すぎる）。テストスイートが十分なケースをカバーしていない場合、コーディングアシスタントは、それらの特定の入力で優れたパフォーマンスを発揮するが、新しい入力では失敗するエージェント実装をエンジニアリングします。これは、単体テストはパスするが本番環境では失敗するコードを書くことに相当するエージェント版です。私たちは、本番環境の失敗を評価にフィードバックし、時間をかけてカバレッジを拡大することでこれを軽減していますが、シナリオ生成が自動化されるまで、テストスイートの成長は私たちの望むよりも遅くなります。

評価者のアライメントはコストがかかります（テストのキャリブレーションには人的労力が必要です）。ゴールデンセットの構築には、ドメインエキスパートが手動で出力を採点する必要があります。これは、まさに私たちが排除しようとしているボトルネックです。そして、これは一度限りのコストではありません。エージェントが進化するにつれて、評価者の再キャリブレーションが必要です。私たちは、評価者の不確実性を測定し、評価者が不明確である特定の例を特定し、人間のラベルが実際に曖昧さを解消するような方法を調査しています。目標は、評価者アライメントのためのアクティブラーニングです。ランダムなサンプルを採点するように専門家に依頼するのではなく、評価者が不確実であり、ドメインエキスパートの入力がその基準を最もシャープにする例のみを提示します。

複数ステップの失敗の帰属は困難です（根本原因分析）。エージェントが10ステップのパイプラインのステップ7で失敗した場合、根本原因はステップ7にあったのか、それともステップ3にあったのか？私たちの評価者は症状を捉えますが、コーディングアシスタントは間違った関数を変更してテストの失敗を修正するなど、間違ったステップをパッチすることがあります。より良い因果トレースは活発な研究分野です。

新しい失敗モードが見逃される（カバレッジギャップ）。coSTARは、評価者がカバーする次元内で最適化します。新しいクラスの失敗が出現し、どの評価者もチェックしない場合、それはコード内のテストで実行されないバグのように、見えなくなります。coSTARはテストスイート内で改善しますが、テストスイートを独自に拡張することはできません。人間は依然として新しい失敗モードに気づき、評価者を追加する必要があります。

主なテイクアウェイ

1. エージェント開発にはテストの問題があります。自動評価なしでは、テストなしでコーディングしていることになり、当然リグレッションが発生します。
2. 評価者にトレースではなくツールを与えます。ターゲットツールを呼び出すエージェント評価者は、集中した単体テストのようなものです。評価者に完全なトレースをダンプすることは、アサーションにプログラム状態をダンプするようなものです。それはスケーリングしません。
3. テストをテストします。LLM評価者は確率的です。仕様に対してテストスイートを検証するのと同じ方法で、人間が採点したゴールデンセットに対してそれらをアラインします。
4. ループを閉じます。真の勝利は、完全なcoSTARループです。信頼できるシナリオ、記録されたトレース、アラインされた評価者、そしてテストがパスするまでエージェントを洗練するコーディングアシスタント。自動評価なしの評価は物語の半分にすぎません。
5. 一度構築し、どこでも監視します。評価で検証するのと同じ評価者が本番環境を監視できます。1回の投資で2倍のリターン。
6. 連携が重要です。エージェントの洗練は、それを駆動する評価者の信頼性と同じくらい信頼できます。coSTARの2つの連携ループ（評価者の信頼を獲得するループと、その信頼を使用してエージェントを洗練するループ）は、自動洗練を単に高速にするだけでなく、意味のあるものにしています。

私たちは、MLflowの一部としてcoSTARを構築しています。同様の問題に取り組んでいる場合は、ぜひお聞かせください。

• coSTAR方法論を使用して出荷した機能を確認するには、Genie Codeを試してください。
• イテレーティブなエージェント洗練のためのLLM評価者の定義と使用を開始するには、MLflowのチュートリアルに従ってください。

(このブログ記事はAI翻訳ツールを使用して翻訳されています) 原文記事