mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-10 00:18:36 +00:00
fix(token-budget): acquire _lock in _apply, before_agent, and _drain_pending_warnings (#3714)
_lock was defined in __init__ and used only in reset(), leaving _apply(), before_agent(), _clear_run_state(), and _drain_pending_warnings() unprotected. Concurrent after_model callbacks for the same run_id could race on usage_accum.input += diff_input, causing lost-update on token counters. Also fix stray bare 'a' character in frontend/src/content/zh/introduction/index.mdx that rendered as visible text on the documentation page. Co-authored-by: kapilvus <kapilvus@gmail.com>
This commit is contained in:
kapil971390
GitHub
parent
e847d738d3
commit
d050fba07f
2 changed files with 72 additions and 69 deletions
|
|
@ -96,31 +96,33 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
|
|||
return str(id(runtime))
|
||||
|
||||
def _clear_run_state(self, run_id: str) -> None:
|
||||
self._warned.pop(run_id, None)
|
||||
self._pending_warnings.pop(run_id, None)
|
||||
self._seen_messages.pop(run_id, None)
|
||||
self._cumulative_usage.pop(run_id, None)
|
||||
with self._lock:
|
||||
self._warned.pop(run_id, None)
|
||||
self._pending_warnings.pop(run_id, None)
|
||||
self._seen_messages.pop(run_id, None)
|
||||
self._cumulative_usage.pop(run_id, None)
|
||||
|
||||
@override
|
||||
def before_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
if not self._config.enabled:
|
||||
return
|
||||
|
||||
# Mark all old messages from previous runss as 'seen' so they don't count toward THIS run's budget
|
||||
# Mark all old messages from previous runs as 'seen' so they don't count toward THIS run's budget
|
||||
messages = state.get("messages", [])
|
||||
if not messages:
|
||||
return
|
||||
|
||||
run_id = self._get_run_id(runtime)
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
with self._lock:
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
|
||||
@override
|
||||
async def abefore_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
|
|
@ -182,68 +184,69 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
|
|||
|
||||
run_id = self._get_run_id(runtime)
|
||||
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
usage_accum = self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
with self._lock:
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
usage_accum = self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
|
||||
# Check what previously recorded for this exact message
|
||||
prev_input, prev_output = seen.get(msg.id, (0, 0))
|
||||
# Check what previously recorded for this exact message
|
||||
prev_input, prev_output = seen.get(msg.id, (0, 0))
|
||||
|
||||
# Calculate if any new tokens were added (handles retroactive subagent tokens)
|
||||
diff_input = max(0, input_tokens - prev_input)
|
||||
diff_output = max(0, output_tokens - prev_output)
|
||||
# Calculate if any new tokens were added (handles retroactive subagent tokens)
|
||||
diff_input = max(0, input_tokens - prev_input)
|
||||
diff_output = max(0, output_tokens - prev_output)
|
||||
|
||||
if diff_input > 0 or diff_output > 0:
|
||||
usage_accum.input += diff_input
|
||||
usage_accum.output += diff_output
|
||||
usage_accum.total += diff_input + diff_output
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
if diff_input > 0 or diff_output > 0:
|
||||
usage_accum.input += diff_input
|
||||
usage_accum.output += diff_output
|
||||
usage_accum.total += diff_input + diff_output
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
|
||||
if usage_accum.total <= 0:
|
||||
return None
|
||||
|
||||
fractions = [("total", usage_accum.total, self._config.max_tokens)]
|
||||
if self._config.max_input_tokens:
|
||||
fractions.append(("input", usage_accum.input, self._config.max_input_tokens))
|
||||
if self._config.max_output_tokens:
|
||||
fractions.append(("output", usage_accum.output, self._config.max_output_tokens))
|
||||
|
||||
highest_fraction = 0.0
|
||||
trigger_reason = ""
|
||||
trigger_used = 0
|
||||
trigger_budget = 0
|
||||
|
||||
for reason, used, limit in fractions:
|
||||
frac = used / limit
|
||||
if frac > highest_fraction:
|
||||
highest_fraction = frac
|
||||
trigger_reason = reason
|
||||
trigger_used = used
|
||||
trigger_budget = limit
|
||||
|
||||
if highest_fraction >= self._config.hard_stop_threshold:
|
||||
logger.warning("Token budget hard stop triggered for run %s: %s limit exceeded", run_id, trigger_reason)
|
||||
stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget)
|
||||
return self._build_hard_stop_update(last_msg, stop_text)
|
||||
|
||||
if highest_fraction >= self._config.warn_threshold and not self._warned.get(run_id, False):
|
||||
self._warned[run_id] = True
|
||||
percent = highest_fraction * 100
|
||||
warn_text = _BUDGET_WARNING_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget, percent=percent)
|
||||
logger.info("Token budget warning triggered for run %s: %s limit at %.1f%%", run_id, trigger_reason, percent)
|
||||
# queue warning for wrap_model_call
|
||||
warnings = self._pending_warnings.setdefault(run_id, [])
|
||||
warnings.append(warn_text)
|
||||
return None
|
||||
|
||||
if usage_accum.total <= 0:
|
||||
return None
|
||||
|
||||
fractions = [("total", usage_accum.total, self._config.max_tokens)]
|
||||
if self._config.max_input_tokens:
|
||||
fractions.append(("input", usage_accum.input, self._config.max_input_tokens))
|
||||
if self._config.max_output_tokens:
|
||||
fractions.append(("output", usage_accum.output, self._config.max_output_tokens))
|
||||
|
||||
highest_fraction = 0.0
|
||||
trigger_reason = ""
|
||||
trigger_used = 0
|
||||
trigger_budget = 0
|
||||
|
||||
for reason, used, limit in fractions:
|
||||
frac = used / limit
|
||||
if frac > highest_fraction:
|
||||
highest_fraction = frac
|
||||
trigger_reason = reason
|
||||
trigger_used = used
|
||||
trigger_budget = limit
|
||||
|
||||
if highest_fraction >= self._config.hard_stop_threshold:
|
||||
logger.warning("Token budget hard stop triggered for run %s: %s limit exceeded", run_id, trigger_reason)
|
||||
stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget)
|
||||
return self._build_hard_stop_update(last_msg, stop_text)
|
||||
|
||||
if highest_fraction >= self._config.warn_threshold and not self._warned.get(run_id, False):
|
||||
self._warned[run_id] = True
|
||||
percent = highest_fraction * 100
|
||||
warn_text = _BUDGET_WARNING_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget, percent=percent)
|
||||
logger.info("Token budget warning triggered for run %s: %s limit at %.1f%%", run_id, trigger_reason, percent)
|
||||
# queue warning for wrap_model_call
|
||||
warnings = self._pending_warnings.setdefault(run_id, [])
|
||||
warnings.append(warn_text)
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
@override
|
||||
def after_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
||||
return self._apply(state, runtime)
|
||||
|
|
@ -257,7 +260,8 @@ class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
|
|||
return []
|
||||
|
||||
run_id = self._get_run_id(runtime)
|
||||
warnings = self._pending_warnings.pop(run_id, None)
|
||||
with self._lock:
|
||||
warnings = self._pending_warnings.pop(run_id, None)
|
||||
return warnings or []
|
||||
|
||||
def _inject_warnings(self, request: ModelRequest, warnings: list[str]) -> ModelRequest:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
title: 简介
|
||||
description: DeerFlow 是一个运行时 Harness，用于构建支持技能、记忆、工具和沙箱执行的长时序 Agent 系统。
|
||||
---
|
||||
a
|
||||
import { Cards } from "nextra/components";
|
||||
|
||||
# 简介
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue