Great end-to-end write-up — the MLflow tracing setup in particular is a pattern more teams should adopt from day one rather than bolting on later.

One thing worth extending that tracing to: per-node cost attribution. Since every ChatDatabricks call in the LangGraph graph is a separate model invocation, the context window grows with each tool result appended to the message list. In a session where the agent calls Genie + Vector Search + Lakebase in sequence, you're often paying 3–5× what a single-query baseline would suggest — because each tool output accumulates in the prompt for all subsequent steps.

MLflow spans already capture token counts per step. Adding a cost_usd attribute to each span (input_tokens * price_in + output_tokens * price_out) and grouping by session_id gives you a per-conversation cost view that's essential for tuning buffer sizes before you hit production scale.

Two optimizations that help: (1) Summarize each tool response before appending it to the graph state — passing a 50-token summary instead of a 500-token raw Genie result keeps context lean across subsequent nodes. (2) Not all nodes need Sonnet-class reasoning. The routing step (which tool to call?) is usually a simpler classification task than the final answer synthesis. Splitting that to a lighter model and reserving claude-sonnet-4-5 for the final generation step can cut per-session cost 30–50% with no visible quality regression on the routing decision itself.