feat(aio): honest per-tool spend attribution and cost-by-context-size

products/ai_observability/backend/api/test/__snapshots__/test_personal_spend.ambr

@@ -82,6 +82,91 @@
   '''
 # ---
 # name: TestPersonalSpendQueries.test_empty_result_when_no_events.2
+  '''
+  /* user_id:0 request:_snapshot_@me_spend_ */
+  SELECT tool AS tool,
+         round(sum(cost_share), 6) AS cost_attributed_usd
+  FROM
+    (SELECT nullIf(arrayJoin(if(empty(tools), [''], tools)), '') AS tool,
+            divide(cost, length(if(empty(tools), [''], tools))) AS cost_share
+     FROM
+       (SELECT arrayDistinct(arrayFilter(x -> ifNull(notEquals(x, ''), 1), splitByChar(',', coalesce(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_tools_called'), ''), 'null'), '^"|"$', ''), '')))) AS tools,
+               accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_total_cost_usd'), ''), 'null'), '^"|"$', ''), 'Float64') AS cost
+        FROM events
+        LEFT OUTER JOIN
+          (SELECT argMax(person_distinct_id_overrides.person_id, person_distinct_id_overrides.version) AS person_id,
+                  person_distinct_id_overrides.distinct_id AS distinct_id
+           FROM person_distinct_id_overrides
+           WHERE equals(person_distinct_id_overrides.team_id, 99999)
+           GROUP BY person_distinct_id_overrides.distinct_id
+           HAVING equals(argMax(person_distinct_id_overrides.is_deleted, person_distinct_id_overrides.version), 0) SETTINGS optimize_aggregation_in_order=1) AS events__override ON equals(events.distinct_id, events__override.distinct_id)
+        LEFT JOIN
+          (SELECT person.id AS id,
+                  replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(person.properties, 'email'), ''), 'null'), '^"|"$', '') AS properties___email
+           FROM person
+           WHERE and(equals(person.team_id, 99999), in(tuple(person.id, person.version),
+                                                         (SELECT person.id AS id, max(person.version) AS version
+                                                          FROM person
+                                                          WHERE equals(person.team_id, 99999)
+                                                          GROUP BY person.id
+                                                          HAVING and(equals(argMax(person.is_deleted, person.version), 0), less(argMax(toTimeZone(person.created_at, 'UTC'), person.version), plus(now64(6, 'UTC'), toIntervalDay(1))))))) SETTINGS optimize_aggregation_in_order=1) AS events__person ON equals(if(not(empty(events__override.distinct_id)), events__override.person_id, events.person_id), events__person.id)
+        WHERE and(equals(events.team_id, 99999), equals(events.event, '$ai_generation'), ifNull(equals(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, 'ai_product'), ''), 'null'), '^"|"$', ''), 'posthog_code'), 0), ifNull(equals(events__person.properties___email, 'user1@posthog.com'), 0), and(greaterOrEquals(events.timestamp, toDateTime64('explicit_redacted_timestamp', 6, 'UTC')), less(events.timestamp, toDateTime64('today', 6, 'UTC'))))))
+  GROUP BY tool
+  LIMIT 100 SETTINGS readonly=2,
+                     max_execution_time=60,
+                     allow_experimental_object_type=1,
+                     max_ast_elements=4000000,
+                     max_expanded_ast_elements=4000000,
+                     max_bytes_before_external_group_by=0,
+                     transform_null_in=1,
+                     optimize_min_equality_disjunction_chain_length=4294967295,
+                     optimize_rewrite_aggregate_function_with_if=0,
+                     optimize_min_inequality_conjunction_chain_length=4294967295,
+                     allow_experimental_join_condition=1,
+                     use_hive_partitioning=0
+  '''
+# ---
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.3
+  '''
+  /* user_id:0 request:_snapshot_@me_spend_ */
+  SELECT multiIf(ifNull(less(coalesce(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_input_tokens'), ''), 'null'), '^"|"$', ''), 'Float64'), 0), 50000), 0), '<50k', ifNull(less(coalesce(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_input_tokens'), ''), 'null'), '^"|"$', ''), 'Float64'), 0), 100000), 0), '50k-100k', ifNull(less(coalesce(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_input_tokens'), ''), 'null'), '^"|"$', ''), 'Float64'), 0), 150000), 0), '100k-150k', ifNull(less(coalesce(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_input_tokens'), ''), 'null'), '^"|"$', ''), 'Float64'), 0), 200000), 0), '150k-200k', ifNull(less(coalesce(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_input_tokens'), ''), 'null'), '^"|"$', ''), 'Float64'), 0), 300000), 0), '200k-300k', '>300k') AS bucket,
+         count() AS generation_count,
+         round(sum(accurateCastOrNull(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_total_cost_usd'), ''), 'null'), '^"|"$', ''), 'Float64')), 6) AS cost_usd
+  FROM events
+  LEFT OUTER JOIN
+    (SELECT argMax(person_distinct_id_overrides.person_id, person_distinct_id_overrides.version) AS person_id,
+            person_distinct_id_overrides.distinct_id AS distinct_id
+     FROM person_distinct_id_overrides
+     WHERE equals(person_distinct_id_overrides.team_id, 99999)
+     GROUP BY person_distinct_id_overrides.distinct_id
+     HAVING equals(argMax(person_distinct_id_overrides.is_deleted, person_distinct_id_overrides.version), 0) SETTINGS optimize_aggregation_in_order=1) AS events__override ON equals(events.distinct_id, events__override.distinct_id)
+  LEFT JOIN
+    (SELECT person.id AS id,
+            replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(person.properties, 'email'), ''), 'null'), '^"|"$', '') AS properties___email
+     FROM person
+     WHERE and(equals(person.team_id, 99999), in(tuple(person.id, person.version),
+                                                   (SELECT person.id AS id, max(person.version) AS version
+                                                    FROM person
+                                                    WHERE equals(person.team_id, 99999)
+                                                    GROUP BY person.id
+                                                    HAVING and(equals(argMax(person.is_deleted, person.version), 0), less(argMax(toTimeZone(person.created_at, 'UTC'), person.version), plus(now64(6, 'UTC'), toIntervalDay(1))))))) SETTINGS optimize_aggregation_in_order=1) AS events__person ON equals(if(not(empty(events__override.distinct_id)), events__override.person_id, events.person_id), events__person.id)
+  WHERE and(equals(events.team_id, 99999), in(events.event, tuple('$ai_generation', '$ai_embedding')), ifNull(equals(replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, 'ai_product'), ''), 'null'), '^"|"$', ''), 'posthog_code'), 0), ifNull(equals(events__person.properties___email, 'user1@posthog.com'), 0), and(greaterOrEquals(events.timestamp, toDateTime64('explicit_redacted_timestamp', 6, 'UTC')), less(events.timestamp, toDateTime64('today', 6, 'UTC'))))
+  GROUP BY bucket
+  LIMIT 100 SETTINGS readonly=2,
+                     max_execution_time=60,
+                     allow_experimental_object_type=1,
+                     max_ast_elements=4000000,
+                     max_expanded_ast_elements=4000000,
+                     max_bytes_before_external_group_by=0,
+                     transform_null_in=1,
+                     optimize_min_equality_disjunction_chain_length=4294967295,
+                     optimize_rewrite_aggregate_function_with_if=0,
+                     optimize_min_inequality_conjunction_chain_length=4294967295,
+                     allow_experimental_join_condition=1,
+                     use_hive_partitioning=0
+  '''
+# ---
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.4
   '''
   /* user_id:0 request:_snapshot_@me_spend_ */
   SELECT replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, 'ai_product'), ''), 'null'), '^"|"$', '') AS product,
