mirror of
https://github.com/langchain-ai/deepagents.git
synced 2026-07-22 01:35:28 -04:00
fix(code): persist _context_tokens via after_model middleware (#3496)
`_context_tokens` was being written from the CLI via a client-side `aupdate_state` call after every turn. Locally we suppressed the resulting standalone `UpdateState` LangSmith run with `tracing_context(enabled=False)`, but over an HTTP `RemoteGraph` the server opens the run and the client cannot opt out — so we skipped the persist entirely on remote agents and accepted stale token counts on resume. This moves the LLM-usage path into `TokenStateMiddleware.aafter_model`, which reads the latest `AIMessage.usage_metadata` and emits a normal state update on the model node's own checkpoint. No separate `UpdateState` run, no client aggregation, no local-vs-remote divergence. The offload path folds the post-offload count into the same `aupdate_state` call that already writes `_summarization_event`, so it rides on an existing (intentionally traced) state transition. `_persist_context_tokens` is kept only for the interrupt-cleanup path where `after_model` never got to run on the partial turn. _Opened collaboratively by Sydney Runkle and open-swe._ --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com> Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Co-authored-by: Mason Daugherty <github@mdrxy.com>
This commit is contained in:
open-swe[bot]
@@ -220,95 +220,7 @@ Always use the latest generally available models when referencing LLMs in docstr
|
||||
|
||||
### Deep Agents Code (`libs/code/`)
|
||||
|
||||
`deepagents-code` is the interactive coding agent — the Textual REPL, headless `-x` mode, MCP integration, skills, sandbox bootstrap, and slash-command surface. Forked from `deepagents-cli` at the 0.1.0 split; treat the layout, conventions, and CI checks below as authoritative for this package.
|
||||
|
||||
#### Textual (terminal UI framework)
|
||||
|
||||
`deepagents-code` uses [Textual](https://textual.textualize.io/).
|
||||
|
||||
**Key Textual resources:**
|
||||
|
||||
- **Guide:** https://textual.textualize.io/guide/
|
||||
- **Widget gallery:** https://textual.textualize.io/widget_gallery/
|
||||
- **CSS reference:** https://textual.textualize.io/styles/
|
||||
- **API reference:** https://textual.textualize.io/api/
|
||||
|
||||
**Styled text in widgets:**
|
||||
|
||||
Prefer Textual's `Content` (`textual.content`) over Rich's `Text` for widget rendering. `Content` is immutable (like `str`) and integrates natively with Textual's rendering pipeline. Rich `Text` is still correct for code that renders via Rich's `Console.print()` (e.g., `non_interactive.py`, `main.py`).
|
||||
|
||||
IMPORTANT: `Content` requires **Textual's** `Style` (`textual.style.Style`) for rendering, not Rich's `Style` (`rich.style.Style`). Mixing Rich `Style` objects into `Content` spans will cause `TypeError` during widget rendering. String styles (`"bold cyan"`, `"dim"`) work for non-link styling. For links, use `TStyle(link=url)`.
|
||||
|
||||
**Never use f-string interpolation in Rich markup** (e.g., `f"[bold]{var}[/bold]"`). If `var` contains square brackets, the markup breaks or throws. Use `Content` methods instead:
|
||||
|
||||
- `Content.from_markup("[bold]$var[/bold]", var=value)` — for inline markup templates. `$var` substitution auto-escapes dynamic content. **Use when the variable is external/user-controlled** (tool args, file paths, user messages, diff content, error messages from exceptions).
|
||||
- `Content.styled(text, "bold")` — single style applied to plain text. No markup parsing. Use for static strings or when the variable is internal/trusted (glyphs, ints, enum-like status values). Avoid `Content.styled(f"..{var}..", style)` when `var` is user-controlled — while `styled` doesn't parse markup, the f-string pattern is fragile and inconsistent with the `from_markup` convention.
|
||||
- `Content.assemble("prefix: ", (text, "bold"), " ", other_content)` — for composing pre-built `Content` objects, `(text, style)` tuples, and plain strings. Plain strings are treated as plain text (no markup parsing). Use for structural composition, especially when parts use `TStyle(link=url)`.
|
||||
- `content.join(parts)` — like `str.join()` for `Content` objects.
|
||||
|
||||
**Decision rule:** if the value could ever come from outside the codebase (user input, tool output, API responses, file contents), use `from_markup` with `$var`. If it's a hardcoded string, glyph, or computed int, `styled` is fine.
|
||||
|
||||
**`App.notify()` defaults to `markup=True`:** Textual's `App.notify(message)` parses the message string as Rich markup by default. Any dynamic content (exception messages, file paths, user input, command strings) containing brackets `[]`, ANSI escape codes, or `=` will cause a `MarkupError` crash in Textual's Toast renderer. Always pass `markup=False` when the message contains f-string interpolated variables. Hardcoded string literals are safe with the default.
|
||||
|
||||
**Rich `console.print()` and number highlighting:**
|
||||
|
||||
`console.print()` defaults to `highlight=True`, which runs `ReprHighlighter` and auto-applies bold + cyan to any detected numbers. This visually overrides subtle styles like `dim` (bold cancels dim in most terminals). Pass `highlight=False` on any `console.print()` call where the content contains numbers and consistent dim/subtle styling matters.
|
||||
|
||||
**Textual patterns used in this codebase:**
|
||||
|
||||
- **Workers** (`@work` decorator) for async operations - see [Workers guide](https://textual.textualize.io/guide/workers/)
|
||||
- **Message passing** for widget communication - see [Events guide](https://textual.textualize.io/guide/events/)
|
||||
- **Reactive attributes** for state management - see [Reactivity guide](https://textual.textualize.io/guide/reactivity/)
|
||||
|
||||
**Testing Textual apps:**
|
||||
|
||||
- Use `textual.pilot` for async UI testing - see [Testing guide](https://textual.textualize.io/guide/testing/)
|
||||
- Snapshot testing available for visual regression - see repo `notes/snapshot_testing.md`
|
||||
|
||||
#### SDK dependency pin
|
||||
|
||||
`deepagents-code` pins an exact `deepagents==X.Y.Z` version in `libs/code/pyproject.toml`. When developing features that depend on new SDK functionality, bump this pin as part of the same PR. A CI check verifies the pin matches the current SDK version at release time (unless bypassed with `dangerous-skip-sdk-pin-check`).
|
||||
|
||||
#### Startup performance
|
||||
|
||||
`deepagents-code` must stay fast to launch. Never import heavy packages (e.g., `deepagents`, LangChain, LangGraph) at module level or in the argument-parsing path. These imports pull in large dependency trees and add seconds to every invocation, including trivial commands like `deepagents-code -v`.
|
||||
|
||||
- Keep top-level imports in `main.py` and other entry-point modules minimal.
|
||||
- Defer heavy imports to the point where they are actually needed (inside functions/methods).
|
||||
- To read another package's version without importing it, use `importlib.metadata.version("package-name")`.
|
||||
- Feature-gate checks on the startup hot path (before background workers fire) must be lightweight — env var lookups, small file reads. Never pull in expensive modules just to decide whether to skip a feature.
|
||||
- When adding logic that already exists elsewhere (e.g., editable-install detection), import the existing cached implementation rather than duplicating it.
|
||||
- Features that run shell commands silently must be opt-in, never default-enabled. Gate behind an explicit env var or config key.
|
||||
- Background workers that spawn subprocesses must set a timeout to avoid blocking indefinitely.
|
||||
|
||||
#### CLI help screen
|
||||
|
||||
The `deepagents-code --help` screen is hand-maintained in `ui.show_help()`, separate from the argparse definitions in `main.parse_args()`. When adding a new CLI flag, update **both** files. A drift-detection test (`test_args.TestHelpScreenDrift`) fails if a flag is registered in argparse but missing from the help screen.
|
||||
|
||||
#### Splash screen tips
|
||||
|
||||
When adding a user-facing CLI feature (new slash command, keybinding, workflow), add a corresponding tip to the `_TIPS` list in `libs/code/deepagents_code/widgets/welcome.py`. Tips are shown randomly on startup to help users discover features. Keep tips short and action-oriented (e.g., `"Press ctrl+x to compose prompts in your external editor"`).
|
||||
|
||||
#### Slash commands
|
||||
|
||||
Slash commands are defined as `SlashCommand` entries in the `COMMANDS` tuple in `libs/code/deepagents_code/command_registry.py`. Each entry declares the command name, description, `bypass_tier` (queue-bypass classification), optional `hidden_keywords` for fuzzy matching, and optional `aliases`. Bypass-tier frozensets and the `SLASH_COMMANDS` autocomplete list are derived automatically — no other file should hard-code command metadata.
|
||||
|
||||
To add a new slash command: (1) add a `SlashCommand` entry to `COMMANDS`, (2) set the appropriate `bypass_tier`, (3) add a handler branch in `_handle_command` in `app.py`, (4) run `make lint && make test` — the drift test will catch any mismatch.
|
||||
|
||||
#### Adding a new model provider
|
||||
|
||||
`deepagents-code` supports LangChain-based chat model providers as optional dependencies. To add a new provider, update these files (all entries alphabetically sorted):
|
||||
|
||||
1. `libs/code/deepagents_code/model_config.py` — add `"provider_name": "ENV_VAR_NAME"` to `PROVIDER_API_KEY_ENV`
|
||||
2. `libs/code/pyproject.toml` — add `provider = ["langchain-provider>=X.Y.Z,<N.0.0"]` to `[project.optional-dependencies]` and include it in the `all-providers` composite extra
|
||||
3. `libs/code/tests/unit_tests/test_model_config.py` — add `assert PROVIDER_API_KEY_ENV["provider_name"] == "ENV_VAR_NAME"` to `TestProviderApiKeyEnv.test_contains_major_providers`
|
||||
|
||||
**Not required** unless the provider's models have a distinctive name prefix (like `gpt-*`, `claude*`, `gemini*`):
|
||||
|
||||
- `detect_provider()` in `config.py` — only needed for auto-detection from bare model names
|
||||
- `Settings.has_*` property in `config.py` — only needed if referenced by `detect_provider()` fallback logic
|
||||
|
||||
Model discovery, credential checking, and UI integration are automatic once `PROVIDER_API_KEY_ENV` is populated and the `langchain-*` package is installed.
|
||||
See `libs/code/AGENTS.md` for package-specific guidance — Textual, startup performance, slash commands, model providers, SDK pin, help-screen drift.
|
||||
|
||||
### Deep Agents CLI (`libs/cli/`)
|
||||
|
||||
|
||||
@@ -0,0 +1,95 @@
|
||||
# `libs/code` agent guide
|
||||
|
||||
`deepagents-code` is the interactive coding agent — the Textual REPL, headless `-x` mode, MCP integration, skills, sandbox bootstrap, and slash-command surface. Forked from `deepagents-cli` at the 0.1.0 split.
|
||||
|
||||
For monorepo-wide conventions (commit titles, lint, testing, docs, CI, benchmarks), see the root `AGENTS.md`.
|
||||
|
||||
## Textual (terminal UI framework)
|
||||
|
||||
`deepagents-code` uses [Textual](https://textual.textualize.io/).
|
||||
|
||||
**Key Textual resources:**
|
||||
|
||||
- **Guide:** https://textual.textualize.io/guide/
|
||||
- **Widget gallery:** https://textual.textualize.io/widget_gallery/
|
||||
- **CSS reference:** https://textual.textualize.io/styles/
|
||||
- **API reference:** https://textual.textualize.io/api/
|
||||
|
||||
### Styled text in widgets
|
||||
|
||||
Prefer Textual's `Content` (`textual.content`) over Rich's `Text` for widget rendering. `Content` is immutable (like `str`) and integrates natively with Textual's rendering pipeline. Rich `Text` is still correct for code that renders via Rich's `Console.print()` (e.g., `non_interactive.py`, `main.py`).
|
||||
|
||||
IMPORTANT: `Content` requires **Textual's** `Style` (`textual.style.Style`) for rendering, not Rich's `Style` (`rich.style.Style`). Mixing Rich `Style` objects into `Content` spans will cause `TypeError` during widget rendering. String styles (`"bold cyan"`, `"dim"`) work for non-link styling. For links, use `TStyle(link=url)`.
|
||||
|
||||
**Never use f-string interpolation in Rich markup** (e.g., `f"[bold]{var}[/bold]"`). If `var` contains square brackets, the markup breaks or throws. Use `Content` methods instead:
|
||||
|
||||
- `Content.from_markup("[bold]$var[/bold]", var=value)` — for inline markup templates. `$var` substitution auto-escapes dynamic content. **Use when the variable is external/user-controlled** (tool args, file paths, user messages, diff content, error messages from exceptions).
|
||||
- `Content.styled(text, "bold")` — single style applied to plain text. No markup parsing. Use for static strings or when the variable is internal/trusted (glyphs, ints, enum-like status values). Avoid `Content.styled(f"..{var}..", style)` when `var` is user-controlled — while `styled` doesn't parse markup, the f-string pattern is fragile and inconsistent with the `from_markup` convention.
|
||||
- `Content.assemble("prefix: ", (text, "bold"), " ", other_content)` — for composing pre-built `Content` objects, `(text, style)` tuples, and plain strings. Plain strings are treated as plain text (no markup parsing). Use for structural composition, especially when parts use `TStyle(link=url)`.
|
||||
- `content.join(parts)` — like `str.join()` for `Content` objects.
|
||||
|
||||
**Decision rule:** if the value could ever come from outside the codebase (user input, tool output, API responses, file contents), use `from_markup` with `$var`. If it's a hardcoded string, glyph, or computed int, `styled` is fine.
|
||||
|
||||
### `App.notify()` defaults to `markup=True`
|
||||
|
||||
Textual's `App.notify(message)` parses the message string as Rich markup by default. Any dynamic content (exception messages, file paths, user input, command strings) containing brackets `[]`, ANSI escape codes, or `=` will cause a `MarkupError` crash in Textual's Toast renderer. Always pass `markup=False` when the message contains f-string interpolated variables. Hardcoded string literals are safe with the default.
|
||||
|
||||
### Rich `console.print()` and number highlighting
|
||||
|
||||
`console.print()` defaults to `highlight=True`, which runs `ReprHighlighter` and auto-applies bold + cyan to any detected numbers. This visually overrides subtle styles like `dim` (bold cancels dim in most terminals). Pass `highlight=False` on any `console.print()` call where the content contains numbers and consistent dim/subtle styling matters.
|
||||
|
||||
### Textual patterns used in this codebase
|
||||
|
||||
- **Workers** (`@work` decorator) for async operations - see [Workers guide](https://textual.textualize.io/guide/workers/)
|
||||
- **Message passing** for widget communication - see [Events guide](https://textual.textualize.io/guide/events/)
|
||||
- **Reactive attributes** for state management - see [Reactivity guide](https://textual.textualize.io/guide/reactivity/)
|
||||
|
||||
### Testing Textual apps
|
||||
|
||||
- Use `textual.pilot` for async UI testing - see [Testing guide](https://textual.textualize.io/guide/testing/)
|
||||
- Snapshot testing available for visual regression - see repo `notes/snapshot_testing.md`
|
||||
|
||||
## SDK dependency pin
|
||||
|
||||
`deepagents-code` pins an exact `deepagents==X.Y.Z` version in `pyproject.toml`. When developing features that depend on new SDK functionality, bump this pin as part of the same PR. A CI check verifies the pin matches the current SDK version at release time (unless bypassed with `dangerous-skip-sdk-pin-check`).
|
||||
|
||||
## Startup performance
|
||||
|
||||
`deepagents-code` must stay fast to launch. Never import heavy packages (e.g., `deepagents`, LangChain, LangGraph) at module level or in the argument-parsing path. These imports pull in large dependency trees and add seconds to every invocation, including trivial commands like `deepagents-code -v`.
|
||||
|
||||
- Keep top-level imports in `main.py` and other entry-point modules minimal.
|
||||
- Defer heavy imports to the point where they are actually needed (inside functions/methods).
|
||||
- To read another package's version without importing it, use `importlib.metadata.version("package-name")`.
|
||||
- Feature-gate checks on the startup hot path (before background workers fire) must be lightweight — env var lookups, small file reads. Never pull in expensive modules just to decide whether to skip a feature.
|
||||
- When adding logic that already exists elsewhere (e.g., editable-install detection), import the existing cached implementation rather than duplicating it.
|
||||
- Features that run shell commands silently must be opt-in, never default-enabled. Gate behind an explicit env var or config key.
|
||||
- Background workers that spawn subprocesses must set a timeout to avoid blocking indefinitely.
|
||||
|
||||
## CLI help screen
|
||||
|
||||
The `deepagents-code --help` screen is hand-maintained in `ui.show_help()`, separate from the argparse definitions in `main.parse_args()`. When adding a new CLI flag, update **both** files. A drift-detection test (`test_args.TestHelpScreenDrift`) fails if a flag is registered in argparse but missing from the help screen.
|
||||
|
||||
## Splash screen tips
|
||||
|
||||
When adding a user-facing CLI feature (new slash command, keybinding, workflow), add a corresponding tip to the `_TIPS` list in `deepagents_code/widgets/welcome.py`. Tips are shown randomly on startup to help users discover features. Keep tips short and action-oriented (e.g., `"Press ctrl+x to compose prompts in your external editor"`).
|
||||
|
||||
## Slash commands
|
||||
|
||||
Slash commands are defined as `SlashCommand` entries in the `COMMANDS` tuple in `deepagents_code/command_registry.py`. Each entry declares the command name, description, `bypass_tier` (queue-bypass classification), optional `hidden_keywords` for fuzzy matching, and optional `aliases`. Bypass-tier frozensets and the `SLASH_COMMANDS` autocomplete list are derived automatically — no other file should hard-code command metadata.
|
||||
|
||||
To add a new slash command: (1) add a `SlashCommand` entry to `COMMANDS`, (2) set the appropriate `bypass_tier`, (3) add a handler branch in `_handle_command` in `app.py`, (4) run `make lint && make test` — the drift test will catch any mismatch.
|
||||
|
||||
## Adding a new model provider
|
||||
|
||||
`deepagents-code` supports LangChain-based chat model providers as optional dependencies. To add a new provider, update these files (all entries alphabetically sorted):
|
||||
|
||||
1. `deepagents_code/model_config.py` — add `"provider_name": "ENV_VAR_NAME"` to `PROVIDER_API_KEY_ENV`
|
||||
2. `pyproject.toml` — add `provider = ["langchain-provider>=X.Y.Z,<N.0.0"]` to `[project.optional-dependencies]` and include it in the `all-providers` composite extra
|
||||
3. `tests/unit_tests/test_model_config.py` — add `assert PROVIDER_API_KEY_ENV["provider_name"] == "ENV_VAR_NAME"` to `TestProviderApiKeyEnv.test_contains_major_providers`
|
||||
|
||||
**Not required** unless the provider's models have a distinctive name prefix (like `gpt-*`, `claude*`, `gemini*`):
|
||||
|
||||
- `detect_provider()` in `config.py` — only needed for auto-detection from bare model names
|
||||
- `Settings.has_*` property in `config.py` — only needed if referenced by `detect_provider()` fallback logic
|
||||
|
||||
Model discovery, credential checking, and UI integration are automatic once `PROVIDER_API_KEY_ENV` is populated and the `langchain-*` package is installed.
|
||||
@@ -1100,9 +1100,9 @@ def create_cli_agent(
|
||||
agent_middleware = []
|
||||
agent_middleware.append(ConfigurableModelMiddleware())
|
||||
|
||||
# Token state: adds _context_tokens to graph state (checkpointed, not
|
||||
# passed to model). Must be registered before any middleware that might
|
||||
# read the channel.
|
||||
# Token state: declares the `_context_tokens` channel and writes it from
|
||||
# `after_model` based on the latest `AIMessage.usage_metadata`. The CLI
|
||||
# reads it back from `state_values` on thread resume.
|
||||
from deepagents_code.token_state import TokenStateMiddleware
|
||||
|
||||
agent_middleware.append(TokenStateMiddleware())
|
||||
|
||||
@@ -5661,9 +5661,15 @@ class DeepAgentsApp(App):
|
||||
|
||||
# Intentionally traced: the summarization event is a meaningful state
|
||||
# transition that should surface in LangSmith alongside real agent turns.
|
||||
# The new `_context_tokens` count rides along on the same update so it
|
||||
# shares a checkpoint with the offload and doesn't create a separate
|
||||
# `UpdateState` run.
|
||||
await self._agent.aupdate_state(
|
||||
config,
|
||||
{"_summarization_event": result.new_event},
|
||||
{
|
||||
"_summarization_event": result.new_event,
|
||||
"_context_tokens": result.tokens_after,
|
||||
},
|
||||
)
|
||||
|
||||
before = format_token_count(result.tokens_before)
|
||||
@@ -5679,9 +5685,6 @@ class DeepAgentsApp(App):
|
||||
)
|
||||
|
||||
self._on_tokens_update(result.tokens_after)
|
||||
from deepagents_code.textual_adapter import _persist_context_tokens
|
||||
|
||||
await _persist_context_tokens(self._agent, config, result.tokens_after)
|
||||
|
||||
except OffloadModelError as exc:
|
||||
logger.warning("Offload model creation failed: %s", exc, exc_info=True)
|
||||
|
||||
@@ -1306,13 +1306,13 @@ async def execute_task_textual(
|
||||
)
|
||||
await adapter._mount_message(AppMessage(message))
|
||||
turn_stats.wall_time_seconds = time.monotonic() - start_time
|
||||
await _report_and_persist_tokens(
|
||||
# Model call already completed (HITL interrupt fires after
|
||||
# the model node); `TokenStateMiddleware.after_model`
|
||||
# persisted the count, so only refresh UI here.
|
||||
_report_tokens(
|
||||
adapter,
|
||||
agent,
|
||||
config,
|
||||
captured_input_tokens,
|
||||
captured_output_tokens,
|
||||
shield=True,
|
||||
)
|
||||
return turn_stats
|
||||
|
||||
@@ -1334,12 +1334,11 @@ async def execute_task_textual(
|
||||
)
|
||||
return turn_stats
|
||||
|
||||
# Update token count and return stats
|
||||
# Update token count and return stats. Persistence is handled inside the
|
||||
# graph by `TokenStateMiddleware.after_model`, so this only refreshes UI.
|
||||
turn_stats.wall_time_seconds = time.monotonic() - start_time
|
||||
await _report_and_persist_tokens(
|
||||
_report_tokens(
|
||||
adapter,
|
||||
agent,
|
||||
config,
|
||||
captured_input_tokens,
|
||||
captured_output_tokens,
|
||||
)
|
||||
@@ -1408,7 +1407,15 @@ async def _handle_interrupt_cleanup(
|
||||
content="[SYSTEM] Task interrupted by user. "
|
||||
"Previous operation was cancelled."
|
||||
)
|
||||
await agent.aupdate_state(config, {"messages": [cancellation_msg]})
|
||||
cancellation_values: dict[str, Any] = {"messages": [cancellation_msg]}
|
||||
# Piggy-back the latest token count on this already-required write
|
||||
# instead of issuing a separate `aupdate_state`. `after_model` never
|
||||
# ran on the partial turn, so without this the count would be stale
|
||||
# on resume.
|
||||
captured_total = captured_input_tokens + captured_output_tokens
|
||||
if captured_total:
|
||||
cancellation_values["_context_tokens"] = captured_total
|
||||
await agent.aupdate_state(config, cancellation_values)
|
||||
except (httpx.TransportError, httpx.TimeoutException) as e:
|
||||
logger.warning("Could not save interrupted state (network): %s", e)
|
||||
except Exception as exc: # interrupt cleanup must not propagate
|
||||
@@ -1435,102 +1442,37 @@ async def _handle_interrupt_cleanup(
|
||||
approximate = interrupted_msg is not None
|
||||
|
||||
turn_stats.wall_time_seconds = time.monotonic() - start_time
|
||||
await _report_and_persist_tokens(
|
||||
_report_tokens(
|
||||
adapter,
|
||||
agent,
|
||||
config,
|
||||
captured_input_tokens,
|
||||
captured_output_tokens,
|
||||
shield=True,
|
||||
approximate=approximate,
|
||||
)
|
||||
|
||||
|
||||
async def _persist_context_tokens(
|
||||
agent: Any, # noqa: ANN401 # Dynamic agent graph type
|
||||
config: RunnableConfig,
|
||||
tokens: int,
|
||||
) -> None:
|
||||
"""Best-effort persist of the context token count into graph state.
|
||||
|
||||
For local agents, the `aupdate_state` call is wrapped in
|
||||
`tracing_context(enabled=False)` so this purely-internal bookkeeping write
|
||||
does not surface as a separate `UpdateState` run in LangSmith.
|
||||
`_context_tokens` is already marked `PrivateStateAttr`, but `aupdate_state`
|
||||
itself creates its own traced run that would otherwise clutter the project's
|
||||
traces.
|
||||
|
||||
For remote agents (e.g. a `langgraph dev` server), `aupdate_state` goes
|
||||
over HTTP and the server creates its own LangSmith trace that the client
|
||||
cannot suppress with `tracing_context`. Skipping the call is safe because
|
||||
persistence is best-effort — a stale count on resume is acceptable.
|
||||
|
||||
Args:
|
||||
agent: The LangGraph agent (must support `aupdate_state`).
|
||||
config: Runnable config with `thread_id`.
|
||||
tokens: Total context tokens to persist.
|
||||
"""
|
||||
from deepagents_code.remote_client import RemoteAgent
|
||||
|
||||
if isinstance(agent, RemoteAgent):
|
||||
return
|
||||
|
||||
from langsmith import tracing_context
|
||||
|
||||
try:
|
||||
with tracing_context(enabled=False):
|
||||
await agent.aupdate_state(config, {"_context_tokens": tokens})
|
||||
except (httpx.TransportError, httpx.TimeoutException) as e:
|
||||
logger.warning(
|
||||
"Could not persist _context_tokens=%d (network): %s; "
|
||||
"token count may be stale on resume",
|
||||
tokens,
|
||||
e,
|
||||
)
|
||||
except Exception: # non-critical; stale count on resume is acceptable
|
||||
logger.warning(
|
||||
"Failed to persist _context_tokens=%d; token count may be stale on resume",
|
||||
tokens,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
|
||||
async def _report_and_persist_tokens(
|
||||
def _report_tokens(
|
||||
adapter: TextualUIAdapter,
|
||||
agent: Any, # noqa: ANN401 # Dynamic agent graph type
|
||||
config: RunnableConfig,
|
||||
captured_input_tokens: int,
|
||||
captured_output_tokens: int,
|
||||
*,
|
||||
shield: bool = False,
|
||||
approximate: bool = False,
|
||||
) -> None:
|
||||
"""Update the token display and best-effort persist to graph state.
|
||||
"""Refresh the token-count UI display.
|
||||
|
||||
Persistence into graph state is owned by `TokenStateMiddleware.after_model`
|
||||
(normal turns), `_handle_offload` (offload turns), and the interrupt-cleanup
|
||||
`aupdate_state` write (partial turns) — never this helper.
|
||||
|
||||
Args:
|
||||
adapter: UI adapter with token callbacks.
|
||||
agent: The LangGraph agent.
|
||||
config: Runnable config with `thread_id` in its configurable dict.
|
||||
captured_input_tokens: Total input tokens captured during the turn.
|
||||
captured_output_tokens: Total output tokens captured during the turn.
|
||||
shield: When `True`, suppress exceptions and `CancelledError` from the
|
||||
persist call so that interrupt handlers can safely await this.
|
||||
approximate: When `True`, signal to the UI that the count is stale
|
||||
(e.g. after an interrupted generation) by appending "+".
|
||||
"""
|
||||
if captured_input_tokens or captured_output_tokens:
|
||||
if adapter._on_tokens_update:
|
||||
adapter._on_tokens_update(captured_input_tokens, approximate=approximate)
|
||||
if shield:
|
||||
try:
|
||||
await _persist_context_tokens(agent, config, captured_input_tokens)
|
||||
except (Exception, asyncio.CancelledError):
|
||||
logger.debug(
|
||||
"Token persist suppressed during interrupt cleanup",
|
||||
exc_info=True,
|
||||
)
|
||||
else:
|
||||
await _persist_context_tokens(agent, config, captured_input_tokens)
|
||||
elif adapter._on_tokens_show:
|
||||
adapter._on_tokens_show(approximate=approximate)
|
||||
|
||||
|
||||
@@ -1,21 +1,32 @@
|
||||
"""Middleware that adds a `_context_tokens` channel to the graph state.
|
||||
"""Middleware that tracks total context tokens in graph state.
|
||||
|
||||
The field is checkpointed (survives across sessions) but not passed to the
|
||||
model (`PrivateStateAttr`). The app writes the latest total-context token
|
||||
count here after every LLM response and context offload, and reads it back
|
||||
when resuming a thread so that `/tokens` and the status bar show accurate
|
||||
values immediately.
|
||||
Registers a `_context_tokens` channel (checkpointed, schema-private) and
|
||||
writes it from `after_model` based on the latest `AIMessage.usage_metadata`.
|
||||
|
||||
Persisting from inside the graph (rather than via a separate client-side
|
||||
`aupdate_state` call) keeps the write on the same checkpoint as the model
|
||||
response and avoids creating a standalone `UpdateState` run in LangSmith.
|
||||
It also works identically against local graphs and remote (HTTP) graphs.
|
||||
|
||||
The CLI reads `_context_tokens` back from `state_values` on thread resume
|
||||
so `/tokens` and the status bar show accurate values immediately, without
|
||||
having to replay or re-tokenize history.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, NotRequired
|
||||
from typing import TYPE_CHECKING, Annotated, Any, NotRequired
|
||||
|
||||
from langchain.agents.middleware.types import (
|
||||
AgentMiddleware,
|
||||
AgentState,
|
||||
ContextT,
|
||||
PrivateStateAttr,
|
||||
)
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
|
||||
class TokenTrackingState(AgentState):
|
||||
@@ -25,7 +36,53 @@ class TokenTrackingState(AgentState):
|
||||
"""Total context tokens reported by the model's last `usage_metadata`."""
|
||||
|
||||
|
||||
class TokenStateMiddleware(AgentMiddleware):
|
||||
"""Schema-only middleware that registers `_context_tokens` in the state schema."""
|
||||
def _extract_context_tokens(message: AIMessage) -> int | None:
|
||||
"""Return the context-token count from an AI message, or `None` if absent.
|
||||
|
||||
Prefers `input_tokens + output_tokens` when both are reported; falls back
|
||||
to `total_tokens` when the model only provides the aggregate.
|
||||
"""
|
||||
usage = getattr(message, "usage_metadata", None)
|
||||
if not usage:
|
||||
return None
|
||||
input_toks = usage.get("input_tokens", 0) or 0
|
||||
output_toks = usage.get("output_tokens", 0) or 0
|
||||
if input_toks or output_toks:
|
||||
return input_toks + output_toks
|
||||
total = usage.get("total_tokens", 0) or 0
|
||||
return total or None
|
||||
|
||||
|
||||
class TokenStateMiddleware(AgentMiddleware[TokenTrackingState, ContextT]):
|
||||
"""Persists the latest context-token count after each model call.
|
||||
|
||||
See the module docstring for why this rides the model node's checkpoint
|
||||
instead of a separate `aupdate_state` (avoids a standalone `UpdateState`
|
||||
run in LangSmith and works identically against remote graphs).
|
||||
"""
|
||||
|
||||
state_schema = TokenTrackingState
|
||||
|
||||
def after_model( # noqa: PLR6301 # AgentMiddleware hook must be an instance method.
|
||||
self,
|
||||
state: TokenTrackingState,
|
||||
runtime: Runtime[ContextT], # noqa: ARG002
|
||||
) -> dict[str, Any] | None:
|
||||
"""Write `_context_tokens` from the most recent `AIMessage.usage_metadata`.
|
||||
|
||||
Args:
|
||||
state: Current agent state; only `messages` is inspected.
|
||||
runtime: LangGraph runtime (unused; required by the hook signature).
|
||||
|
||||
Returns:
|
||||
State update `{"_context_tokens": <int>}` when usage is reported on
|
||||
the latest `AIMessage`; otherwise `None`.
|
||||
"""
|
||||
messages = state.get("messages") or []
|
||||
for msg in reversed(messages):
|
||||
if isinstance(msg, AIMessage):
|
||||
tokens = _extract_context_tokens(msg)
|
||||
if tokens is not None:
|
||||
return {"_context_tokens": tokens}
|
||||
return None
|
||||
return None
|
||||
|
||||
@@ -288,14 +288,17 @@ class TestOffloadSuccess:
|
||||
await pilot.pause()
|
||||
|
||||
mock_agent = app._agent
|
||||
# Two aupdate_state calls: _summarization_event + _context_tokens
|
||||
assert mock_agent.aupdate_state.call_count == 2 # type: ignore[union-attr]
|
||||
# Single aupdate_state call: _summarization_event + _context_tokens
|
||||
# ride along together to share a checkpoint and avoid a separate
|
||||
# standalone `UpdateState` LangSmith run.
|
||||
assert mock_agent.aupdate_state.call_count == 1 # type: ignore[union-attr]
|
||||
|
||||
update_values = mock_agent.aupdate_state.call_args_list[0][0][1] # type: ignore[union-attr]
|
||||
event = update_values["_summarization_event"]
|
||||
assert event["cutoff_index"] == 4
|
||||
assert event["summary_message"] is not None
|
||||
assert event["file_path"] == "/conversation_history/test-thread.md"
|
||||
assert update_values["_context_tokens"] == result.tokens_after
|
||||
|
||||
async def test_offload_shows_feedback_message(self) -> None:
|
||||
"""Should display feedback with message count and token change."""
|
||||
@@ -513,7 +516,7 @@ class TestReOffload:
|
||||
await pilot.pause()
|
||||
|
||||
mock_agent = app._agent
|
||||
assert mock_agent.aupdate_state.call_count == 2 # type: ignore[union-attr]
|
||||
assert mock_agent.aupdate_state.call_count == 1 # type: ignore[union-attr]
|
||||
|
||||
update_values = mock_agent.aupdate_state.call_args_list[0][0][1] # type: ignore[union-attr]
|
||||
event = update_values["_summarization_event"]
|
||||
@@ -598,7 +601,7 @@ class TestOffloadErrorHandling:
|
||||
await pilot.pause()
|
||||
|
||||
mock_agent = app._agent
|
||||
assert mock_agent.aupdate_state.call_count == 2 # type: ignore[union-attr]
|
||||
assert mock_agent.aupdate_state.call_count == 1 # type: ignore[union-attr]
|
||||
|
||||
update_values = mock_agent.aupdate_state.call_args_list[0][0][1] # type: ignore[union-attr]
|
||||
event = update_values["_summarization_event"]
|
||||
@@ -1190,8 +1193,8 @@ class TestOffloadProfileOverride:
|
||||
await app._handle_offload()
|
||||
await pilot.pause()
|
||||
|
||||
# State should have been updated (offload + _context_tokens)
|
||||
assert app._agent.aupdate_state.call_count == 2 # type: ignore[union-attr]
|
||||
# Single state update folds offload + _context_tokens together.
|
||||
assert app._agent.aupdate_state.call_count == 1 # type: ignore[union-attr]
|
||||
kwargs = mock_perform.call_args.kwargs
|
||||
assert kwargs["context_limit"] == 4096
|
||||
|
||||
|
||||
@@ -25,7 +25,6 @@ from deepagents_code.textual_adapter import (
|
||||
_build_interrupted_ai_message,
|
||||
_handle_interrupt_cleanup,
|
||||
_is_summarization_chunk,
|
||||
_persist_context_tokens,
|
||||
execute_task_textual,
|
||||
format_token_count,
|
||||
print_usage_table,
|
||||
@@ -310,43 +309,177 @@ class TestInterruptCleanup:
|
||||
)
|
||||
|
||||
|
||||
class TestPersistContextTokens:
|
||||
"""Tests for `_persist_context_tokens`."""
|
||||
class TestInterruptCleanupTokenPersist:
|
||||
"""`_context_tokens` rides on the cancellation `aupdate_state` write."""
|
||||
|
||||
async def test_skips_aupdate_state_for_remote_agent(self) -> None:
|
||||
"""Remote agents must not call aupdate_state for _context_tokens.
|
||||
async def test_includes_context_tokens_in_cancellation_update(self) -> None:
|
||||
"""The cancellation HumanMessage write carries the latest token count."""
|
||||
captured: list[dict[str, Any]] = []
|
||||
|
||||
When the agent is a RemoteAgent, aupdate_state goes over HTTP to the
|
||||
dev server, which creates its own LangSmith trace that the client
|
||||
cannot suppress with tracing_context(enabled=False). Skipping the call
|
||||
is safe because persistence is best-effort.
|
||||
async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029
|
||||
captured.append(values)
|
||||
|
||||
agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture))
|
||||
adapter = TextualUIAdapter(
|
||||
mount_message=AsyncMock(),
|
||||
update_status=_noop_status,
|
||||
request_approval=_mock_approval,
|
||||
set_spinner=AsyncMock(),
|
||||
set_active_message=MagicMock(),
|
||||
)
|
||||
|
||||
await _handle_interrupt_cleanup(
|
||||
adapter=adapter,
|
||||
agent=agent,
|
||||
config={"configurable": {"thread_id": "t-1"}},
|
||||
pending_text_by_namespace={},
|
||||
captured_input_tokens=4321,
|
||||
captured_output_tokens=0,
|
||||
turn_stats=SessionStats(),
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
# Only the cancellation write happens (no partial AI message in this test);
|
||||
# it carries both `messages` and `_context_tokens`.
|
||||
assert len(captured) == 1
|
||||
assert captured[0]["_context_tokens"] == 4321
|
||||
assert "messages" in captured[0]
|
||||
|
||||
async def test_omits_context_tokens_when_no_usage_captured(self) -> None:
|
||||
"""Zero tokens means we never saw `usage_metadata`; preserve the prior value."""
|
||||
captured: list[dict[str, Any]] = []
|
||||
|
||||
async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029
|
||||
captured.append(values)
|
||||
|
||||
agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture))
|
||||
adapter = TextualUIAdapter(
|
||||
mount_message=AsyncMock(),
|
||||
update_status=_noop_status,
|
||||
request_approval=_mock_approval,
|
||||
set_spinner=AsyncMock(),
|
||||
set_active_message=MagicMock(),
|
||||
)
|
||||
|
||||
await _handle_interrupt_cleanup(
|
||||
adapter=adapter,
|
||||
agent=agent,
|
||||
config={"configurable": {"thread_id": "t-1"}},
|
||||
pending_text_by_namespace={},
|
||||
captured_input_tokens=0,
|
||||
captured_output_tokens=0,
|
||||
turn_stats=SessionStats(),
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
assert "_context_tokens" not in captured[0]
|
||||
|
||||
async def test_includes_context_tokens_for_output_only_turn(self) -> None:
|
||||
"""Output-only AI turns (no input usage) still persist a count."""
|
||||
captured: list[dict[str, Any]] = []
|
||||
|
||||
async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029
|
||||
captured.append(values)
|
||||
|
||||
agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture))
|
||||
adapter = TextualUIAdapter(
|
||||
mount_message=AsyncMock(),
|
||||
update_status=_noop_status,
|
||||
request_approval=_mock_approval,
|
||||
set_spinner=AsyncMock(),
|
||||
set_active_message=MagicMock(),
|
||||
)
|
||||
|
||||
await _handle_interrupt_cleanup(
|
||||
adapter=adapter,
|
||||
agent=agent,
|
||||
config={"configurable": {"thread_id": "t-1"}},
|
||||
pending_text_by_namespace={},
|
||||
captured_input_tokens=0,
|
||||
captured_output_tokens=500,
|
||||
turn_stats=SessionStats(),
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
assert len(captured) == 1
|
||||
assert captured[0]["_context_tokens"] == 500
|
||||
|
||||
async def test_remote_agent_interrupt_write_carries_context_tokens(self) -> None:
|
||||
"""Remote agents are not skipped on the interrupt-cleanup write.
|
||||
|
||||
Locks in the deletion of the old `_persist_context_tokens` `RemoteAgent`
|
||||
short-circuit so a future refactor cannot silently re-introduce it.
|
||||
"""
|
||||
from deepagents_code.remote_client import RemoteAgent
|
||||
|
||||
agent = MagicMock(spec=RemoteAgent)
|
||||
agent.aupdate_state = AsyncMock()
|
||||
captured: list[dict[str, Any]] = []
|
||||
|
||||
await _persist_context_tokens(
|
||||
agent, {"configurable": {"thread_id": "t-1"}}, 1234
|
||||
async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029
|
||||
captured.append(values)
|
||||
|
||||
agent = MagicMock(spec=RemoteAgent)
|
||||
agent.aupdate_state = AsyncMock(side_effect=_capture)
|
||||
adapter = TextualUIAdapter(
|
||||
mount_message=AsyncMock(),
|
||||
update_status=_noop_status,
|
||||
request_approval=_mock_approval,
|
||||
set_spinner=AsyncMock(),
|
||||
set_active_message=MagicMock(),
|
||||
)
|
||||
|
||||
agent.aupdate_state.assert_not_called()
|
||||
await _handle_interrupt_cleanup(
|
||||
adapter=adapter,
|
||||
agent=agent,
|
||||
config={"configurable": {"thread_id": "t-1"}},
|
||||
pending_text_by_namespace={},
|
||||
captured_input_tokens=1234,
|
||||
captured_output_tokens=88,
|
||||
turn_stats=SessionStats(),
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
async def test_calls_aupdate_state_for_local_agent(self) -> None:
|
||||
"""Local agents should still call aupdate_state for _context_tokens."""
|
||||
called_with: list[object] = []
|
||||
assert isinstance(agent, RemoteAgent)
|
||||
assert len(captured) == 1
|
||||
assert captured[0]["_context_tokens"] == 1322
|
||||
|
||||
def _capture(*args: object, **_kwargs: object) -> None:
|
||||
called_with.append(args[1]) # the values dict
|
||||
async def test_partial_ai_message_write_does_not_carry_tokens(self) -> None:
|
||||
"""Only the cancellation write carries `_context_tokens`."""
|
||||
captured: list[dict[str, Any]] = []
|
||||
|
||||
async def _capture(_config: object, values: dict[str, Any]) -> None: # noqa: RUF029
|
||||
captured.append(values)
|
||||
|
||||
tool_widget = MagicMock()
|
||||
tool_widget._tool_name = "read_file"
|
||||
tool_widget._args = {"path": "notes.txt"}
|
||||
|
||||
agent = SimpleNamespace(aupdate_state=AsyncMock(side_effect=_capture))
|
||||
adapter = TextualUIAdapter(
|
||||
mount_message=AsyncMock(),
|
||||
update_status=_noop_status,
|
||||
request_approval=_mock_approval,
|
||||
set_spinner=AsyncMock(),
|
||||
set_active_message=MagicMock(),
|
||||
)
|
||||
adapter._current_tool_messages = {"call-1": tool_widget}
|
||||
|
||||
await _persist_context_tokens(
|
||||
agent, {"configurable": {"thread_id": "t-2"}}, 9999
|
||||
await _handle_interrupt_cleanup(
|
||||
adapter=adapter,
|
||||
agent=agent,
|
||||
config={"configurable": {"thread_id": "t-1"}},
|
||||
pending_text_by_namespace={},
|
||||
captured_input_tokens=7777,
|
||||
captured_output_tokens=0,
|
||||
turn_stats=SessionStats(),
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
assert len(called_with) == 1
|
||||
assert called_with[0] == {"_context_tokens": 9999}
|
||||
assert len(captured) == 2
|
||||
# First write is the interrupted AI message; should not be polluted.
|
||||
assert "_context_tokens" not in captured[0]
|
||||
# Second write is the cancellation HumanMessage; carries the token count.
|
||||
assert captured[1]["_context_tokens"] == 7777
|
||||
|
||||
|
||||
class TestBuildStreamConfig:
|
||||
|
||||
@@ -1,9 +1,16 @@
|
||||
"""Tests for token state persistence and display callbacks."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
from deepagents_code.app import DeepAgentsApp
|
||||
from deepagents_code.token_state import TokenStateMiddleware, TokenTrackingState
|
||||
from deepagents_code.token_state import (
|
||||
TokenStateMiddleware,
|
||||
TokenTrackingState,
|
||||
_extract_context_tokens,
|
||||
)
|
||||
|
||||
|
||||
class TestTokenTrackingState:
|
||||
@@ -17,6 +24,121 @@ class TestTokenTrackingState:
|
||||
assert TokenStateMiddleware.state_schema is TokenTrackingState
|
||||
|
||||
|
||||
class TestExtractContextTokens:
|
||||
"""Tests for `_extract_context_tokens`."""
|
||||
|
||||
def test_prefers_input_plus_output(self) -> None:
|
||||
msg = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 25,
|
||||
"total_tokens": 200, # deliberately inconsistent
|
||||
},
|
||||
)
|
||||
assert _extract_context_tokens(msg) == 125
|
||||
|
||||
def test_falls_back_to_total_tokens(self) -> None:
|
||||
msg = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"total_tokens": 999,
|
||||
},
|
||||
)
|
||||
assert _extract_context_tokens(msg) == 999
|
||||
|
||||
def test_returns_none_without_usage_metadata(self) -> None:
|
||||
msg = AIMessage(content="hi")
|
||||
assert _extract_context_tokens(msg) is None
|
||||
|
||||
def test_returns_none_for_zero_usage(self) -> None:
|
||||
msg = AIMessage(
|
||||
content="hi",
|
||||
usage_metadata={
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
"total_tokens": 0,
|
||||
},
|
||||
)
|
||||
assert _extract_context_tokens(msg) is None
|
||||
|
||||
|
||||
class TestAfterModelHook:
|
||||
"""Tests for the `after_model` persistence hook."""
|
||||
|
||||
async def test_writes_context_tokens_from_last_ai_message(self) -> None:
|
||||
middleware = TokenStateMiddleware()
|
||||
state: dict[str, Any] = {
|
||||
"messages": [
|
||||
HumanMessage(content="hi"),
|
||||
AIMessage(
|
||||
content="response",
|
||||
usage_metadata={
|
||||
"input_tokens": 1500,
|
||||
"output_tokens": 200,
|
||||
"total_tokens": 1700,
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
result = middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
||||
assert result == {"_context_tokens": 1700}
|
||||
|
||||
async def test_returns_none_when_no_ai_message(self) -> None:
|
||||
middleware = TokenStateMiddleware()
|
||||
state: dict[str, Any] = {"messages": [HumanMessage(content="hi")]}
|
||||
result = middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
||||
assert result is None
|
||||
|
||||
async def test_returns_none_when_last_ai_lacks_usage(self) -> None:
|
||||
middleware = TokenStateMiddleware()
|
||||
state: dict[str, Any] = {
|
||||
"messages": [
|
||||
HumanMessage(content="hi"),
|
||||
AIMessage(content="no usage info"),
|
||||
],
|
||||
}
|
||||
result = middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
||||
assert result is None
|
||||
|
||||
async def test_handles_empty_messages(self) -> None:
|
||||
middleware = TokenStateMiddleware()
|
||||
result = middleware.after_model({"messages": []}, runtime=None) # type: ignore[arg-type]
|
||||
assert result is None
|
||||
|
||||
async def test_skips_intervening_tool_messages(self) -> None:
|
||||
"""Picks up the most recent AIMessage even when followed by tool turns."""
|
||||
from langchain_core.messages import ToolMessage
|
||||
|
||||
middleware = TokenStateMiddleware()
|
||||
state: dict[str, Any] = {
|
||||
"messages": [
|
||||
HumanMessage(content="hi"),
|
||||
AIMessage(
|
||||
content="older",
|
||||
usage_metadata={
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 10,
|
||||
"total_tokens": 110,
|
||||
},
|
||||
),
|
||||
ToolMessage(content="tool out", tool_call_id="t1"),
|
||||
AIMessage(
|
||||
content="newer",
|
||||
usage_metadata={
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 50,
|
||||
"total_tokens": 550,
|
||||
},
|
||||
),
|
||||
],
|
||||
}
|
||||
result = middleware.after_model(state, runtime=None) # type: ignore[arg-type]
|
||||
assert result == {"_context_tokens": 550}
|
||||
|
||||
|
||||
class TestTokenDisplayCallbacks:
|
||||
"""Verify the callback-based token tracking that replaced TextualTokenTracker."""
|
||||
|
||||
@@ -93,58 +215,3 @@ class TestTokenDisplayCallbacks:
|
||||
|
||||
assert app._context_tokens == 0
|
||||
assert display_calls == [0]
|
||||
|
||||
|
||||
class TestPersistContextTokens:
|
||||
"""Tests for the `_persist_context_tokens` helper."""
|
||||
|
||||
async def test_calls_aupdate_state_with_token_count(self):
|
||||
"""Happy path: persists the count via `aupdate_state`."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from deepagents_code.textual_adapter import _persist_context_tokens
|
||||
|
||||
agent = AsyncMock()
|
||||
config = {"configurable": {"thread_id": "t-1"}}
|
||||
|
||||
await _persist_context_tokens(agent, config, 4200) # type: ignore[arg-type]
|
||||
|
||||
agent.aupdate_state.assert_awaited_once_with(config, {"_context_tokens": 4200})
|
||||
|
||||
async def test_suppresses_exceptions(self):
|
||||
"""Failures should be swallowed (non-critical persistence)."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from deepagents_code.textual_adapter import _persist_context_tokens
|
||||
|
||||
agent = AsyncMock()
|
||||
agent.aupdate_state.side_effect = RuntimeError("checkpointer down")
|
||||
config = {"configurable": {"thread_id": "t-1"}}
|
||||
|
||||
# Should not raise
|
||||
await _persist_context_tokens(agent, config, 1000) # type: ignore[arg-type]
|
||||
|
||||
async def test_disables_tracing_during_update(self):
|
||||
"""`aupdate_state` should run with LangSmith tracing disabled.
|
||||
|
||||
The token-count write is a private bookkeeping update; surfacing it as a
|
||||
standalone `UpdateState` run in LangSmith clutters traces.
|
||||
"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from langsmith import get_tracing_context
|
||||
|
||||
from deepagents_code.textual_adapter import _persist_context_tokens
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def _capture(*_args: object, **_kwargs: object) -> None: # noqa: RUF029 # AsyncMock side_effect must be a coroutine function
|
||||
captured["enabled"] = get_tracing_context().get("enabled")
|
||||
|
||||
agent = AsyncMock()
|
||||
agent.aupdate_state.side_effect = _capture
|
||||
config = {"configurable": {"thread_id": "t-1"}}
|
||||
|
||||
await _persist_context_tokens(agent, config, 4200) # type: ignore[arg-type]
|
||||
|
||||
assert captured["enabled"] is False
|
||||
|
||||
@@ -2,13 +2,12 @@
|
||||
|
||||
## [0.6.3](https://github.com/langchain-ai/deepagents/compare/deepagents==0.6.2...deepagents==0.6.3) (2026-05-20)
|
||||
|
||||
|
||||
### Bug Fixes
|
||||
|
||||
* **sdk:** anchor ripgrep glob to search root ([#3454](https://github.com/langchain-ai/deepagents/issues/3454)) ([e50fa3f](https://github.com/langchain-ai/deepagents/commit/e50fa3f00ab1b1a84bbaed74bf7e89118b7c2d82))
|
||||
* **sdk:** assign UUIDs to ID-less messages in _messages_delta_reducer ([#3513](https://github.com/langchain-ai/deepagents/issues/3513)) ([6d959ad](https://github.com/langchain-ai/deepagents/commit/6d959ade30655eae3967c9809994434e0bbd1148))
|
||||
* **sdk:** clarify skill source labels in system prompt ([#3464](https://github.com/langchain-ai/deepagents/issues/3464)) ([fc6a24f](https://github.com/langchain-ai/deepagents/commit/fc6a24f18829cf3f36089945226edfa50d52ab9e))
|
||||
* **sdk:** strip HTML comments from memory content before system prompt injection ([#3462](https://github.com/langchain-ai/deepagents/issues/3462)) ([bfbb8bc](https://github.com/langchain-ai/deepagents/commit/bfbb8bc5575ebd1ba9aa29430f6d2f86c24b7d3c))
|
||||
* Anchor ripgrep glob to search root ([#3454](https://github.com/langchain-ai/deepagents/issues/3454)) ([e50fa3f](https://github.com/langchain-ai/deepagents/commit/e50fa3f00ab1b1a84bbaed74bf7e89118b7c2d82))
|
||||
* Assign UUIDs to ID-less messages in _messages_delta_reducer ([#3513](https://github.com/langchain-ai/deepagents/issues/3513)) ([6d959ad](https://github.com/langchain-ai/deepagents/commit/6d959ade30655eae3967c9809994434e0bbd1148))
|
||||
* Clarify skill source labels in system prompt ([#3464](https://github.com/langchain-ai/deepagents/issues/3464)) ([fc6a24f](https://github.com/langchain-ai/deepagents/commit/fc6a24f18829cf3f36089945226edfa50d52ab9e))
|
||||
* Strip HTML comments from memory content before system prompt injection ([#3462](https://github.com/langchain-ai/deepagents/issues/3462)) ([bfbb8bc](https://github.com/langchain-ai/deepagents/commit/bfbb8bc5575ebd1ba9aa29430f6d2f86c24b7d3c))
|
||||
|
||||
## [0.6.2](https://github.com/langchain-ai/deepagents/compare/deepagents==0.6.1...deepagents==0.6.2) (2026-05-18)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user