Log aggregate token usage to a local file - #81
Conversation
There was a problem hiding this comment.
💡 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".
| "callbacks": [ | ||
| build_token_usage_callback_handler( | ||
| thread_id=thread_id, | ||
| project_root=project_root, | ||
| ) |
There was a problem hiding this comment.
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 👍 / 👎.
| total_tokens = ( | ||
| input_tokens + output_tokens | ||
| if raw_total_tokens is None | ||
| else (_token_count(raw_total_tokens) or 0) |
There was a problem hiding this comment.
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 👍 / 👎.
| "thread_id": self._thread_id, | ||
| "request_id": str(run_id), |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
💡 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".
| if isinstance(message, AIMessage) and message.usage_metadata: | ||
| usage = message.usage_metadata | ||
| if usage is None and isinstance(response.llm_output, dict): |
There was a problem hiding this comment.
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 👍 / 👎.
| """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") |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
💡 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".
| "api_key": api_key or "deepagent", | ||
| "temperature": resolved_profile.temperature, | ||
| "disable_streaming": resolved_profile.disable_streaming, | ||
| "stream_usage": True, |
There was a problem hiding this comment.
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 👍 / 👎.
| """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") |
There was a problem hiding this comment.
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 👍 / 👎.
Summary
.files/token-usage.jsonlAIMessage.usage_metadata, with provider token-usage fallback and malformed-value filteringImpact
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.venv/bin/python -m compileall -q chainagents tests/test_token_usage_logging.pygit diff --check