@@ -122,7 +207,7 @@
                     use_hive_partitioning=0
   '''
 # ---
-# name: TestPersonalSpendQueries.test_empty_result_when_no_events.3
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.5
   '''
   /* user_id:0 request:_snapshot_@me_spend_ */
   SELECT replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_model'), ''), 'null'), '^"|"$', '') AS model,
@@ -165,7 +250,7 @@
                     use_hive_partitioning=0
   '''
 # ---
-# name: TestPersonalSpendQueries.test_empty_result_when_no_events.4
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.6
   '''
   /* user_id:0 request:_snapshot_@me_spend_ */
   SELECT toDate(events.timestamp) AS day,
@@ -208,7 +293,7 @@
                     use_hive_partitioning=0
   '''
 # ---
-# name: TestPersonalSpendQueries.test_empty_result_when_no_events.5
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.7
   '''
   /* user_id:0 request:_snapshot_@me_spend_ */
   SELECT replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(events.properties, '$ai_model'), ''), 'null'), '^"|"$', '') AS model
@@ -247,7 +332,7 @@
                    use_hive_partitioning=0
   '''
 # ---
-# name: TestPersonalSpendQueries.test_empty_result_when_no_events.6
+# name: TestPersonalSpendQueries.test_empty_result_when_no_events.8
   '''
   /* user_id:0 request:_snapshot_@me_spend_ */
   SELECT toDate(events.timestamp) AS day,

