Skip to content

Log aggregate token usage to a local file - #81

Closed
AminMahpour wants to merge 4 commits into
masterfrom
dev/token-usage-file-logging
Closed

Log aggregate token usage to a local file#81
AminMahpour wants to merge 4 commits into
masterfrom
dev/token-usage-file-logging

Conversation

@AminMahpour

Copy link
Copy Markdown
Owner

Summary

  • add an always-on, request-scoped LangChain callback that appends aggregate token usage to .files/token-usage.jsonl
  • cover Chainlit, CLI, TUI, API, reflection, and nested model calls through the shared run configuration
  • prefer normalized AIMessage.usage_metadata, with provider token-usage fallback and malformed-value filtering
  • serialize concurrent appends, write once per root request, and keep logging failures from interrupting agent runs
  • document the privacy-safe JSONL contract and add focused plus interface-level regression coverage

Impact

Each top-level request now produces one local JSONL record containing only timestamp, thread/request identifiers, completion status, and aggregate input/output/total token counts. Prompts, responses, tool arguments, and per-model breakdowns are never logged. The log follows the active runtime project root and remains ignored by Git under .files/.

Validation

  • .venv/bin/python -m pytest -q — 343 passed, 2 pre-existing dependency deprecation warnings
  • real LangChain callback lifecycle probe — one record with 11 input + 4 output = 15 total
  • .venv/bin/python -m compileall -q chainagents tests/test_token_usage_logging.py
  • git diff --check
  • read-only code review and follow-up review — no remaining Critical, Important, or Minor issues

@AminMahpour
AminMahpour marked this pull request as ready for review July 30, 2026 14:55

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2fabf4a106

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +3690 to +3694
"callbacks": [
build_token_usage_callback_handler(
thread_id=thread_id,
project_root=project_root,
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Attach the usage callback to Agent Server runs

When requests use the documented langgraph dev deployment, including the configured async-researcher, they bypass this run-config builder: chainagents/langgraph/app.py exports raw graphs and Agent Protocol calls cannot inherit this in-process callback. Consequently, long-running background subagent model calls produce no record in .files/token-usage.jsonl, despite the new aggregate log claiming nested-call coverage. Instrument the exported Agent Server graphs or their request configuration as well.

Useful? React with 👍 / 👎.

Comment thread chainagents/runtime/token_usage.py Outdated
Comment on lines +97 to +100
total_tokens = (
input_tokens + output_tokens
if raw_total_tokens is None
else (_token_count(raw_total_tokens) or 0)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Fall back to component totals when the reported total is invalid

When a provider returns valid input/output counts but a malformed total_tokens value, this branch records the total as zero rather than filtering the malformed field and deriving input_tokens + output_tokens. For example, {input_tokens: 5, output_tokens: 2, total_tokens: "7"} produces an internally inconsistent record with a total of 0, corrupting aggregate usage data.

Useful? React with 👍 / 👎.

Comment thread chainagents/runtime/token_usage.py Outdated
Comment on lines +143 to +144
"thread_id": self._thread_id,
"request_id": str(run_id),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep cross-user thread IDs outside the agent workspace

In a shared stateful deployment, this writes every user's thread ID into .files/token-usage.jsonl under the project root, while build_deepagent_backend() exposes that same root through the agent-readable /workspace/ route. Any user can therefore ask their agent to read the log and obtain other users' checkpoint identifiers; if the unauthenticated API is also reachable, its caller-controlled thread_id can then be used to resume or modify those conversations. Store the log outside the agent filesystem, isolate it per user, or avoid recording reusable thread IDs.

Useful? React with 👍 / 👎.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 38a426473a

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment on lines +78 to +80
if isinstance(message, AIMessage) and message.usage_metadata:
usage = message.usage_metadata
if usage is None and isinstance(response.llm_output, dict):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Enable usage metadata for OpenAI-compatible streams

When provider = "openai_compatible" uses the normal streaming path, every interface calls astream_events, but build_model() supplies a custom base_url, causing ChatOpenAI to leave stream_usage disabled by default. Endpoints such as OpenAI that emit streaming usage only when stream_options.include_usage is requested therefore produce an LLMResult with neither usage_metadata nor fallback token_usage, so this handler silently writes zero tokens for the entire request. Explicitly enable streaming usage where supported or add another accounting path.

Useful? React with 👍 / 👎.

Comment thread chainagents/runtime/token_usage.py Outdated
Comment on lines +128 to +131
"""Append aggregate usage when the root request fails."""
if parent_run_id is not None:
return
self._write_once(run_id=run_id, status="error")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Finalize usage records when streamed requests are cancelled

When a user stops a Chainlit run, a CLI invocation is interrupted, or an API client disconnects, the interfaces close the astream_events generator; closing a LangGraph event stream does not emit a root on_chain_end or on_chain_error. Because these are the only paths that call _write_once, the request produces no log record and any tokens accumulated before cancellation are lost. Add explicit cancellation/finalization handling around streamed invocations.

Useful? React with 👍 / 👎.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8cc6b21f49

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread chainagents/runtime/core.py Outdated
"api_key": api_key or "deepagent",
"temperature": resolved_profile.temperature,
"disable_streaming": resolved_profile.disable_streaming,
"stream_usage": True,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Avoid forcing usage options on every compatible endpoint

When provider = "openai_compatible" points to a server that does not implement OpenAI's stream_options.include_usage extension, every normally streamed request now sends that unsupported parameter and can fail with a provider-side 400 instead of producing a response. LangChain deliberately leaves stream_usage disabled for custom base URLs because many compatible APIs lack this extension, so usage collection needs an endpoint capability/opt-in or a fallback rather than enabling it unconditionally.

Useful? React with 👍 / 👎.

Comment thread chainagents/runtime/token_usage.py Outdated
"""Append aggregate usage when the root request fails."""
if parent_run_id is not None:
return
self._write_once(run_id=run_id, status="error")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Classify root cancellation errors as cancelled

When a streamed LangGraph run is cancelled while a node is active, LangGraph invokes the root on_chain_error with asyncio.CancelledError before the interface catches cancellation, so this writes an error record and sets _written; the later finalize_cancelled_token_usage() call cannot correct it. Fresh evidence after the cancellation fix is that cancelling a real StateGraph.astream_events run follows exactly this callback order, causing user stops and client disconnects to be misclassified as failures.

Useful? React with 👍 / 👎.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant