mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 10:56:02 +00:00
* feat: show real-time context window usage in chat UI (#3125) Adds a `context_usage` block to `GET /api/threads/{id}/token-usage` (token count from the live checkpoint, the thread model's `context_window`, and a percentage), introduces a new `ModelConfig.context_window` distinct from the per-call `max_tokens` output cap, and surfaces the percentage in the chat header — inside `TokenUsageIndicator` when token-usage tracking is on, or as a standalone badge when it's off so context capacity stays visible independent of cost tracking. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * feat: per-category breakdown for context window usage Replace the single-number context_usage payload with a Claude-Code-style breakdown — messages, system prompt, skills, system/MCP tools (active + deferred), custom agents, memory injection, autocompact buffer, and free space — and surface it in the chat UI with a segmented progress bar and per-row table. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * docs(config): document context_window across model examples Add `context_window` to every example model in config.example.yaml so the new chat-UI "% context used" indicator works out of the box for whichever example a user adopts. Each value is the published default at the time of writing; users are pointed at the official model spec to verify. Bumps config_version to 11 so `make config-upgrade` flags outdated user configs. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style: ruff format (line-length 240) No behavior change — collapses two multi-line expressions that fit on one line under the project's 240-char limit. Picked up by `make format`. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * review: address Copilot bot comments on #3183 - token-usage-indicator: switch `{contextPercentage && (...)}` to an explicit `!= null` check. (The string `"0"` is actually truthy in JS so the original code wasn't buggy, but the explicit check is clearer.) - context-usage-breakdown: drop the `useMemo` around segments/totals — the computation is O(n) over a handful of rows and the previous memo deps omitted `t.contextUsage.categories`, so the bar's tooltips/aria-labels could stay in the old language after a locale switch. - context_usage._split_tools: snapshot MCP names from `get_cached_mcp_tools()` directly instead of re-reading `extensions_config.json` after `get_available_tools()` already loaded it. Removes redundant file I/O on every `/token-usage` poll. (`get_available_tools()` still emits its own INFO logs — silencing those is out of scope here.) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style(frontend): prettier --write context-usage-breakdown CI's `pnpm format` (prettier --check) caught two lines previously formatted by hand. Collapses one comma to fit on one line; no behavior change. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(gateway): correct context-usage breakdown + add exact token counting The context-usage indicator shipped two bugs that silently zeroed whole breakdown rows (both caught by try/except, so the feature looked alive but produced wrong numbers): 1. _count_system_prompt passed app_config= to get_deferred_tools_prompt_section, which only accepts deferred_names -> TypeError swallowed -> system_prompt row always 0, and used_tokens/percentage undercounted by the full prompt. Also subtracted the deferred section twice (the rendered prompt already excluded it). Fix: derive deferred names deterministically and pass them to apply_prompt_template; drop the redundant subtraction. 2. _split_tools imported a non-existent get_deferred_registry -> ImportError swallowed -> all four tool-category rows always 0. Fix: classify via the public is_mcp_tool predicate + tool_search.enabled (mirrors build_deferred_tool_setup); the MCP tag is set by get_available_tools. Added token_usage.counting (approximate|exact). 'exact' routes text/schema/ message counting through the model tokenizer (tiktoken cl100k_base) via the existing memory-module machinery (lazy load + cache + cooldown + CJK-aware fallback), so CJK-heavy threads stop being undercounted by chars//4. Regression + e2e tests added; 6621 backend tests pass. * fix(gateway): harden context usage accounting * fix(gateway): count promoted MCP tools as active in context usage Promoted tools (deferred MCP tools the thread has fetched via tool_search) have their full schema bound on every subsequent turn by DeferredToolFilterMiddleware, so they consume context like any active tool. The breakdown previously left them in the reserved *_deferred rows, under- counting the thread's used_tokens. Classification now treats a tool as deferred only when tool_search is enabled, it is MCP-sourced, AND it has not been promoted. The promoted set is read from the checkpoint's channel_values and scoped by catalog hash — matching the runtime middleware, so a stale promotion from MCP-config drift cannot inflate the active count. The static system prompt still lists all deferred tool names (promotions only affect schema binding, not the prompt), so _count_system_prompt's deferred rendering is intentionally left unchanged. 8 new tests cover classification, catalog-hash scoping (match / drift / compute-failure / malformed), and checkpoint extraction. * fix(context): address review feedback * fix(context): count structured message payloads * fix(context): harden usage accounting * fix(config): bump schema for context usage fields * refactor: narrow context usage to core indicator --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
17461ee52e
commit
85c3909c2e
@ -891,6 +891,8 @@ The Web UI shows the active goal above the composer. The same command is availab
|
||||
|
||||
Use `/compact` in the Web UI composer to summarize older context for the current thread. DeerFlow keeps the full chat visible, but future model calls use the compacted summary plus recent messages. The command is ignored when there is not enough history to compact, and it is blocked while the thread has a run in flight, including when that run is owned by another Gateway worker. If a multi-worker reservation loses its lease, DeerFlow cancels the checkpoint writer before the replacing run proceeds and returns a retryable conflict after cleanup. Thread-title edits are serialized through the same state-write boundary and show a conflict without closing the rename dialog when a run is active.
|
||||
|
||||
The chat header also shows a context-window gauge when the selected model has a positive `context_window` configured. It estimates the latest materialized checkpoint's message tokens and keeps the previous same-thread percentage visible while data refetches, independently of the cumulative token-usage setting.
|
||||
|
||||
### Sub-Agents
|
||||
|
||||
Sub-agents are an optimization, not the default response to a complex request.
|
||||
|
||||
@ -495,7 +495,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores `
|
||||
| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - serve artifacts; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types |
|
||||
| **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`<think>...</think>`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing |
|
||||
| **Input Polish** (`/api/input-polish`) | `POST /` - rewrite a composer draft before it is sent. This is a short authenticated `runs:create` LLM request using `input_polish` config; it does not create a LangGraph run, persist a message, or modify thread state. Shares the non-graph one-shot LLM path (`deerflow.utils.oneshot_llm.run_oneshot_llm`) with the suggestions route so model build + Langfuse metadata + invoke stay in one place; validates the same stripped view of the draft it sends to the model, and preserves literal `<think>` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) |
|
||||
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block. Before the first journaled run, an empty run-event message feed is seeded from an existing checkpoint head so legacy checkpoint-only history receives earlier thread-global sequence numbers and remains visible after the new run; a thread with no checkpoint or an already-populated feed skips this compatibility path. `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest completed or interrupted assistant answer, carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `POST /edit-regenerate/prepare` - prepare a checkpoint replay from the latest editable human turn with a replacement user message and edit replay metadata; it carries the current thread title the same way, but only when the replay base already has one — an untitled base belongs to a thread the title middleware has not named yet, so pinning the current title there would keep a name generated from the prompt the edit just replaced; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate/edit-replay filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens |
|
||||
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block. Before the first journaled run, an empty run-event message feed is seeded from an existing checkpoint head so legacy checkpoint-only history receives earlier thread-global sequence numbers and remains visible after the new run; a thread with no checkpoint or an already-populated feed skips this compatibility path. `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest completed or interrupted assistant answer, carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `POST /edit-regenerate/prepare` - prepare a checkpoint replay from the latest editable human turn with a replacement user message and edit replay metadata; it carries the current thread title the same way, but only when the replay base already has one — an untitled base belongs to a thread the title middleware has not named yet, so pinning the current title there would keep a name generated from the prompt the edit just replaced; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate/edit-replay filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens plus an optional `context_usage` percentage. Context usage approximately counts messages from the latest materialized thread state through `build_thread_checkpoint_state_accessor`, so full and delta checkpoint modes expose the same input. The percentage uses the latest run's model and its configured `context_window`. |
|
||||
| **Feedback** (`/api/threads/{id}/runs/{rid}/feedback`) | `PUT /` - upsert feedback; `DELETE /` - delete user feedback; `POST /` - create feedback; `GET /` - list feedback; `GET /stats` - aggregate stats; `DELETE /{fid}` - delete specific |
|
||||
| **Runs** (`/api/runs`) | `POST /stream` - stateless run + SSE; `POST /wait` - stateless run + block; `GET /{rid}/messages` - paginated messages by run_id `{data, has_more}` (cursor: `after_seq`/`before_seq`); `GET /{rid}/feedback` - list feedback by run_id |
|
||||
| **GitHub Webhooks** (`/api/webhooks/github`) | `POST /` - receive GitHub App / repo webhook deliveries. Verifies `X-Hub-Signature-256` against `GITHUB_WEBHOOK_SECRET`; exempt from auth + CSRF because authenticity is enforced by HMAC. The route is fail-closed: mounted only when `GITHUB_WEBHOOK_SECRET` is set, or when explicit dev opt-in `DEER_FLOW_ALLOW_UNVERIFIED_GITHUB_WEBHOOKS=1` is set. Recognized events include `ping`, `issues`, `issue_comment`, `pull_request`, `pull_request_review`, and `pull_request_review_comment`; unknown events return 200 with `handled=false`. Fan-out runtime failures return 503, keeping the delivery recorded as failed for manual/API/scripted redelivery (GitHub does not automatically retry any failed delivery, 5xx included); permanent/non-retryable conditions such as `channels.github.enabled: false`, unknown events, malformed payloads, or unavailable channel service return 200 with a skipped/handled response. |
|
||||
|
||||
86
backend/app/gateway/context_usage.py
Normal file
86
backend/app/gateway/context_usage.py
Normal file
@ -0,0 +1,86 @@
|
||||
"""Compute the current message-context usage for a thread."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
from app.gateway.deps import get_config
|
||||
from app.gateway.services import build_thread_checkpoint_state_accessor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _count_messages_approximately(messages: list[Any]) -> int:
|
||||
"""Count checkpoint messages with LangChain's network-free heuristic."""
|
||||
if not messages:
|
||||
return 0
|
||||
from langchain_core.messages.utils import count_tokens_approximately
|
||||
|
||||
return int(count_tokens_approximately(messages))
|
||||
|
||||
|
||||
async def _load_checkpoint_messages(accessor: Any, config: dict[str, Any]) -> list[Any]:
|
||||
"""Read materialized messages so full and delta checkpoints behave alike."""
|
||||
snapshot = await accessor.aget(config)
|
||||
values = getattr(snapshot, "values", None) or {}
|
||||
if not isinstance(values, dict):
|
||||
return []
|
||||
return list(values.get("messages") or [])
|
||||
|
||||
|
||||
async def _resolve_thread_model_name(run_store: Any, thread_id: str, app_config: Any) -> str | None:
|
||||
"""Prefer the latest run's model, then fall back to the first configured model."""
|
||||
try:
|
||||
runs = await run_store.list_by_thread(thread_id, limit=1)
|
||||
except Exception:
|
||||
runs = []
|
||||
if runs:
|
||||
latest = runs[0]
|
||||
name = latest.get("model_name") if isinstance(latest, dict) else getattr(latest, "model_name", None)
|
||||
if isinstance(name, str) and name:
|
||||
return name
|
||||
models = getattr(app_config, "models", None) or []
|
||||
return models[0].name if models else None
|
||||
|
||||
|
||||
def build_context_usage_payload(*, token_count: int, max_context_tokens: int | None) -> dict[str, Any]:
|
||||
"""Build the stable API payload for a message count and model capacity."""
|
||||
percentage: float | None = None
|
||||
if max_context_tokens and max_context_tokens > 0:
|
||||
percentage = round(token_count / max_context_tokens * 100, 1)
|
||||
return {
|
||||
"token_count": token_count,
|
||||
"max_context_tokens": max_context_tokens,
|
||||
"percentage": percentage,
|
||||
}
|
||||
|
||||
|
||||
async def build_context_usage(request: Request, thread_id: str, run_store: Any) -> dict[str, Any] | None:
|
||||
"""Return approximate usage for the latest materialized thread checkpoint."""
|
||||
try:
|
||||
app_config = get_config()
|
||||
except HTTPException:
|
||||
return None
|
||||
|
||||
try:
|
||||
accessor, checkpoint_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id)
|
||||
messages = await _load_checkpoint_messages(accessor, checkpoint_config)
|
||||
except Exception:
|
||||
logger.warning("Failed to load checkpoint for context usage on thread %s", thread_id, exc_info=True)
|
||||
return None
|
||||
|
||||
try:
|
||||
token_count = await asyncio.to_thread(_count_messages_approximately, messages)
|
||||
except Exception:
|
||||
logger.warning("Failed to count context messages for thread %s", thread_id, exc_info=True)
|
||||
return None
|
||||
|
||||
model_name = await _resolve_thread_model_name(run_store, thread_id, app_config)
|
||||
model_config = app_config.get_model_config(model_name) if model_name else None
|
||||
configured_window = getattr(model_config, "context_window", None) if model_config is not None else None
|
||||
max_context_tokens = int(configured_window) if configured_window else None
|
||||
return build_context_usage_payload(token_count=token_count, max_context_tokens=max_context_tokens)
|
||||
@ -33,6 +33,7 @@ from app.gateway.checkpoint_lineage import (
|
||||
find_checkpoint_before_message_chronologically,
|
||||
is_duration_only_checkpoint,
|
||||
)
|
||||
from app.gateway.context_usage import build_context_usage
|
||||
from app.gateway.deps import get_current_user, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge
|
||||
from app.gateway.pagination import trim_run_message_page
|
||||
from app.gateway.run_models import RunCreateRequest
|
||||
@ -180,6 +181,12 @@ class ThreadTokenUsageCallerBreakdown(BaseModel):
|
||||
middleware: int = 0
|
||||
|
||||
|
||||
class ThreadContextUsage(BaseModel):
|
||||
token_count: int = 0
|
||||
max_context_tokens: int | None = None
|
||||
percentage: float | None = None
|
||||
|
||||
|
||||
class ThreadTokenUsageResponse(BaseModel):
|
||||
thread_id: str
|
||||
total_tokens: int = 0
|
||||
@ -188,6 +195,7 @@ class ThreadTokenUsageResponse(BaseModel):
|
||||
total_runs: int = 0
|
||||
by_model: dict[str, ThreadTokenUsageModelBreakdown] = Field(default_factory=dict)
|
||||
by_caller: ThreadTokenUsageCallerBreakdown = Field(default_factory=ThreadTokenUsageCallerBreakdown)
|
||||
context_usage: ThreadContextUsage | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -1459,4 +1467,5 @@ async def thread_token_usage(
|
||||
agg = await run_store.aggregate_tokens_by_thread(thread_id, include_active=True)
|
||||
else:
|
||||
agg = await run_store.aggregate_tokens_by_thread(thread_id)
|
||||
return ThreadTokenUsageResponse(thread_id=thread_id, **agg)
|
||||
context_usage = await build_context_usage(request, thread_id, run_store)
|
||||
return ThreadTokenUsageResponse(thread_id=thread_id, context_usage=context_usage, **agg)
|
||||
|
||||
@ -32,6 +32,16 @@ class ModelConfig(BaseModel):
|
||||
description="Extra settings to be passed to the model when thinking is disabled",
|
||||
)
|
||||
supports_vision: bool = Field(default_factory=lambda: False, description="Whether the model supports vision/image inputs")
|
||||
context_window: int | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
description=(
|
||||
"Positive total context window size in tokens (prompt + completion). Used to compute the real-time "
|
||||
"context usage percentage displayed in the chat UI. Distinct from `max_tokens`, which is the "
|
||||
"per-call output cap passed to the provider. Leave unset if unknown; the UI will hide the "
|
||||
"percentage."
|
||||
),
|
||||
)
|
||||
stream_chunk_timeout: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
|
||||
@ -219,6 +219,9 @@ def create_chat_model(name: str | None = None, thinking_enabled: bool = False, *
|
||||
"when_thinking_disabled",
|
||||
"thinking",
|
||||
"supports_vision",
|
||||
# Runtime/UI metadata used to size the context indicator. Provider
|
||||
# clients do not accept this as a model-constructor argument.
|
||||
"context_window",
|
||||
# Presentation-only metadata (consumed by the console's cost
|
||||
# display) — must never reach the provider client, which would
|
||||
# forward unknown kwargs into the completion request payload.
|
||||
|
||||
@ -1,7 +1,8 @@
|
||||
"""Dual-mode (full/delta) parity for the gateway thread-state endpoints.
|
||||
|
||||
Drives ``GET /api/threads/{id}``, ``GET /api/threads/{id}/state`` and
|
||||
``POST /api/threads/{id}/history`` through the real route stack
|
||||
Drives ``GET /api/threads/{id}``, ``GET /api/threads/{id}/state``,
|
||||
``POST /api/threads/{id}/history``, and the context-usage checkpoint reader
|
||||
through the real materialization stack
|
||||
(``build_thread_checkpoint_state_accessor`` -> factory-built graph ->
|
||||
``CheckpointStateAccessor``) against a real ``InMemorySaver``, once per
|
||||
checkpoint channel mode, and asserts the wire responses are identical apart
|
||||
@ -24,6 +25,7 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
|
||||
from app.gateway import context_usage
|
||||
from app.gateway import services as gateway_services
|
||||
from app.gateway.routers import threads
|
||||
from deerflow.agents.thread_state import get_thread_state_schema
|
||||
@ -113,6 +115,46 @@ def test_thread_state_endpoints_are_mode_invariant(_stub_app_config, monkeypatch
|
||||
assert any(full["history_messages"]), "expected history snapshots with messages"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["full", "delta"])
|
||||
def test_context_usage_reads_materialized_messages_in_both_modes(
|
||||
mode: str,
|
||||
_stub_app_config,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Context usage must not read raw delta ``channel_values``."""
|
||||
app = make_authed_test_app()
|
||||
store = InMemoryStore()
|
||||
checkpointer = InMemorySaver()
|
||||
app.state.store = store
|
||||
app.state.checkpointer = checkpointer
|
||||
app.state.thread_store = MemoryThreadMetaStore(store)
|
||||
app.state.checkpoint_channel_mode = mode
|
||||
app.state.run_event_store = SimpleNamespace()
|
||||
|
||||
graph = _build_reply_graph(mode, checkpointer)
|
||||
monkeypatch.setattr(
|
||||
gateway_services,
|
||||
"resolve_agent_factory",
|
||||
lambda assistant_id=None: lambda config: graph,
|
||||
)
|
||||
|
||||
config: dict[str, Any] = {"configurable": {"thread_id": _THREAD_ID}}
|
||||
inject_checkpoint_mode(config, mode)
|
||||
for i in range(2):
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content=f"question-{i}", id=f"h{i}")]}, config))
|
||||
|
||||
request = SimpleNamespace(app=app)
|
||||
accessor, read_config = asyncio.run(gateway_services.build_thread_checkpoint_state_accessor(request, thread_id=_THREAD_ID))
|
||||
messages = asyncio.run(context_usage._load_checkpoint_messages(accessor, read_config))
|
||||
|
||||
assert [(message.type, message.content, message.id) for message in messages] == [
|
||||
("human", "question-0", "h0"),
|
||||
("ai", "answer-1", "a1"),
|
||||
("human", "question-1", "h1"),
|
||||
("ai", "answer-3", "a3"),
|
||||
]
|
||||
|
||||
|
||||
def test_full_mode_gateway_rejects_delta_thread_with_409(_stub_app_config, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Fail-closed gate at the HTTP boundary, against a real checkpointer.
|
||||
|
||||
|
||||
@ -1,3 +1,6 @@
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from deerflow.config.model_config import ModelConfig
|
||||
|
||||
|
||||
@ -28,3 +31,13 @@ def test_responses_api_fields_round_trip_in_model_dump():
|
||||
|
||||
assert dumped["use_responses_api"] is True
|
||||
assert dumped["output_version"] == "responses/v1"
|
||||
|
||||
|
||||
def test_context_window_round_trips_when_positive():
|
||||
assert _make_model(context_window=128_000).context_window == 128_000
|
||||
|
||||
|
||||
@pytest.mark.parametrize("context_window", [0, -1])
|
||||
def test_context_window_rejects_non_positive_capacity(context_window):
|
||||
with pytest.raises(ValidationError, match="context_window"):
|
||||
_make_model(context_window=context_window)
|
||||
|
||||
@ -141,6 +141,19 @@ def test_pricing_metadata_never_reaches_the_provider_client(monkeypatch):
|
||||
assert "pricing" not in FakeChatModel.captured_kwargs
|
||||
|
||||
|
||||
def test_context_window_never_reaches_the_provider_client(monkeypatch):
|
||||
"""Context sizing metadata belongs to DeerFlow, not the provider SDK."""
|
||||
model = _make_model("large-context")
|
||||
model.context_window = 200_000
|
||||
cfg = _make_app_config([model])
|
||||
_patch_factory(monkeypatch, cfg)
|
||||
|
||||
FakeChatModel.captured_kwargs = {}
|
||||
factory_module.create_chat_model(name="large-context")
|
||||
|
||||
assert "context_window" not in FakeChatModel.captured_kwargs
|
||||
|
||||
|
||||
def test_appends_all_tracing_callbacks(monkeypatch):
|
||||
cfg = _make_app_config([_make_model("alpha")])
|
||||
_patch_factory(monkeypatch, cfg)
|
||||
|
||||
@ -1,46 +1,20 @@
|
||||
"""Tests for thread-level token usage aggregation API."""
|
||||
"""Tests for thread-level token usage and context-window usage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from _router_auth_helpers import make_authed_test_app
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from app.gateway import context_usage
|
||||
from app.gateway.routers import thread_runs
|
||||
|
||||
|
||||
def _make_app(run_store: MagicMock):
|
||||
app = make_authed_test_app()
|
||||
app.include_router(thread_runs.router)
|
||||
app.state.run_store = run_store
|
||||
return app
|
||||
|
||||
|
||||
def test_thread_token_usage_returns_stable_shape():
|
||||
run_store = MagicMock()
|
||||
run_store.aggregate_tokens_by_thread = AsyncMock(
|
||||
return_value={
|
||||
"total_tokens": 150,
|
||||
"total_input_tokens": 90,
|
||||
"total_output_tokens": 60,
|
||||
"total_runs": 2,
|
||||
"by_model": {"unknown": {"tokens": 150, "runs": 2}},
|
||||
"by_caller": {
|
||||
"lead_agent": 120,
|
||||
"subagent": 25,
|
||||
"middleware": 5,
|
||||
},
|
||||
},
|
||||
)
|
||||
app = _make_app(run_store)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/api/threads/thread-1/token-usage")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"thread_id": "thread-1",
|
||||
def _aggregate_result() -> dict:
|
||||
return {
|
||||
"total_tokens": 150,
|
||||
"total_input_tokens": 90,
|
||||
"total_output_tokens": 60,
|
||||
@ -52,31 +26,157 @@ def test_thread_token_usage_returns_stable_shape():
|
||||
"middleware": 5,
|
||||
},
|
||||
}
|
||||
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1")
|
||||
|
||||
|
||||
def test_thread_token_usage_can_include_active_runs():
|
||||
def _make_run_store(*, model_name: str | None = None) -> MagicMock:
|
||||
run_store = MagicMock()
|
||||
run_store.aggregate_tokens_by_thread = AsyncMock(
|
||||
return_value={
|
||||
"total_tokens": 175,
|
||||
"total_input_tokens": 120,
|
||||
"total_output_tokens": 55,
|
||||
"total_runs": 3,
|
||||
"by_model": {"unknown": {"tokens": 175, "runs": 3}},
|
||||
"by_caller": {
|
||||
"lead_agent": 145,
|
||||
"subagent": 25,
|
||||
"middleware": 5,
|
||||
},
|
||||
},
|
||||
)
|
||||
run_store.aggregate_tokens_by_thread = AsyncMock(return_value=_aggregate_result())
|
||||
runs = [{"model_name": model_name}] if model_name else []
|
||||
run_store.list_by_thread = AsyncMock(return_value=runs)
|
||||
return run_store
|
||||
|
||||
|
||||
def _make_app(run_store: MagicMock):
|
||||
app = make_authed_test_app()
|
||||
app.include_router(thread_runs.router)
|
||||
app.state.run_store = run_store
|
||||
return app
|
||||
|
||||
|
||||
def test_thread_token_usage_returns_stable_shape(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
run_store = _make_run_store()
|
||||
build_context_usage = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(thread_runs, "build_context_usage", build_context_usage)
|
||||
app = _make_app(run_store)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/api/threads/thread-1/token-usage")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {
|
||||
"thread_id": "thread-1",
|
||||
**_aggregate_result(),
|
||||
"context_usage": None,
|
||||
}
|
||||
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1")
|
||||
build_context_usage.assert_awaited_once()
|
||||
|
||||
|
||||
def test_thread_token_usage_can_include_active_runs(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
run_store = _make_run_store()
|
||||
build_context_usage = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(thread_runs, "build_context_usage", build_context_usage)
|
||||
app = _make_app(run_store)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/api/threads/thread-1/token-usage?include_active=true")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["total_tokens"] == 175
|
||||
assert response.json()["total_runs"] == 3
|
||||
run_store.aggregate_tokens_by_thread.assert_awaited_once_with("thread-1", include_active=True)
|
||||
|
||||
|
||||
def test_thread_token_usage_serializes_context_percentage(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
run_store = _make_run_store()
|
||||
monkeypatch.setattr(
|
||||
thread_runs,
|
||||
"build_context_usage",
|
||||
AsyncMock(
|
||||
return_value={
|
||||
"token_count": 350,
|
||||
"max_context_tokens": 1000,
|
||||
"percentage": 35.0,
|
||||
}
|
||||
),
|
||||
)
|
||||
app = _make_app(run_store)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/api/threads/thread-1/token-usage")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["context_usage"] == {
|
||||
"token_count": 350,
|
||||
"max_context_tokens": 1000,
|
||||
"percentage": 35.0,
|
||||
}
|
||||
|
||||
|
||||
def test_build_context_usage_payload_computes_percentage() -> None:
|
||||
assert context_usage.build_context_usage_payload(token_count=350, max_context_tokens=1000) == {
|
||||
"token_count": 350,
|
||||
"max_context_tokens": 1000,
|
||||
"percentage": 35.0,
|
||||
}
|
||||
|
||||
|
||||
def test_build_context_usage_payload_handles_unknown_capacity() -> None:
|
||||
assert context_usage.build_context_usage_payload(token_count=350, max_context_tokens=None) == {
|
||||
"token_count": 350,
|
||||
"max_context_tokens": None,
|
||||
"percentage": None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_thread_model_prefers_latest_run() -> None:
|
||||
run_store = _make_run_store(model_name="thread-model")
|
||||
app_config = SimpleNamespace(models=[SimpleNamespace(name="fallback-model")])
|
||||
|
||||
assert await context_usage._resolve_thread_model_name(run_store, "thread-1", app_config) == "thread-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_thread_model_falls_back_to_first_configured_model() -> None:
|
||||
run_store = _make_run_store()
|
||||
app_config = SimpleNamespace(models=[SimpleNamespace(name="fallback-model")])
|
||||
|
||||
assert await context_usage._resolve_thread_model_name(run_store, "thread-1", app_config) == "fallback-model"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_context_usage_counts_materialized_messages(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
messages = [SimpleNamespace(content="hello")]
|
||||
snapshot = SimpleNamespace(values={"messages": messages})
|
||||
accessor = SimpleNamespace(aget=AsyncMock(return_value=snapshot))
|
||||
monkeypatch.setattr(
|
||||
context_usage,
|
||||
"build_thread_checkpoint_state_accessor",
|
||||
AsyncMock(return_value=(accessor, {"configurable": {"thread_id": "thread-1"}})),
|
||||
)
|
||||
model_config = SimpleNamespace(context_window=1000)
|
||||
app_config = SimpleNamespace(
|
||||
models=[SimpleNamespace(name="fallback-model")],
|
||||
get_model_config=lambda name: model_config if name == "thread-model" else None,
|
||||
)
|
||||
monkeypatch.setattr(context_usage, "get_config", lambda: app_config)
|
||||
monkeypatch.setattr(context_usage, "_count_messages_approximately", lambda value: 250 if value == messages else 0)
|
||||
|
||||
result = await context_usage.build_context_usage(
|
||||
request=SimpleNamespace(app=SimpleNamespace()),
|
||||
thread_id="thread-1",
|
||||
run_store=_make_run_store(model_name="thread-model"),
|
||||
)
|
||||
|
||||
assert result == {
|
||||
"token_count": 250,
|
||||
"max_context_tokens": 1000,
|
||||
"percentage": 25.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_context_usage_returns_none_when_checkpoint_read_fails(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
context_usage,
|
||||
"build_thread_checkpoint_state_accessor",
|
||||
AsyncMock(side_effect=RuntimeError("checkpoint unavailable")),
|
||||
)
|
||||
monkeypatch.setattr(context_usage, "get_config", lambda: SimpleNamespace())
|
||||
|
||||
result = await context_usage.build_context_usage(
|
||||
request=SimpleNamespace(app=SimpleNamespace()),
|
||||
thread_id="thread-1",
|
||||
run_store=_make_run_store(),
|
||||
)
|
||||
|
||||
assert result is None
|
||||
|
||||
@ -94,7 +94,15 @@ max_recursion_limit: 1000
|
||||
# ============================================================================
|
||||
# Models Configuration
|
||||
# ============================================================================
|
||||
# Configure available LLM models for the agent to use
|
||||
# Configure available LLM models for the agent to use.
|
||||
#
|
||||
# Two token fields look similar but mean different things:
|
||||
# - `max_tokens` is the per-call OUTPUT cap passed to the provider.
|
||||
# - `context_window` is a positive integer for the total context capacity
|
||||
# (prompt + completion) and drives the real-time "% context used" indicator
|
||||
# in the chat UI.
|
||||
# Leave `context_window` unset if the provider limit is unknown; the percentage
|
||||
# will not render. Verify configured values against the provider's model docs.
|
||||
#
|
||||
# Optional per-model pricing (powers the real-cost display on the workspace
|
||||
# console). Add a `pricing` block to any model entry; use ONE currency across
|
||||
@ -118,6 +126,7 @@ models:
|
||||
# api_key: $VOLCENGINE_API_KEY
|
||||
# timeout: 600.0
|
||||
# max_retries: 2
|
||||
# context_window: 262144 # Total prompt + completion capacity
|
||||
# supports_thinking: true
|
||||
# supports_vision: true
|
||||
# supports_reasoning_effort: true
|
||||
@ -182,7 +191,8 @@ models:
|
||||
# api_key: $OPENAI_API_KEY # Use environment variable
|
||||
# request_timeout: 600.0
|
||||
# max_retries: 2
|
||||
# max_tokens: 4096
|
||||
# max_tokens: 4096 # Per-call output cap
|
||||
# context_window: 128000 # Total prompt + completion capacity
|
||||
# temperature: 0.7
|
||||
# supports_vision: true # Enable vision support for view_image tool
|
||||
|
||||
@ -196,6 +206,7 @@ models:
|
||||
# max_retries: 2
|
||||
# use_responses_api: true
|
||||
# output_version: responses/v1
|
||||
# context_window: 400000
|
||||
# supports_vision: true
|
||||
|
||||
# Example: Ollama (native provider — preserves thinking/reasoning content)
|
||||
@ -216,6 +227,7 @@ models:
|
||||
# num_predict: 8192
|
||||
# temperature: 0.7
|
||||
# reasoning: true # Passes think:true to Ollama native API
|
||||
# context_window: 32768 # Match the context length served by Ollama
|
||||
# supports_thinking: true
|
||||
# supports_vision: false
|
||||
#
|
||||
@ -227,6 +239,7 @@ models:
|
||||
# num_predict: 8192
|
||||
# temperature: 0.7
|
||||
# reasoning: true
|
||||
# context_window: 32768
|
||||
# supports_thinking: true
|
||||
# supports_vision: true
|
||||
#
|
||||
@ -245,7 +258,8 @@ models:
|
||||
# api_key: $ANTHROPIC_API_KEY
|
||||
# default_request_timeout: 600.0
|
||||
# max_retries: 2
|
||||
# max_tokens: 16000
|
||||
# max_tokens: 16000 # Per-call output cap
|
||||
# context_window: 200000 # Total prompt + completion capacity
|
||||
# supports_vision: true
|
||||
# supports_thinking: true
|
||||
# when_thinking_enabled:
|
||||
@ -265,6 +279,7 @@ models:
|
||||
# timeout: 600.0
|
||||
# max_retries: 2
|
||||
# max_tokens: 8192
|
||||
# context_window: 1048576
|
||||
# supports_vision: true
|
||||
|
||||
# Example: Gemini model via OpenAI-compatible gateway (with thinking support)
|
||||
@ -281,6 +296,7 @@ models:
|
||||
# request_timeout: 600.0
|
||||
# max_retries: 2
|
||||
# max_tokens: 16384
|
||||
# context_window: 1048576
|
||||
# supports_thinking: true
|
||||
# supports_vision: true
|
||||
# when_thinking_enabled:
|
||||
@ -739,7 +755,7 @@ tools:
|
||||
# use: deerflow.community.groundroute.tools:web_search_tool
|
||||
# max_results: 5 # Clamped to 1-50 by GroundRoute
|
||||
# # api_key: $GROUNDROUTE_API_KEY # Optional if the env var is set
|
||||
|
||||
|
||||
# Web search tool (uses fastCRW - Firecrawl-compatible web scraper, single binary,
|
||||
# self-host or cloud. Cloud requires CRW_API_KEY; self-host may need no key.)
|
||||
# - name: web_search
|
||||
@ -913,7 +929,7 @@ tools:
|
||||
# group: web
|
||||
# use: deerflow.community.groundroute.tools:web_fetch_tool
|
||||
# # api_key: $GROUNDROUTE_API_KEY
|
||||
|
||||
|
||||
# Web fetch tool (uses fastCRW - Firecrawl-compatible web scraper, single binary,
|
||||
# self-host or cloud. Cloud requires CRW_API_KEY; self-host may need no key.)
|
||||
# NOTE: Only one web_fetch provider can be active at a time.
|
||||
|
||||
@ -79,6 +79,8 @@ The frontend is a stateful chat application. Users create **threads** (conversat
|
||||
`["mcpConfig"]` only after success.
|
||||
6. Components subscribe to thread state and render updates
|
||||
|
||||
The chat header's context-window control is intentionally persistent: while `context_usage` is unavailable, `ContextUsageBadge` renders a gauge placeholder rather than unmounting; once data arrives, the same position shows the percentage. `useThreadTokenUsage` retains placeholder data only when the response `thread_id` still matches the active route, so same-thread refetches do not flicker and cross-thread navigation never displays the previous chat's usage.
|
||||
|
||||
Run duration is run-scoped UI metadata even though the compatibility field `additional_kwargs.turn_duration` is repeated on historical AI messages. `core/messages/run-duration.ts` folds those copies into one display anchored after the run's last visible message group. `MessageList` owns the temporary client-side duration for a just-completed live turn until authoritative history arrives. The duration is total run wall-clock time, not per-message reasoning time; reasoning disclosure and run activity/duration are rendered separately.
|
||||
|
||||
The workspace-change card follows the same rule: it is resolved from `(threadId, runId)` alone, so every AI message of a run would render an identical copy. A run ends in more than one terminal assistant bubble whenever the model emits answer text that never gains a tool call, so `core/messages/workspace-change-anchor.ts` picks the run's last assistant bubble and `MessageListItem` renders the badge only for that anchor (#4555). Any future run-scoped display belongs in the same place — do not hang one off every message. The two anchor helpers deliberately differ in which group types they accept as a run's last position, because an anchor is only useful where the display is actually rendered: run duration is emitted by `MessageList` around every group, so it accepts any type, while the workspace-change card comes from `MessageListItem` and so restricts to `assistant`. Keep a new helper's candidate set matched to its own render site rather than unifying them.
|
||||
|
||||
@ -10,6 +10,7 @@ import { SidebarTrigger } from "@/components/ui/sidebar";
|
||||
import { AgentWelcome } from "@/components/workspace/agent-welcome";
|
||||
import { ArtifactTrigger } from "@/components/workspace/artifacts";
|
||||
import { ChatBox, useThreadChat } from "@/components/workspace/chats";
|
||||
import { ContextUsageBadge } from "@/components/workspace/context-usage-badge";
|
||||
import { ExportTrigger } from "@/components/workspace/export-trigger";
|
||||
import { GoalStatus } from "@/components/workspace/goal-status";
|
||||
import {
|
||||
@ -47,7 +48,10 @@ import {
|
||||
useThreadStream,
|
||||
useThreadTokenUsage,
|
||||
} from "@/core/threads/hooks";
|
||||
import { threadTokenUsageToTokenUsage } from "@/core/threads/token-usage";
|
||||
import {
|
||||
selectContextUsage,
|
||||
threadTokenUsageToTokenUsage,
|
||||
} from "@/core/threads/token-usage";
|
||||
import { textOfMessage } from "@/core/threads/utils";
|
||||
import { env } from "@/env";
|
||||
import { cn } from "@/lib/utils";
|
||||
@ -73,13 +77,14 @@ export default function AgentChatPage() {
|
||||
const { tokenUsageEnabled } = useModels();
|
||||
const threadTokenUsage = useThreadTokenUsage(
|
||||
isNewThread || isMock ? undefined : threadId,
|
||||
{ enabled: tokenUsageEnabled && !isMock },
|
||||
{ enabled: !isMock },
|
||||
);
|
||||
const threadMetadata = useThreadMetadata(threadId, {
|
||||
enabled: !isNewThread && !isMock,
|
||||
isMock,
|
||||
});
|
||||
const backendTokenUsage = threadTokenUsageToTokenUsage(threadTokenUsage.data);
|
||||
const contextUsage = selectContextUsage(threadTokenUsage.data);
|
||||
|
||||
const { showNotification } = useNotification();
|
||||
|
||||
@ -276,17 +281,22 @@ export default function AgentChatPage() {
|
||||
<span className="hidden sm:inline">{t.agents.newChat}</span>
|
||||
</Button>
|
||||
</Tooltip>
|
||||
<TokenUsageIndicator
|
||||
threadId={isNewThread ? undefined : threadId}
|
||||
backendUsage={backendTokenUsage}
|
||||
enabled={tokenUsageEnabled}
|
||||
messages={thread.messages}
|
||||
pendingMessages={pendingUsageMessages}
|
||||
preferences={localSettings.tokenUsage}
|
||||
onPreferencesChange={(preferences) =>
|
||||
setLocalSettings("tokenUsage", preferences)
|
||||
}
|
||||
/>
|
||||
{tokenUsageEnabled ? (
|
||||
<TokenUsageIndicator
|
||||
threadId={isNewThread ? undefined : threadId}
|
||||
backendUsage={backendTokenUsage}
|
||||
contextUsage={contextUsage}
|
||||
enabled={tokenUsageEnabled}
|
||||
messages={thread.messages}
|
||||
pendingMessages={pendingUsageMessages}
|
||||
preferences={localSettings.tokenUsage}
|
||||
onPreferencesChange={(preferences) =>
|
||||
setLocalSettings("tokenUsage", preferences)
|
||||
}
|
||||
/>
|
||||
) : (
|
||||
<ContextUsageBadge contextUsage={contextUsage} />
|
||||
)}
|
||||
<SidecarTrigger />
|
||||
<ExportTrigger threadId={threadId} />
|
||||
<ArtifactTrigger />
|
||||
|
||||
@ -13,6 +13,7 @@ import {
|
||||
useSpecificChatMode,
|
||||
useThreadChat,
|
||||
} from "@/components/workspace/chats";
|
||||
import { ContextUsageBadge } from "@/components/workspace/context-usage-badge";
|
||||
import { ExportTrigger } from "@/components/workspace/export-trigger";
|
||||
import { GoalStatus } from "@/components/workspace/goal-status";
|
||||
import {
|
||||
@ -52,7 +53,10 @@ import {
|
||||
useThreadStream,
|
||||
useThreadTokenUsage,
|
||||
} from "@/core/threads/hooks";
|
||||
import { threadTokenUsageToTokenUsage } from "@/core/threads/token-usage";
|
||||
import {
|
||||
selectContextUsage,
|
||||
threadTokenUsageToTokenUsage,
|
||||
} from "@/core/threads/token-usage";
|
||||
import { textOfMessage } from "@/core/threads/utils";
|
||||
import { env } from "@/env";
|
||||
import { cn } from "@/lib/utils";
|
||||
@ -74,7 +78,7 @@ export default function ChatPage() {
|
||||
const { tokenUsageEnabled } = useModels();
|
||||
const threadTokenUsage = useThreadTokenUsage(
|
||||
isNewThread || isMock ? undefined : threadId,
|
||||
{ enabled: tokenUsageEnabled && !isMock },
|
||||
{ enabled: !isMock },
|
||||
);
|
||||
const threadMetadata = useThreadMetadata(threadId, {
|
||||
enabled: !isNewThread && !isMock,
|
||||
@ -82,6 +86,7 @@ export default function ChatPage() {
|
||||
});
|
||||
const branchThread = useBranchThread();
|
||||
const backendTokenUsage = threadTokenUsageToTokenUsage(threadTokenUsage.data);
|
||||
const contextUsage = selectContextUsage(threadTokenUsage.data);
|
||||
const mountedRef = useRef(false);
|
||||
useSpecificChatMode();
|
||||
|
||||
@ -288,17 +293,22 @@ export default function ChatPage() {
|
||||
{!isNewThread && (
|
||||
<ThreadScheduledTasksLink threadId={threadId} />
|
||||
)}
|
||||
<TokenUsageIndicator
|
||||
threadId={isNewThread ? undefined : threadId}
|
||||
backendUsage={backendTokenUsage}
|
||||
enabled={tokenUsageEnabled}
|
||||
messages={thread.messages}
|
||||
pendingMessages={pendingUsageMessages}
|
||||
preferences={localSettings.tokenUsage}
|
||||
onPreferencesChange={(preferences) =>
|
||||
setLocalSettings("tokenUsage", preferences)
|
||||
}
|
||||
/>
|
||||
{tokenUsageEnabled ? (
|
||||
<TokenUsageIndicator
|
||||
threadId={isNewThread ? undefined : threadId}
|
||||
backendUsage={backendTokenUsage}
|
||||
contextUsage={contextUsage}
|
||||
enabled={tokenUsageEnabled}
|
||||
messages={thread.messages}
|
||||
pendingMessages={pendingUsageMessages}
|
||||
preferences={localSettings.tokenUsage}
|
||||
onPreferencesChange={(preferences) =>
|
||||
setLocalSettings("tokenUsage", preferences)
|
||||
}
|
||||
/>
|
||||
) : (
|
||||
<ContextUsageBadge contextUsage={contextUsage} />
|
||||
)}
|
||||
<SidecarTrigger />
|
||||
{browserEnabled && <BrowserTrigger />}
|
||||
<ExportTrigger threadId={threadId} />
|
||||
|
||||
54
frontend/src/components/workspace/context-usage-badge.tsx
Normal file
54
frontend/src/components/workspace/context-usage-badge.tsx
Normal file
@ -0,0 +1,54 @@
|
||||
"use client";
|
||||
|
||||
import { GaugeIcon } from "lucide-react";
|
||||
|
||||
import { useI18n } from "@/core/i18n/hooks";
|
||||
import type { ContextUsage } from "@/core/threads/token-usage";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
import { formatContextUsagePercentage } from "./context-usage-format";
|
||||
|
||||
interface ContextUsageBadgeProps {
|
||||
contextUsage: ContextUsage | null;
|
||||
className?: string;
|
||||
}
|
||||
|
||||
export function ContextUsageBadge({
|
||||
contextUsage,
|
||||
className,
|
||||
}: ContextUsageBadgeProps) {
|
||||
const { t } = useI18n();
|
||||
|
||||
const formatted = formatContextUsagePercentage(contextUsage?.percentage);
|
||||
if (formatted == null) {
|
||||
return (
|
||||
<div
|
||||
role="status"
|
||||
data-context-usage-placeholder="true"
|
||||
aria-label={t.contextUsage.title}
|
||||
title={t.contextUsage.title}
|
||||
className={cn(
|
||||
"text-muted-foreground bg-background/70 flex size-7 items-center justify-center rounded-full border",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
<GaugeIcon size={14} />
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div
|
||||
role="status"
|
||||
aria-label={t.contextUsage.badgeAriaLabel(formatted)}
|
||||
className={cn(
|
||||
"text-muted-foreground bg-background/70 flex h-auto items-center gap-1.5 rounded-full border px-2 py-1 text-xs font-normal",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
<GaugeIcon size={14} />
|
||||
<span>{t.contextUsage.label}</span>
|
||||
<span className="font-mono">{formatted}%</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
19
frontend/src/components/workspace/context-usage-format.ts
Normal file
19
frontend/src/components/workspace/context-usage-format.ts
Normal file
@ -0,0 +1,19 @@
|
||||
/**
|
||||
* Format a context-usage percentage for display.
|
||||
*
|
||||
* Returns `null` when the percentage is unknown so callers can choose to
|
||||
* hide the indicator entirely rather than render a placeholder.
|
||||
*
|
||||
* Whole-number percentages are rendered without a decimal point; otherwise
|
||||
* a single decimal place is shown to stay readable at high resolution
|
||||
* (e.g. ``35.4%``).
|
||||
*/
|
||||
export function formatContextUsagePercentage(
|
||||
percentage: number | null | undefined,
|
||||
): string | null {
|
||||
if (typeof percentage !== "number" || !Number.isFinite(percentage)) {
|
||||
return null;
|
||||
}
|
||||
const clamped = Math.max(0, percentage);
|
||||
return Number.isInteger(clamped) ? `${clamped}` : clamped.toFixed(1);
|
||||
}
|
||||
@ -26,13 +26,17 @@ import {
|
||||
type TokenUsagePreferences,
|
||||
type TokenUsageViewPreset,
|
||||
} from "@/core/messages/usage-model";
|
||||
import type { ContextUsage } from "@/core/threads/token-usage";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
import { formatContextUsagePercentage } from "./context-usage-format";
|
||||
|
||||
interface TokenUsageIndicatorProps {
|
||||
threadId?: string;
|
||||
messages: Message[];
|
||||
pendingMessages?: Message[];
|
||||
backendUsage?: TokenUsage | null;
|
||||
contextUsage?: ContextUsage | null;
|
||||
enabled?: boolean;
|
||||
preferences: TokenUsagePreferences;
|
||||
onPreferencesChange: (preferences: TokenUsagePreferences) => void;
|
||||
@ -44,6 +48,7 @@ export function TokenUsageIndicator({
|
||||
messages,
|
||||
pendingMessages,
|
||||
backendUsage,
|
||||
contextUsage,
|
||||
enabled = false,
|
||||
preferences,
|
||||
onPreferencesChange,
|
||||
@ -61,6 +66,9 @@ export function TokenUsageIndicator({
|
||||
[backendUsage, messages, pendingMessages, threadId],
|
||||
);
|
||||
const preset = getTokenUsageViewPreset(preferences);
|
||||
const contextPercentage = formatContextUsagePercentage(
|
||||
contextUsage?.percentage,
|
||||
);
|
||||
|
||||
if (!enabled) {
|
||||
return null;
|
||||
@ -86,6 +94,14 @@ export function TokenUsageIndicator({
|
||||
: "-"
|
||||
: t.tokenUsage.presets[presetKeyToTranslationKey(preset)]}
|
||||
</span>
|
||||
{contextPercentage != null && (
|
||||
<span
|
||||
className="text-muted-foreground/80 border-l pl-1.5 font-mono"
|
||||
aria-label={t.contextUsage.badgeAriaLabel(contextPercentage)}
|
||||
>
|
||||
{contextPercentage}%
|
||||
</span>
|
||||
)}
|
||||
<ChevronDownIcon className="size-3" />
|
||||
</Button>
|
||||
</DropdownMenuTrigger>
|
||||
|
||||
@ -699,6 +699,13 @@ export const enUS: Translations = {
|
||||
removeTodo: (content: string) => `Remove To-do: ${content}`,
|
||||
},
|
||||
|
||||
contextUsage: {
|
||||
label: "Context",
|
||||
title: "Context window",
|
||||
badgeAriaLabel: (percentage: string) =>
|
||||
`Context window ${percentage}% full`,
|
||||
},
|
||||
|
||||
// Shortcuts
|
||||
shortcuts: {
|
||||
searchActions: "Search actions...",
|
||||
|
||||
@ -582,6 +582,12 @@ export interface Translations {
|
||||
removeTodo: (content: string) => string;
|
||||
};
|
||||
|
||||
contextUsage: {
|
||||
label: string;
|
||||
title: string;
|
||||
badgeAriaLabel: (percentage: string) => string;
|
||||
};
|
||||
|
||||
// Shortcuts
|
||||
shortcuts: {
|
||||
searchActions: string;
|
||||
|
||||
@ -672,6 +672,12 @@ export const zhCN: Translations = {
|
||||
removeTodo: (content: string) => `移除 To-do:${content}`,
|
||||
},
|
||||
|
||||
contextUsage: {
|
||||
label: "上下文",
|
||||
title: "上下文窗口",
|
||||
badgeAriaLabel: (percentage: string) => `上下文窗口已使用 ${percentage}%`,
|
||||
},
|
||||
|
||||
// Shortcuts
|
||||
shortcuts: {
|
||||
searchActions: "搜索操作...",
|
||||
|
||||
@ -41,7 +41,10 @@ import {
|
||||
filterThreadSearchResults,
|
||||
type ThreadSearchParams,
|
||||
} from "./thread-search-query";
|
||||
import { threadTokenUsageQueryKey } from "./token-usage";
|
||||
import {
|
||||
retainThreadTokenUsagePlaceholder,
|
||||
threadTokenUsageQueryKey,
|
||||
} from "./token-usage";
|
||||
import type {
|
||||
AgentThread,
|
||||
AgentThreadState,
|
||||
@ -2771,6 +2774,10 @@ export function useThreadTokenUsage(
|
||||
enabled: enabled && Boolean(threadId),
|
||||
retry: false,
|
||||
refetchOnWindowFocus: false,
|
||||
// Keep same-thread data visible during refetches without carrying usage
|
||||
// from the previous route into a newly selected thread.
|
||||
placeholderData: (previous) =>
|
||||
retainThreadTokenUsagePlaceholder(previous, threadId),
|
||||
});
|
||||
}
|
||||
|
||||
|
||||
@ -6,6 +6,13 @@ export function threadTokenUsageQueryKey(threadId?: string | null) {
|
||||
return ["thread-token-usage", threadId] as const;
|
||||
}
|
||||
|
||||
export function retainThreadTokenUsagePlaceholder(
|
||||
previous: ThreadTokenUsageResponse | null | undefined,
|
||||
threadId?: string | null,
|
||||
): ThreadTokenUsageResponse | undefined {
|
||||
return previous && previous.thread_id === threadId ? previous : undefined;
|
||||
}
|
||||
|
||||
export function threadTokenUsageToTokenUsage(
|
||||
usage: ThreadTokenUsageResponse | null | undefined,
|
||||
): TokenUsage | null {
|
||||
@ -18,3 +25,23 @@ export function threadTokenUsageToTokenUsage(
|
||||
totalTokens: usage.total_tokens ?? 0,
|
||||
};
|
||||
}
|
||||
|
||||
export interface ContextUsage {
|
||||
tokenCount: number;
|
||||
maxContextTokens: number | null;
|
||||
percentage: number | null;
|
||||
}
|
||||
|
||||
export function selectContextUsage(
|
||||
usage: ThreadTokenUsageResponse | null | undefined,
|
||||
): ContextUsage | null {
|
||||
if (!usage?.context_usage) {
|
||||
return null;
|
||||
}
|
||||
const { token_count, max_context_tokens, percentage } = usage.context_usage;
|
||||
return {
|
||||
tokenCount: token_count ?? 0,
|
||||
maxContextTokens: max_context_tokens ?? null,
|
||||
percentage: percentage ?? null,
|
||||
};
|
||||
}
|
||||
|
||||
@ -62,6 +62,12 @@ export interface RunMessage {
|
||||
created_at: string;
|
||||
}
|
||||
|
||||
export interface ThreadContextUsage {
|
||||
token_count: number;
|
||||
max_context_tokens: number | null;
|
||||
percentage: number | null;
|
||||
}
|
||||
|
||||
export interface ThreadTokenUsageResponse {
|
||||
thread_id: string;
|
||||
total_tokens: number;
|
||||
@ -74,4 +80,5 @@ export interface ThreadTokenUsageResponse {
|
||||
subagent: number;
|
||||
middleware: number;
|
||||
};
|
||||
context_usage?: ThreadContextUsage | null;
|
||||
}
|
||||
|
||||
@ -0,0 +1,42 @@
|
||||
import { expect, rs, test } from "@rstest/core";
|
||||
import { createElement } from "react";
|
||||
import { renderToStaticMarkup } from "react-dom/server";
|
||||
|
||||
import { ContextUsageBadge } from "@/components/workspace/context-usage-badge";
|
||||
|
||||
rs.mock("@/core/i18n/hooks", () => ({
|
||||
useI18n: () => ({
|
||||
t: {
|
||||
contextUsage: {
|
||||
title: "Context window",
|
||||
label: "Context",
|
||||
badgeAriaLabel: (percentage: string) =>
|
||||
`Context window ${percentage}% full`,
|
||||
},
|
||||
},
|
||||
}),
|
||||
}));
|
||||
|
||||
test("keeps a gauge placeholder visible while context usage is unavailable", () => {
|
||||
const html = renderToStaticMarkup(
|
||||
createElement(ContextUsageBadge, { contextUsage: null }),
|
||||
);
|
||||
|
||||
expect(html).toContain('data-context-usage-placeholder="true"');
|
||||
expect(html).toContain('aria-label="Context window"');
|
||||
});
|
||||
|
||||
test("renders the current percentage when context usage is available", () => {
|
||||
const html = renderToStaticMarkup(
|
||||
createElement(ContextUsageBadge, {
|
||||
contextUsage: {
|
||||
tokenCount: 35_000,
|
||||
maxContextTokens: 100_000,
|
||||
percentage: 35,
|
||||
},
|
||||
}),
|
||||
);
|
||||
|
||||
expect(html).toContain("35%");
|
||||
expect(html).toContain('aria-label="Context window 35% full"');
|
||||
});
|
||||
@ -0,0 +1,24 @@
|
||||
import { expect, test } from "@rstest/core";
|
||||
|
||||
import { formatContextUsagePercentage } from "@/components/workspace/context-usage-format";
|
||||
|
||||
test("returns null for unknown percentages", () => {
|
||||
expect(formatContextUsagePercentage(null)).toBeNull();
|
||||
expect(formatContextUsagePercentage(undefined)).toBeNull();
|
||||
expect(formatContextUsagePercentage(Number.NaN)).toBeNull();
|
||||
});
|
||||
|
||||
test("renders whole numbers without a decimal point", () => {
|
||||
expect(formatContextUsagePercentage(0)).toBe("0");
|
||||
expect(formatContextUsagePercentage(35)).toBe("35");
|
||||
expect(formatContextUsagePercentage(100)).toBe("100");
|
||||
});
|
||||
|
||||
test("renders fractional percentages with one decimal place", () => {
|
||||
expect(formatContextUsagePercentage(35.4)).toBe("35.4");
|
||||
expect(formatContextUsagePercentage(35.46)).toBe("35.5");
|
||||
});
|
||||
|
||||
test("clamps negative percentages to zero", () => {
|
||||
expect(formatContextUsagePercentage(-1)).toBe("0");
|
||||
});
|
||||
@ -1,6 +1,12 @@
|
||||
import { expect, test } from "@rstest/core";
|
||||
import { QueryClient, QueryObserver } from "@tanstack/react-query";
|
||||
|
||||
import { threadTokenUsageToTokenUsage } from "@/core/threads/token-usage";
|
||||
import {
|
||||
retainThreadTokenUsagePlaceholder,
|
||||
selectContextUsage,
|
||||
threadTokenUsageQueryKey,
|
||||
threadTokenUsageToTokenUsage,
|
||||
} from "@/core/threads/token-usage";
|
||||
import type { ThreadTokenUsageResponse } from "@/core/threads/types";
|
||||
|
||||
test("maps backend thread token usage to UI token usage", () => {
|
||||
@ -29,3 +35,150 @@ test("returns null when backend thread token usage is unavailable", () => {
|
||||
expect(threadTokenUsageToTokenUsage(null)).toBeNull();
|
||||
expect(threadTokenUsageToTokenUsage(undefined)).toBeNull();
|
||||
});
|
||||
|
||||
test("retains placeholder usage only for the current thread", () => {
|
||||
const response: ThreadTokenUsageResponse = {
|
||||
thread_id: "thread-1",
|
||||
total_input_tokens: 90,
|
||||
total_output_tokens: 60,
|
||||
total_tokens: 150,
|
||||
total_runs: 2,
|
||||
by_model: { unknown: { tokens: 150, runs: 2 } },
|
||||
by_caller: {
|
||||
lead_agent: 120,
|
||||
subagent: 25,
|
||||
middleware: 5,
|
||||
},
|
||||
};
|
||||
|
||||
expect(retainThreadTokenUsagePlaceholder(response, "thread-1")).toBe(
|
||||
response,
|
||||
);
|
||||
expect(
|
||||
retainThreadTokenUsagePlaceholder(response, "thread-2"),
|
||||
).toBeUndefined();
|
||||
expect(retainThreadTokenUsagePlaceholder(null, undefined)).toBeUndefined();
|
||||
});
|
||||
|
||||
test("query observer keeps same-thread data but drops it while a new thread is pending", async () => {
|
||||
const queryClient = new QueryClient({
|
||||
defaultOptions: { queries: { retry: false } },
|
||||
});
|
||||
const threadA: ThreadTokenUsageResponse = {
|
||||
thread_id: "thread-a",
|
||||
total_input_tokens: 90,
|
||||
total_output_tokens: 60,
|
||||
total_tokens: 150,
|
||||
total_runs: 2,
|
||||
by_model: { unknown: { tokens: 150, runs: 2 } },
|
||||
by_caller: { lead_agent: 120, subagent: 25, middleware: 5 },
|
||||
};
|
||||
const threadB: ThreadTokenUsageResponse = {
|
||||
...threadA,
|
||||
thread_id: "thread-b",
|
||||
total_tokens: 200,
|
||||
};
|
||||
let queryResult = Promise.resolve(threadA);
|
||||
let resolveRefresh: (value: ThreadTokenUsageResponse) => void = () =>
|
||||
undefined;
|
||||
let resolveThreadB: (value: ThreadTokenUsageResponse) => void = () =>
|
||||
undefined;
|
||||
const observer = new QueryObserver<ThreadTokenUsageResponse | null>(
|
||||
queryClient,
|
||||
{
|
||||
queryKey: threadTokenUsageQueryKey("thread-a"),
|
||||
queryFn: () => queryResult,
|
||||
placeholderData: (previous) =>
|
||||
retainThreadTokenUsagePlaceholder(previous, "thread-a"),
|
||||
},
|
||||
);
|
||||
const unsubscribe = observer.subscribe(() => undefined);
|
||||
|
||||
try {
|
||||
await observer.refetch();
|
||||
expect(observer.getCurrentResult().data).toBe(threadA);
|
||||
|
||||
queryResult = new Promise((resolve) => {
|
||||
resolveRefresh = resolve;
|
||||
});
|
||||
const sameThreadRefetch = observer.refetch();
|
||||
expect(observer.getCurrentResult().data).toBe(threadA);
|
||||
resolveRefresh(threadA);
|
||||
await sameThreadRefetch;
|
||||
|
||||
const pendingThreadB = new Promise<ThreadTokenUsageResponse>((resolve) => {
|
||||
resolveThreadB = resolve;
|
||||
});
|
||||
observer.setOptions({
|
||||
queryKey: threadTokenUsageQueryKey("thread-b"),
|
||||
queryFn: () => pendingThreadB,
|
||||
retry: false,
|
||||
placeholderData: (previous) =>
|
||||
retainThreadTokenUsagePlaceholder(previous, "thread-b"),
|
||||
});
|
||||
expect(observer.getCurrentResult().data).toBeUndefined();
|
||||
expect(observer.getCurrentResult().isPlaceholderData).toBe(false);
|
||||
|
||||
resolveThreadB(threadB);
|
||||
await observer.refetch();
|
||||
expect(observer.getCurrentResult().data).toBe(threadB);
|
||||
} finally {
|
||||
resolveRefresh(threadA);
|
||||
resolveThreadB(threadB);
|
||||
unsubscribe();
|
||||
queryClient.clear();
|
||||
}
|
||||
});
|
||||
|
||||
const _baseResponse = {
|
||||
thread_id: "thread-1",
|
||||
total_input_tokens: 0,
|
||||
total_output_tokens: 0,
|
||||
total_tokens: 0,
|
||||
total_runs: 0,
|
||||
by_model: {},
|
||||
by_caller: { lead_agent: 0, subagent: 0, middleware: 0 },
|
||||
} satisfies ThreadTokenUsageResponse;
|
||||
|
||||
test("selectContextUsage projects the backend block to UI shape", () => {
|
||||
const response: ThreadTokenUsageResponse = {
|
||||
..._baseResponse,
|
||||
context_usage: {
|
||||
token_count: 350,
|
||||
max_context_tokens: 1000,
|
||||
percentage: 35,
|
||||
},
|
||||
};
|
||||
|
||||
expect(selectContextUsage(response)).toEqual({
|
||||
tokenCount: 350,
|
||||
maxContextTokens: 1000,
|
||||
percentage: 35,
|
||||
});
|
||||
});
|
||||
|
||||
test("selectContextUsage preserves nullable capacity and percentage", () => {
|
||||
const response: ThreadTokenUsageResponse = {
|
||||
..._baseResponse,
|
||||
context_usage: {
|
||||
token_count: 200,
|
||||
max_context_tokens: null,
|
||||
percentage: null,
|
||||
},
|
||||
};
|
||||
|
||||
expect(selectContextUsage(response)).toEqual({
|
||||
tokenCount: 200,
|
||||
maxContextTokens: null,
|
||||
percentage: null,
|
||||
});
|
||||
});
|
||||
|
||||
test("selectContextUsage returns null when context_usage is missing", () => {
|
||||
expect(selectContextUsage(_baseResponse)).toBeNull();
|
||||
expect(
|
||||
selectContextUsage({ ..._baseResponse, context_usage: null }),
|
||||
).toBeNull();
|
||||
expect(selectContextUsage(null)).toBeNull();
|
||||
expect(selectContextUsage(undefined)).toBeNull();
|
||||
});
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user