products/ai_observability/backend/api/test/test_personal_spend.py

@@ -278,9 +278,10 @@ def test_empty_result_when_no_events(self) -> None:
         assert body["by_product"] == {"items": [], "truncated": False}
         assert body["by_tool"] == {"items": [], "truncated": False}
         assert body["by_model"] == {"items": [], "truncated": False}
+        assert body["by_input_size"] == {"items": [], "truncated": False}
         assert body["by_day"] == {"items": [], "truncated": False}
         assert body["by_day_model"] == []
-        assert body["top_traces"] == {"items": [], "truncated": False}
+        assert "top_traces" not in body
 
     def test_summary_reports_cross_product_totals_alongside_scoped(self) -> None:
         self._create_generation(ai_product="posthog_code", cost=2.0)
@@ -358,16 +359,97 @@ def test_by_tool_splits_comma_separated_tools(self) -> None:
         # Scoped total stays $3.50 — the multi-tool generation is only counted once there.
         assert response.json()["summary"]["scoped_cost_usd"] == 3.5
 
-    def test_top_traces_always_empty(self) -> None:
-        # `top_traces` is deprecated — kept in the response shape so existing consumers don't
-        # crash, but always returns empty. Trace IDs are opaque strings that aren't actionable
-        # in the UI.
+    def test_by_tool_cost_attributed_splits_evenly_across_distinct_tools(self) -> None:
+        # Same fixture as test_by_tool_splits_comma_separated_tools: a multi-tool generation
+        # plus two single-tool generations. Unlike `cost_usd` (which double-counts the
+        # multi-tool generation into both rows), `cost_attributed_usd` must split its $2.0
+        # evenly across Bash and Read so the per-tool sum reconciles to scoped_cost_usd
+        # instead of overcounting -- this is the regression the whole PR exists to fix.
+        self._create_generation(tool="Bash,Read", cost=2.0, trace_id="multi")
+        self._create_generation(tool="Bash", cost=1.0, trace_id="bash-only")
+        self._create_generation(tool="Read", cost=0.5, trace_id="read-only")
+        flush_persons_and_events()
+
+        response = self.client.get(f"{ENDPOINT}?product=posthog_code")
+        body = response.json()
+        rows = {r["tool"]: r for r in body["by_tool"]["items"]}
+
+        # Bash: $1.0 (bash-only) + $1.0 (half of the $2.0 multi-tool generation) = $2.0.
+        assert rows["Bash"]["cost_attributed_usd"] == 2.0
+        # Read: $0.5 (read-only) + $1.0 (half of the $2.0 multi-tool generation) = $1.5.
+        assert rows["Read"]["cost_attributed_usd"] == 1.5
+        # `cost_usd` still double-counts (unchanged, existing behavior) -- the two fields
+        # must genuinely differ, not just both exist.
+        assert rows["Bash"]["cost_usd"] == 3.0
+        assert rows["Read"]["cost_usd"] == 2.5
+
+        # Attributed cost reconciles to scoped_cost_usd; cost_usd does not ($5.5 > $3.5).
+        scoped = body["summary"]["scoped_cost_usd"]
+        assert scoped == 3.5
+        assert sum(r["cost_attributed_usd"] for r in rows.values()) == scoped
+        assert sum(r["cost_usd"] for r in rows.values()) != scoped
+
+        # share_attributed is cost_attributed_usd normalized the same way, and reconciles too.
+        assert rows["Bash"]["share_attributed"] == 2.0 / 3.5
+        assert rows["Read"]["share_attributed"] == 1.5 / 3.5
+        assert sum(r["share_attributed"] for r in rows.values()) == 1.0
+
+    def test_by_tool_cost_attributed_falls_back_to_null_tool_for_no_tools(self) -> None:
+        # A generation with no tools attributes its full cost to the same NULL `tool`
+        # bucket the existing (unattributed) rows already use for "no tool called".
+        self._create_generation(tool=None, cost=1.0)
+        flush_persons_and_events()
+
+        response = self.client.get(f"{ENDPOINT}?product=posthog_code")
+        rows = {r["tool"]: r for r in response.json()["by_tool"]["items"]}
+        assert rows.keys() == {None}
+        assert rows[None]["cost_attributed_usd"] == 1.0
+
+    def test_by_input_size_buckets_generations_by_input_tokens(self) -> None:
+        # One generation per bucket boundary, each with a distinguishing cost so the
+        # buckets can't be confused for one another.
+        self._create_generation(cost=1.0, input_tokens=10_000, trace_id="under-50k")
+        self._create_generation(cost=2.0, input_tokens=75_000, trace_id="50k-100k")
+        self._create_generation(cost=4.0, input_tokens=350_000, trace_id="over-300k")
+        # An embedding counts too -- by_input_size uses the same event set as by_model.
+        self._create_generation(cost=0.5, input_tokens=5_000, event_name="$ai_embedding", trace_id="embedding")
+        flush_persons_and_events()
+
+        response = self.client.get(ENDPOINT_OK)
+        body = response.json()
+        by_bucket = {r["bucket"]: r for r in body["by_input_size"]["items"]}
+
+        assert by_bucket["<50k"]["generation_count"] == 2  # the 10k generation + the 5k embedding
+        assert by_bucket["<50k"]["cost_usd"] == 1.5
+        assert by_bucket["<50k"]["min_input_tokens"] == 0
+        assert by_bucket["50k-100k"]["generation_count"] == 1
+        assert by_bucket["50k-100k"]["cost_usd"] == 2.0
+        assert by_bucket["50k-100k"]["min_input_tokens"] == 50_000
+        assert by_bucket[">300k"]["generation_count"] == 1
+        assert by_bucket[">300k"]["cost_usd"] == 4.0
+        assert by_bucket[">300k"]["min_input_tokens"] == 300_000
+        # Buckets with no events are omitted, not zero-filled.
+        assert set(by_bucket) == {"<50k", "50k-100k", ">300k"}
+
+        # Ordered ascending by min_input_tokens, not by cost.
+        assert [r["bucket"] for r in body["by_input_size"]["items"]] == ["<50k", "50k-100k", ">300k"]
+
+        # cost_usd across buckets reconciles to scoped_cost_usd.
+        scoped = body["summary"]["scoped_cost_usd"]
+        assert scoped == 7.5
+        assert sum(r["cost_usd"] for r in body["by_input_size"]["items"]) == scoped
+        assert sum(r["share_of_scoped"] for r in body["by_input_size"]["items"]) == 1.0
+        assert by_bucket["<50k"]["avg_cost_per_generation"] == 0.75
+
+    def test_top_traces_removed_from_response(self) -> None:
+        # `top_traces` was dropped entirely -- confirmed unused by the PostHog/code client,
+        # which already stopped rendering it. A stale client reading the field just sees it
+        # missing rather than the old always-empty placeholder.
         self._create_generation(trace_id="cheap", cost=0.5)
