feat: track per-session cumulative cost and token usage on CLI#2424
feat: track per-session cumulative cost and token usage on CLI#2424praisonai-triage-agent[bot] wants to merge 2 commits into
Conversation
…2421) Accumulate per-call token usage (already aggregated by the core token collector) into the active session, price it via the CLI cost tracker, and persist cumulative cost/tokens in the session record so they survive resume. - project_sessions.py: add accumulate_session_usage/read_session_usage/ format_usage_footer helpers that merge usage deltas into session metadata. - run.py: record + display a compact usage footer after each prompt run. - resume.py: rehydrate cumulative usage totals into RehydratedSession. - session.py: add Tokens/Cost columns to `session list` and a usage line to the resume panel. Co-authored-by: MervinPraison <MervinPraison@users.noreply.github.com>
|
@coderabbitai review |
|
/review |
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
✅ Action performedReview finished.
|
|
Important Review skippedBot user detected. To trigger a single review, invoke the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds per-session cumulative token and cost tracking to the CLI. New helpers in Session Usage Tracking and Display
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 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 |
|
@copilot Do a thorough review of this PR. Read ALL existing reviewer comments above from Qodo, Coderabbit, and Gemini first — incorporate their findings. Review areas:
|
Greptile SummaryThis PR adds cumulative CLI session usage tracking. The main changes are:
Confidence Score: 4/5This is close, but the global session list path should be fixed before merging.
src/praisonai/praisonai/cli/commands/session.py Important Files Changed
Reviews (2): Last reviewed commit: "fix: price session usage per-model and r..." | Re-trigger Greptile |
| total_tokens = getattr(s, "total_tokens", 0) or 0 | ||
| cost = getattr(s, "cost", 0.0) or 0.0 |
There was a problem hiding this comment.
When praisonai session list --all or a backend-specific list uses SessionManager.list(), each row is a SessionMetadata object without total_tokens or cost attributes. These new columns therefore render as empty for that path even when usage was persisted in session metadata, so global/backend session listings report no usage.
|
|
||
| if bridge is not None: | ||
| bridge.emit_run_result(result, ok=True) | ||
| _record_session_usage(session_id or auto_save_name, model, output) |
There was a problem hiding this comment.
The new usage recording only runs after the cold prompt branches, but _run_prompt can return earlier when the warm runtime handles the prompt. A successful warm-runtime session then never persists or displays its token/cost totals, and the process-global collector is left dirty for a later recorded run.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/praisonai/praisonai/cli/commands/session.py`:
- Around line 92-101: The manager.list() path in session handling is not
mirroring persisted usage totals, so custom-backend and --all listings still
emit empty tokens/cost values. Update the SessionMetadata/listing flow in the
session command to map any usage dict into total_tokens and cost the same way
the project-scoped branch does, or extend SessionMetadata so those fields are
populated before table/JSON output is generated. Use the same usage-to-totals
logic already present in the session loading code to keep both paths consistent.
In `@src/praisonai/praisonai/cli/state/project_sessions.py`:
- Around line 194-206: `read_session_usage()` (and the matching write path
around `update_session_usage()`/the referenced block at 290-297) is using only
`get_project_session_store(project_path)`, while `rehydrate_session()` can fall
back to the global session store; update these helpers to resolve the session
through the same store search order as `rehydrate_session()` so reads and writes
target the actual store that holds the restored session. Keep the cumulative
usage lookup and persistence aligned with the resolved store/session location
instead of always defaulting to the current project store.
- Around line 262-277: The cost calculation in the session update flow is using
the CLI-selected model instead of the per-model breakdown from the summary,
which can persist incorrect totals for multi-model runs. Update the pricing
logic in the project session aggregation path that reads total_metrics so it
iterates over summary["by_model"] (or equivalent per-model totals) and
calculates cost for each model separately, then sum the deltas before writing
them back. Keep the existing token delta handling in the same update method, but
replace the single get_pricing(model or "default") call with model-specific
pricing derived from each entry in by_model.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 9e9e35ac-ebc5-4225-86f5-69a4380666bf
📒 Files selected for processing (5)
src/praisonai/praisonai/cli/commands/run.pysrc/praisonai/praisonai/cli/commands/session.pysrc/praisonai/praisonai/cli/session/resume.pysrc/praisonai/praisonai/cli/state/project_sessions.pysrc/praisonai/tests/unit/test_session_resume.py
|
|
||
| # Cumulative usage totals persisted per session (Issue #2421). | ||
| usage = data.get("usage") | ||
| if isinstance(usage, dict): | ||
| self.total_tokens = usage.get("total_tokens") or 0 | ||
| self.cost = usage.get("cost") or 0.0 | ||
| else: | ||
| self.total_tokens = data.get("total_tokens") or 0 | ||
| self.cost = data.get("cost") or 0.0 | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -i 'sessions.py' src/praisonai/praisonai/cli/state
ast-grep outline src/praisonai/praisonai/cli/state/sessions.py --view expanded || true
rg -n -C3 'class .*Session|def to_dict|def list\(' src/praisonai/praisonai/cli/state/sessions.py src/praisonai/praisonai/cliRepository: MervinPraison/PraisonAI
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '1,220p' src/praisonai/praisonai/cli/commands/session.py | cat -n
printf '\n--- sessions.py ---\n'
sed -n '1,260p' src/praisonai/praisonai/cli/state/sessions.py | cat -n
printf '\n--- project_sessions.py ---\n'
sed -n '1,260p' src/praisonai/praisonai/cli/state/project_sessions.py | cat -nRepository: MervinPraison/PraisonAI
Length of output: 31272
Mirror usage totals in the manager.list() path src/praisonai/praisonai/cli/commands/session.py:128-160
SessionMetadata never carries total_tokens or cost, so --all and custom-backend listings still show empty tokens/cost in both the table and JSON output. Wrap manager.list() results with the same usage → totals mapping used by the project-scoped branch, or add those fields to SessionMetadata.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/cli/commands/session.py` around lines 92 - 101, The
manager.list() path in session handling is not mirroring persisted usage totals,
so custom-backend and --all listings still emit empty tokens/cost values. Update
the SessionMetadata/listing flow in the session command to map any usage dict
into total_tokens and cost the same way the project-scoped branch does, or
extend SessionMetadata so those fields are populated before table/JSON output is
generated. Use the same usage-to-totals logic already present in the session
loading code to keep both paths consistent.
| def read_session_usage( | ||
| session_id: str, project_path: Optional[str] = None | ||
| ) -> Dict[str, Any]: | ||
| """Read the persisted cumulative usage totals for a session. | ||
|
|
||
| Returns a dict with ``input_tokens``/``output_tokens``/``cached_tokens``/ | ||
| ``total_tokens``/``cost``/``requests``. Missing values default to zero so | ||
| callers can render a total even for never-tracked sessions (Issue #2421). | ||
| """ | ||
| usage = _empty_usage() | ||
| try: | ||
| store = get_project_session_store(project_path) | ||
| if not store.session_exists(session_id): |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Resolve usage against the same store search order as rehydrate_session().
rehydrate_session() can restore a session from the global fallback store, but these helpers always read and write through get_project_session_store(project_path). If the restored session only exists in that fallback store, the next accumulation starts from zero and gets persisted into the current project store instead of the original record, so cumulative usage splits/resets after resume.
Also applies to: 290-297
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/cli/state/project_sessions.py` around lines 194 -
206, `read_session_usage()` (and the matching write path around
`update_session_usage()`/the referenced block at 290-297) is using only
`get_project_session_store(project_path)`, while `rehydrate_session()` can fall
back to the global session store; update these helpers to resolve the session
through the same store search order as `rehydrate_session()` so reads and writes
target the actual store that holds the restored session. Keep the cumulative
usage lookup and persistence aligned with the resolved store/session location
instead of always defaulting to the current project store.
| totals = (summary or {}).get("total_metrics") or {} | ||
| delta_in = int(totals.get("input_tokens", 0) or 0) | ||
| delta_out = int(totals.get("output_tokens", 0) or 0) | ||
| delta_cached = int(totals.get("cached_tokens", 0) or 0) | ||
| delta_total = int(totals.get("total_tokens", 0) or 0) | ||
| delta_requests = int((summary or {}).get("total_interactions", 0) or 0) | ||
|
|
||
| if not (delta_in or delta_out or delta_total): | ||
| return current | ||
|
|
||
| delta_cost = 0.0 | ||
| try: | ||
| from ..features.cost_tracker import get_pricing | ||
|
|
||
| pricing = get_pricing(model or "default") | ||
| delta_cost = pricing.calculate_cost(delta_in, delta_out) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Price deltas from summary["by_model"], not the CLI-selected model.
The token collector already returns per-model token totals, but this code prices the entire run with model or "default". Any run that fans out across multiple models will persist the wrong cumulative cost even when token totals are correct.
Suggested fix
- delta_cost = 0.0
- try:
- from ..features.cost_tracker import get_pricing
-
- pricing = get_pricing(model or "default")
- delta_cost = pricing.calculate_cost(delta_in, delta_out)
+ delta_cost = 0.0
+ try:
+ from ..features.cost_tracker import get_pricing
+
+ by_model = (summary or {}).get("by_model") or {}
+ if by_model:
+ for model_name, metrics in by_model.items():
+ pricing = get_pricing(model_name)
+ delta_cost += pricing.calculate_cost(
+ int(metrics.get("input_tokens", 0) or 0),
+ int(metrics.get("output_tokens", 0) or 0),
+ )
+ else:
+ pricing = get_pricing(model or "default")
+ delta_cost = pricing.calculate_cost(delta_in, delta_out)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| totals = (summary or {}).get("total_metrics") or {} | |
| delta_in = int(totals.get("input_tokens", 0) or 0) | |
| delta_out = int(totals.get("output_tokens", 0) or 0) | |
| delta_cached = int(totals.get("cached_tokens", 0) or 0) | |
| delta_total = int(totals.get("total_tokens", 0) or 0) | |
| delta_requests = int((summary or {}).get("total_interactions", 0) or 0) | |
| if not (delta_in or delta_out or delta_total): | |
| return current | |
| delta_cost = 0.0 | |
| try: | |
| from ..features.cost_tracker import get_pricing | |
| pricing = get_pricing(model or "default") | |
| delta_cost = pricing.calculate_cost(delta_in, delta_out) | |
| delta_cost = 0.0 | |
| try: | |
| from ..features.cost_tracker import get_pricing | |
| by_model = (summary or {}).get("by_model") or {} | |
| if by_model: | |
| for model_name, metrics in by_model.items(): | |
| pricing = get_pricing(model_name) | |
| delta_cost += pricing.calculate_cost( | |
| int(metrics.get("input_tokens", 0) or 0), | |
| int(metrics.get("output_tokens", 0) or 0), | |
| ) | |
| else: | |
| pricing = get_pricing(model or "default") | |
| delta_cost = pricing.calculate_cost(delta_in, delta_out) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/praisonai/praisonai/cli/state/project_sessions.py` around lines 262 -
277, The cost calculation in the session update flow is using the CLI-selected
model instead of the per-model breakdown from the summary, which can persist
incorrect totals for multi-model runs. Update the pricing logic in the project
session aggregation path that reads total_metrics so it iterates over
summary["by_model"] (or equivalent per-model totals) and calculates cost for
each model separately, then sum the deltas before writing them back. Keep the
existing token delta handling in the same update method, but replace the single
get_pricing(model or "default") call with model-specific pricing derived from
each entry in by_model.
|
@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Focus ONLY on Python packages (praisonaiagents, praisonai). Do NOT modify praisonai-rust or praisonai-ts. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding. Phase 1: Review per AGENTS.md
Phase 2: FIX Valid Issues Phase 3: Final Verdict |
|
Claude finished @MervinPraison's task —— View job Final Architecture Review — PR #2424Verdict: Approve (with fixes pushed). Two valid Major issues from CodeRabbit/Greptile fixed and pushed to this branch ( Phase 1 — Architecture (AGENTS.md)
Phase 2 — Fixes implemented
Phase 3 — Skipped (with reasons)
Files modified
Scope honored: Python wrapper ( Job run · Branch |
Addresses CodeRabbit/Greptile review feedback on PR #2424: - accumulate_session_usage now prices deltas from summary["by_model"] so multi-model runs persist correct cumulative cost (falls back to CLI model). - read/write usage now resolve the session via the same project->global store search order as rehydrate_session, so usage for globally-stored sessions accumulates into the right record after resume instead of splitting. Adds 2 tests (12 total, all passing). Co-authored-by: Mervin Praison <MervinPraison@users.noreply.github.com>
|
@claude You are the FINAL architecture reviewer. If the branch is under MervinPraison/PraisonAI (not a fork), you are able to make modifications to this branch and push directly. SCOPE: Focus ONLY on Python packages (praisonaiagents, praisonai). Do NOT modify praisonai-rust or praisonai-ts. Read ALL comments above from Gemini, Qodo, CodeRabbit, and Copilot carefully before responding. Phase 1: Review per AGENTS.md
Phase 2: FIX Valid Issues Phase 3: Final Verdict |
|
Claude finished @MervinPraison's task —— View job Final Architecture Review — PR #2424Verdict: ✅ Approve. All valid Major issues were already fixed in commit Phase 1 — Architecture (AGENTS.md)
Phase 2 — Verified fixes (already on branch, commit
|
| Reviewer | Severity | Issue | Status |
|---|---|---|---|
| CodeRabbit | 🟠 Major | Cost priced with single CLI model — wrong for multi-model runs | ✅ Fixed & verified — accumulate_session_usage iterates summary["by_model"], prices each model, sums deltas, falls back to CLI model when empty (project_sessions.py:303-322). |
| CodeRabbit / Greptile | 🟠 Major | Read/write always used the project store, but resume can restore from the global store → cumulative usage splits/resets after resume | ✅ Fixed & verified — _resolve_usage_store() mirrors rehydrate_session's project→global search order; both read and write target the resolved store (project_sessions.py:194-212, :337). |
Phase 3 — Re-validated deferred findings (no change, correctly out of scope)
- Greptile "Warm Runtime Skips Accounting" — Not a valid bug. Warm-runtime eligibility (
run.py:991-994) requiresno_save=Trueand nosession/continue_session/forkflags. Withno_save,auto_save_namestaysNone(run.py:978), sosession_id or auto_save_nameisNone→ there is no session to record and_record_session_usagecorrectly no-ops. The runtime executes out-of-process so the in-process collector isn't populated/left dirty. - CodeRabbit / Greptile "Global
manager.list()lacks Tokens/Cost" — Minor, correctly deferred. CoreSessionMetadata.to_dict()(state/sessions.py:35) never carries usage, and that path never persists usage. Table output degrades gracefully to-viagetattr(..., 0)(session.py:152-153); JSON simply omits the keys. Plumbing usage into the core storage model for a path that doesn't write usage would be scope creep touching core. Out of scope for this PR.
Tests
tests/unit/test_session_resume.py ............ 12 passed in 0.26s
Scope honored: Python wrapper (praisonai) only; no praisonai-rust / praisonai-ts touched. No new commits needed — branch is already in a mergeable state.
Job run · Branch claude/issue-2421-20260628-0841
| total_tokens = getattr(s, "total_tokens", 0) or 0 | ||
| cost = getattr(s, "cost", 0.0) or 0.0 |
There was a problem hiding this comment.
--all and backend-specific listings still use SessionManager.list(), which returns SessionMetadata objects instead of the local wrapper created above. Those objects do not expose total_tokens or cost, so these lookups fall back to zero and the new columns render as - even when usage was persisted for the session. JSON output has the same gap because SessionMetadata.to_dict() does not include the usage fields. The global/backend list path needs to attach the persisted usage before rendering.
Fixes #2421
Summary
Per-session cumulative cost and token accounting is now a first-class,
persisted, displayed field on the CLI. Previously the core LLM layer returned
per-call usage but the wrapper never aggregated, persisted, or showed how much a
session had consumed.
What changed (wrapper-only)
cli/state/project_sessions.py— new helpers:accumulate_session_usage()reads the per-call usage already aggregated bypraisonaiagents.telemetry.token_collector, prices it via the existingcli/features/cost_tracker.get_pricing(), and merges the deltas into thesession's persisted
usagemetadata (also mirroring the flattotal_tokens/costfields thatlist_sessionsalready surfaces). Resetsthe collector after each run to avoid double-counting.
read_session_usage()/format_usage_footer()for reading + rendering.cli/commands/run.py— after each prompt run with an active session,accumulate usage and print a compact footer (e.g.
1,240 in / 3,980 out · $0.0140).Best-effort and silent in
--jsonmode.cli/session/resume.py—RehydratedSessiongains ausagefield socumulative totals are restored on
--continue/--session <id>and keepaccumulating instead of resetting.
cli/commands/session.py—praisonai session listgains Tokens andCost columns; the resume panel shows a usage summary line.
Layer placement
This is the wrapper, exactly as the issue specifies. Core already exposes
per-call usage (
token_collector); aggregation belongs to the sessionabstraction (
.praisonai/sessions/, resume,praisonai session). No newdependencies; everything is lazy-imported and best-effort so usage accounting
never breaks a run.
Testing
Added 4 focused tests (10 total in
test_session_resume.py, all passing):Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Documentation