How are you attributing agent cost per successful task, not just per token?

[DoTA]

Hi everyone, I hope this helps start a useful comparison of notes for those running agents on Databricks in production.

Token cost per endpoint is easy to get. Tag the serving endpoints, and the custom tags propagate to system.billing.usage, where usage_metadata.endpoint_name lets you aggregate by team or app (the model serving cost docs and the example cost attribution dashboard cover this). What I have found much harder is the number the business actually asks about: what does one successful task cost?

Where it gets awkward:

1. A single user request can fan out into several model calls, tool calls and retries, so billing rows (hourly or per endpoint) do not map cleanly to one request. MLflow traces do map to one request, but they do not carry dollars.
2. "Successful" needs a definition outside billing. Thumbs up, a downstream action completing, or an LLM judge score all give different denominators, and the cost per success can swing a lot depending on which you pick.
3. Retries and long tool loops quietly dominate spend on the bad days, and they are invisible in a per-token average.

The approach I am considering: log token counts per call into the trace, price them with a small rate table, join to the outcome signal by trace or request id, and use the billing system table only as a reconciliation check on the total. I have not got this running end to end yet, so I would like to hear what others do.

Questions for the community:

• Are you computing cost per task from traces and token counts, from billing tables, or something else?
• What are you using as the success signal?
• Has anyone found a clean way to attribute shared endpoint spend back to individual agents when several agents call the same model endpoint?

[coolbeans201]

Great question (also looking at this as we're building an internal AI agent that can do similar tasks). I think the key is to log every step of a process, including the metadata associated with those tasks themselves. If Jira tickets and/or PR links are part of what an agent can do, make sure it's logged. If you can associate token usage with successful outcomes (and somehow even associate that with dev hours saved), then it's much easier to justify spend.

[SumeshKashyap]

These are great questions. The third one in particular is where most setups I've seen get stuck, so thanks for starting this thread. Here's what tends to work, taking your questions one at a time.

Q1: Are you computing cost per task from traces and token counts, from billing tables, or something else?

Both, but each for a different job. Your instinct is right: price at the request level and keep system.billing.usage as the reconciliation check. The one thing I'd change is where the token counts come from. Rather than logging them into the trace yourself, let the gateway record them. With Unity Gateway, every call lands in system.ai_gateway.usage with input_tokens, output_tokens, the destination model and any request_tags you send in the Databricks-Ai-Gateway-Request-Tags header. If you send task_id and trace_id in that header, grouping by task gives you the full fan-out including retries. It also catches calls that happen outside your traced code.

The MLflow trace is still where you look when you want to know why a task was expensive. MLflow 3 also rolls up token usage across the LLM spans (trace.info.token_usage), so you get a cross-check for free. Multiply the tokens by a small versioned rate table, then once a month compare the total against system.billing.usage filtered on usage_metadata.endpoint_name. A gap of a few percent is normal. A big gap usually means some calls are going out untagged.

One thing worth deciding early: Vector Search, provisioned throughput, app compute and warehouse time used by tools don't bill per request. Either spread them as overhead or leave them out and call the metric "model cost per success". Just don't mix the two without saying so.

Q2: What are you using as the success signal?

I'd log the outcome as feedback on the trace with mlflow.log_feedback, so cost and outcome join on trace id without needing a separate system. Thumbs up/down go in with a HUMAN source and LLM-judge scores come from production monitoring. "The downstream thing actually happened" (ticket closed, PR merged, record updated) gets written by a small job or webhook with a CODE source.

Keep all three, but report against one. Downstream completion is usually the one the business doesn't argue with. Thumbs are too sparse and judge scores drift, so they work better as early warnings.

The bigger point is the formula itself. Cost per success should be the cost of every attempt divided by the number of successes, not the average cost of the tasks that succeeded. Otherwise the retries and long tool loops you mentioned simply disappear. Put p95 cost per task and the share of spend that went to failed tasks next to it, because that's where the bad days actually show up.

Q3: Has anyone found a clean way to attribute shared endpoint spend back to individual agents when several agents call the same model endpoint?

Yes, and it's the same trick as in Q1: tag the request, not the endpoint. Endpoint tags can only tell you who owns the endpoint. Add a header like this on each call:

Databricks-Ai-Gateway-Request-Tags: {"agent": "invoice_agent", "task_id": "...", "trace_id": "...", "attempt": "2"}

That lets you split one shared endpoint's tokens exactly by agent. If you're still on the legacy AI Gateway usage tracking, usage_context and client_request_id in the request body do the same job and end up in system.serving.endpoint_usage. Either way, these tables are admin-only by default, so you'll probably need someone to grant you SELECT or expose a view.

The query ends up roughly like this:

WITH calls AS (

  SELECT request_tags['agent'] AS agent, request_tags['task_id'] AS task_id,

         input_tokens * r.input_rate + output_tokens * r.output_rate AS cost

  FROM system.ai_gateway.usage u

  JOIN my_rates r ON r.model = u.destination_model

  WHERE u.event_time >= current_date() - INTERVAL 30 DAYS

),

tasks AS (

  SELECT agent, task_id, SUM(cost) AS task_cost

  FROM calls GROUP BY ALL

)

SELECT t.agent,

       SUM(task_cost) / NULLIF(COUNT_IF(o.success), 0) AS cost_per_success,

       percentile(task_cost, 0.95) AS p95_task_cost,

       SUM(IF(o.success, 0, task_cost)) / SUM(task_cost) AS wasted_share

FROM tasks t LEFT JOIN task_outcomes o USING (task_id)

GROUP BY t.agent;

Here task_outcomes is the success feedback pulled out of your trace table.

Hope this helps. I'd be curious which success definition you settle on once it's running end to end.