-        self._create_generation(trace_id="expensive", cost=5.0)
         flush_persons_and_events()
 
         response = self.client.get(ENDPOINT_OK)
-        assert response.json()["top_traces"] == {"items": [], "truncated": False}
+        assert "top_traces" not in response.json()
 
     def test_other_users_spend_is_not_visible(self) -> None:
         other_email = "other-user@example.com"
@@ -396,10 +478,10 @@ def test_date_window_passes_through_to_query_layer(self) -> None:
             response = self.client.get(f"{ENDPOINT}?{PRODUCT_QS}&date_from=-7d")
 
         assert response.status_code == status.HTTP_200_OK
-        # 7 fetchers run once (summary, by_product, by_tool, by_model, by_day, and the two
-        # by_day_model queries). top_traces
-        # is deprecated and returned empty without a query.
-        assert mock_exec.call_count == 7
+        # 9 queries run once: summary, by_product, by_tool (+ its cost-attribution
+        # sub-query), by_model, by_input_size, by_day, and the two by_day_model queries.
+        # top_traces is deprecated and returned empty without a query.
+        assert mock_exec.call_count == 9
 
     def test_second_call_serves_from_cache(self) -> None:
         with patch("products.ai_observability.backend.api.personal_spend.execute_hogql_query") as mock_exec: