fix(agent-sessions): net usage without copying lookup tables per reporter - #1197
Conversation
…rter The list and distributions reads looked up each reporter's ancestors and charged claims inside a lambda that named the trace's link maps and the session's claims. ClickHouse copies a captured column once per element, so memory grew with reporters x table size, per trace for every trace in the window, and a window of long agent runs exceeded the list profile's 1.5 GB. The lookups are now a sort-merge over arrays (table entries and needles sorted by key, arrayFill carrying the value down each run), written as one expression per lookup so the query tree stays linear. The session's reporters are capped by groupArrayArray instead of flatten-then-slice. Same rows: identical output on a production week and on randomized span trees against the previous SQL. A trace of 1,500 model calls now nets inside 100 MB in the ClickHouse e2e; the previous SQL fails that case.
Maple review🟢 Confidence 4/5 · likely safe to merge Rebuilds the Agent Sessions usage netting so ancestors and child claims come from sort-merge array lookups instead of lambdas that name a table per reporter, which is what blew the list read's 1.5 GB ceiling. The netting rules and the SQL-text, e2e and new low-memory assertions all still hold.
What was checked
|
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (5)
Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 2 remain after this review. 📝 WalkthroughWalkthroughAI session usage aggregation now resolves token and cost ancestry per trace, gathers reporters across traces, and nets claims at the session level. Tests cover the updated SQL and a trace with 1,500 model-call spans. ChangesAI usage aggregation
Priority: ⬆️ High Estimated code review effort: 3 (Moderate) | ~25 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to This change rewrites how Agent Sessions nets token and cost usage, to avoid memory-limit failures on long agent runs. No concrete correctness or stability defect was established at the current head, and the one open question about link truncation was shown to match the previous behavior. 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
1 participant
Problem
The Agent Sessions list (
aiSessionsPage) and the sidebar distributions (aiSessionsDistributions) fail withMEMORY_LIMIT_EXCEEDED(code 241) at thelistprofile's 1.5 GB for orgs with long agent runs. The facets read, which does no netting, succeeds over the same window.Cause: the usage netting did its lookups inside lambdas that name a column.
LIMIT):arrayMap(r -> … tokenLinks[tokenLinks[…[r.2]]] …, usageReporters)indexOf(childClaims…, r.1)andhas(reporterIds, …)ClickHouse copies a captured column once per array element, so memory is reporters x table size. ClickHouse names the first one in the error:
while executing 'FUNCTION arrayElement(tokenLinks, tupleElement(r, 2))'.Fix
lookupExpr: the same keyed lookup as a sort-merge over arrays. Table entries and needles are sorted together by key, andarrayFillcarries the table's value down each run. No lambda names a column, so memory is linear.t<SpanId>/c<SpanId>keys for tokens / cost).sumMap, read back in one lookup.bind), because same-level aliases expand wherever they are named: the first attempt with aliases hitQuery tree is too bigand, once under the limit, ran 4x slower.groupArrayArray(2000)instead ofarraySlice(arrayFlatten(groupArray(…))), so the cap bounds the aggregate state too.Netting rules are unchanged.
Verification
13809315504893199302db.duration_ms(2-thread interactive profile, first run of each)max_memory_usage = 100 MBai-trace-index-materialization.clickhouse.e2e.test.ts: 10/10catalog.clickhouse.e2e.test.ts: 405/405packages/query-engine-integrationssrc/ai: 309/309Not covered
deepestFailureCountstill namesfailedSpansinside its lambda (has(tupleElement(failedSpans, 2), f.1)). Same class of cost for a trace with many failed spans; not what failed here, left for a follow-up.Need help on this PR? Tag
@codesmith-botwith what you need. Autofix is disabled.Summary by CodeRabbit