feat(langchain): support custom token_counter in ContextEditingMiddleware

[syyy44]

Description

ContextEditingMiddleware hardcodes its token counter inside wrap_model_call / awrap_model_call: token_count_method only allows "approximate" (which calls count_tokens_approximately with its defaults) or "model". There is no way to inject a custom counter, even though:

• the module already defines a public TokenCounter type alias, and ContextEdit.apply() already receives count_tokens as a parameter — the extension point exists everywhere except the middleware's public API;
• the sibling SummarizationMiddleware already accepts token_counter: TokenCounter = count_tokens_approximately, so this change simply brings the two middlewares to API parity.

Why this matters: count_tokens_approximately assumes ~4 chars/token, which fits English prose but systematically underestimates CJK-heavy workloads by 2–3x (we measured 1.4–2.1 chars/token against provider-reported usage_metadata on Chinese financial filings and HTML tables). With the estimate reading 2–3x low, ClearToolUsesEdit.trigger fires far later than configured, silently defeating the purpose of threshold-based context management. The "model" counting mode is not a general fallback: many OpenAI-compatible providers do not implement get_num_tokens_from_messages.

Changes

• Add optional token_counter: TokenCounter | None = None to ContextEditingMiddleware.__init__; when provided it takes precedence over token_count_method (mirroring SummarizationMiddleware).
• Deduplicate the counter-selection logic shared by the sync/async paths into a private _resolve_token_counter() helper.
• Unit tests covering: a custom counter drives trigger evaluation (sync and async), and it takes precedence over token_count_method.

Example:

from functools import partial
from langchain_core.messages.utils import count_tokens_approximately

middleware = ContextEditingMiddleware(
    edits=[ClearToolUsesEdit(trigger=120_000, keep=12)],
    token_counter=partial(count_tokens_approximately, chars_per_token=2.0),
)

Backwards compatible: no behavior change unless the new argument is passed. No new dependencies.

Issue

Fixes #39757

[codspeed-hq]

Merging this PR will not alter performance

✅ 2 untouched benchmarks
⏩ 15 skipped benchmarks¹

---

_{Comparing syyy44:feat/context-editing-token-counter (f966f53) with master (e92c6db)²}

Footnotes

1. 15 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

2. No successful run was found on master (2019bf5) during the generation of this report, so e92c6db was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[open-swe]

✅ Open SWE Review: No issues found

Open SWE reviewed this PR and found no potential bugs to report.

Open in Web • View Open SWE trace