diff --git a/agents/translation_node.py b/agents/translation_node.py
index 6ad93edc..2867066f 100644
--- a/agents/translation_node.py
+++ b/agents/translation_node.py
@@ -8,8 +8,8 @@
settings,
)
-from utils.runtime_state_sync import (
- refresh_runtime_state,
+from utils.token_usage import (
+ record_token_usage,
)
@@ -45,20 +45,32 @@ async def run(
translated_text
)
- await refresh_runtime_state(
- context,
- runtime_id=(
- settings.TRANSLATOR_MODEL_UID
- ),
- add_tokens=usage.get(
- "total_tokens",
- 0,
- ),
- max_tokens=(
- settings.TRANSLATOR_CONTEXT_WINDOW
- ),
- last_error=None,
- status="online",
- )
+ if usage:
+ record_token_usage(
+ context,
+ runtime_id=(
+ settings.TRANSLATOR_MODEL_UID
+ ),
+ role="translator",
+ kind="service",
+ prompt_tokens=(
+ usage.get(
+ "prompt_tokens",
+ 0,
+ )
+ ),
+ completion_tokens=(
+ usage.get(
+ "completion_tokens",
+ 0,
+ )
+ ),
+ total_tokens=(
+ usage.get(
+ "total_tokens",
+ 0,
+ )
+ ),
+ )
state.translated_input = translated_text
diff --git a/contracts/context_contract.py b/contracts/context_contract.py
index aed4beb4..d263cc11 100644
--- a/contracts/context_contract.py
+++ b/contracts/context_contract.py
@@ -13,33 +13,14 @@
)
-def cdata(
- value: str,
-) -> str:
-
- return (
- "', ']]]]>')}"
- "]]>"
- )
-
-
def format_xml_field(
tag: str,
value,
) -> str:
- if tag.endswith(
- "_ACTION"
- ):
- rendered_value = cdata(
- str(value)
- )
-
- else:
- rendered_value = escape(
- str(value)
- )
+ rendered_value = escape(
+ str(value)
+ )
return f"<{tag}>{rendered_value}{tag}>"
@@ -55,7 +36,7 @@ def format_available_actions(
(
" "
f""
- f"{cdata(template)}"
+ f"{template}"
""
)
for name, template
@@ -96,13 +77,12 @@ def build_runtime_fields(self) -> str:
fields = {
"TIMESTAMP": self.timestamp,
"WEEKDAY": self.weekday,
- "YEAR": self.year,
}
available_actions = []
if self.can_deep_thought:
- fields["DEEP_THOUGHT_COUNTER"] = self.deep_thought_count
+ fields["DEEP_THOUGHT_COUNTER"] = str(self.deep_thought_count)
available_actions.append(
(
RUNTIME_ACTION_DEEP_THOUGHT,
diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py
index b466a017..bf9c2c42 100644
--- a/runtime/runtime_context.py
+++ b/runtime/runtime_context.py
@@ -39,3 +39,7 @@ class RuntimeContext:
runtime_action_events: list[dict] = field(
default_factory=list
)
+
+ runtime_usage_events: list[dict] = field(
+ default_factory=list
+ )
diff --git a/runtime/runtime_stream.py b/runtime/runtime_stream.py
index 172f980e..78c14fb9 100644
--- a/runtime/runtime_stream.py
+++ b/runtime/runtime_stream.py
@@ -12,8 +12,12 @@
StreamHandler,
)
+from utils.token_usage import (
+ record_stream_token_usage,
+)
+
from utils.tokens import (
- estimate_stream_tokens,
+ estimate_stream_live_tokens,
)
@@ -45,6 +49,7 @@ def __init__(
self.log_method = log_method
self.emit_to_chat = emit_to_chat
+ self.context_snapshot = context_snapshot or {}
self.stream = StreamHandler(
self.websocket,
@@ -58,6 +63,113 @@ def __init__(
),
)
+ def build_input_prompt_text(self) -> str:
+
+ if not isinstance(
+ self.context_snapshot,
+ dict,
+ ):
+ return ""
+
+ parts = []
+
+ for key in (
+ "system_prompt",
+ "user_prompt",
+ "context_payload",
+ ):
+
+ value = self.context_snapshot.get(
+ key,
+ "",
+ )
+
+ if value:
+ parts.append(
+ str(value)
+ )
+
+ return "\n".join(
+ parts
+ )
+
+ def is_brain_context(self) -> bool:
+
+ if not isinstance(
+ self.context_snapshot,
+ dict,
+ ):
+ return False
+
+ return (
+ self.context_snapshot.get(
+ "context_role"
+ )
+ == "brain"
+ )
+
+ def estimate_live_tokens(self) -> int:
+
+ return estimate_stream_live_tokens(
+ self.stream,
+ prompt_text=(
+ self.build_input_prompt_text()
+ ),
+ )
+
+ async def refresh_token_usage(self):
+
+ if not self.is_brain_context():
+ return
+
+ used_tokens = (
+ self.estimate_live_tokens()
+ )
+
+ if not used_tokens:
+ return
+
+ await refresh_runtime_state(
+ self.context,
+ runtime_id=(
+ self.runtime_id
+ ),
+ used_tokens=(
+ used_tokens
+ ),
+ max_tokens=(
+ self.context_window
+ ),
+ last_error=None,
+ status="online",
+ )
+
+ def record_token_usage(self):
+ is_brain_context = self.is_brain_context()
+
+ record_stream_token_usage(
+ self.context,
+ runtime_id=(
+ self.runtime_id
+ ),
+ role=(
+ "brain"
+ if is_brain_context
+ else self.role
+ ),
+ kind=(
+ "brain"
+ if is_brain_context
+ else "service"
+ ),
+ stream=(
+ self.stream
+ ),
+ prompt_text=(
+ self.build_input_prompt_text()
+ ),
+ )
+
def build_action_log(
self,
action_event_offset: int,
@@ -133,6 +245,8 @@ async def run(
emit=self.emit_to_chat
)
+ await self.refresh_token_usage()
+
await self.logger.log_runtime(
f"[STREAM START] role={self.role}"
)
@@ -171,6 +285,8 @@ async def run(
emit=self.emit_to_chat,
)
+ await self.refresh_token_usage()
+
continue
# -------------------------------------------------
@@ -196,30 +312,14 @@ async def run(
return None
+ await self.refresh_token_usage()
+
await self.stream.finish(
emit=self.emit_to_chat
)
- used_tokens = (
- estimate_stream_tokens(
- self.stream
- )
- )
-
- await refresh_runtime_state(
- self.context,
- runtime_id=(
- self.runtime_id
- ),
- used_tokens=(
- used_tokens
- ),
- max_tokens=(
- self.context_window
- ),
- last_error=None,
- status="online",
- )
+ await self.refresh_token_usage()
+ self.record_token_usage()
log_response = self.stream.response
diff --git a/static/logger.js b/static/logger.js
index 4ea66116..1e6cd790 100644
--- a/static/logger.js
+++ b/static/logger.js
@@ -226,6 +226,11 @@ function appendLog(
"text-emerald-500";
}
+ if (tag.includes("FLOW TELEMETRY")) {
+ tagClass =
+ "text-purple-400";
+ }
+
if (tag.includes("ERROR")) {
tagClass =
"text-red-500 font-bold";
diff --git a/tests/test_brain_runtime_actions.py b/tests/test_brain_runtime_actions.py
index 6625ab02..e64aa33c 100644
--- a/tests/test_brain_runtime_actions.py
+++ b/tests/test_brain_runtime_actions.py
@@ -72,13 +72,18 @@ def test_prompt_uses_passed_agent_runtime_actions(self):
self.assertIn(
(
- ''
f"{SEARCH_ACTION_TEMPLATE}"
- "]]>"
+ ""
),
prompt,
)
+ self.assertNotIn(
+ "",
+ prompt,
+ )
+
self.assertNotIn(
"RUNTIME_STATE",
prompt,
@@ -135,13 +145,18 @@ def test_prompt_can_flip_agent_actions_dynamically(self):
self.assertIn(
(
- ''
f"{DEEP_THOUGHT_ACTION}"
- "]]>"
+ ""
),
prompt,
)
+ self.assertNotIn(
+ " int:
+
+ try:
+ return int(
+ value or 0
+ )
+
+ except (
+ TypeError,
+ ValueError,
+ ):
+ return 0
+
+
+def record_token_usage(
+ context,
+ *,
+ runtime_id: str,
+ role: str,
+ kind: str = "service",
+ prompt_tokens: int = 0,
+ completion_tokens: int = 0,
+ total_tokens: int = 0,
+ context_tokens: int = 0,
+):
+
+ usage_events = getattr(
+ context,
+ "runtime_usage_events",
+ None,
+ )
+
+ if usage_events is None:
+ usage_events = []
+ context.runtime_usage_events = usage_events
+
+ usage_events.append({
+ "runtime_id": runtime_id,
+ "role": role,
+ "kind": kind,
+ "prompt_tokens": _as_int(
+ prompt_tokens
+ ),
+ "completion_tokens": _as_int(
+ completion_tokens
+ ),
+ "total_tokens": _as_int(
+ total_tokens
+ ),
+ "context_tokens": _as_int(
+ context_tokens
+ ),
+ })
+
+
+def record_stream_token_usage(
+ context,
+ *,
+ runtime_id: str,
+ role: str,
+ kind: str = "service",
+ stream,
+ prompt_text: str = "",
+):
+
+ prompt_tokens = (
+ _as_int(
+ getattr(
+ stream,
+ "prompt_tokens",
+ 0,
+ )
+ )
+ or estimate_stream_input_tokens(
+ stream,
+ prompt_text=prompt_text,
+ )
+ )
+
+ completion_tokens = (
+ _as_int(
+ getattr(
+ stream,
+ "completion_tokens",
+ 0,
+ )
+ )
+ or (
+ estimate_optional_tokens(
+ getattr(
+ stream,
+ "response",
+ "",
+ )
+ )
+ + estimate_optional_tokens(
+ getattr(
+ stream,
+ "reasoning",
+ "",
+ )
+ )
+ )
+ )
+
+ total_tokens = (
+ _as_int(
+ getattr(
+ stream,
+ "total_tokens",
+ 0,
+ )
+ )
+ or prompt_tokens
+ + completion_tokens
+ )
+ context_tokens = estimate_stream_live_tokens(
+ stream,
+ prompt_text=prompt_text,
+ )
+
+ record_token_usage(
+ context,
+ runtime_id=runtime_id,
+ role=role,
+ kind=kind,
+ prompt_tokens=prompt_tokens,
+ completion_tokens=completion_tokens,
+ total_tokens=total_tokens,
+ context_tokens=context_tokens,
+ )
+
+
+def summarize_token_usage(
+ context,
+ *,
+ kind: str | None = None,
+) -> dict:
+
+ summary = {
+ "prompt_tokens": 0,
+ "completion_tokens": 0,
+ "total_tokens": 0,
+ }
+
+ for event in getattr(
+ context,
+ "runtime_usage_events",
+ [],
+ ):
+ if (
+ kind is not None
+ and event.get(
+ "kind"
+ )
+ != kind
+ ):
+ continue
+
+ summary["prompt_tokens"] += _as_int(
+ event.get(
+ "prompt_tokens",
+ 0,
+ )
+ )
+ summary["completion_tokens"] += _as_int(
+ event.get(
+ "completion_tokens",
+ 0,
+ )
+ )
+ summary["total_tokens"] += _as_int(
+ event.get(
+ "total_tokens",
+ 0,
+ )
+ )
+
+ return summary
+
+
+def summarize_token_usage_by_role(
+ context,
+ *,
+ kind: str | None = None,
+) -> list[dict]:
+
+ grouped = {}
+
+ for event in getattr(
+ context,
+ "runtime_usage_events",
+ [],
+ ):
+ if (
+ kind is not None
+ and event.get(
+ "kind"
+ )
+ != kind
+ ):
+ continue
+
+ key = (
+ event.get(
+ "role",
+ "unknown",
+ ),
+ event.get(
+ "runtime_id",
+ "unknown",
+ ),
+ )
+
+ if key not in grouped:
+ grouped[key] = {
+ "role": key[0],
+ "runtime_id": key[1],
+ "prompt_tokens": 0,
+ "completion_tokens": 0,
+ "total_tokens": 0,
+ "context_tokens": 0,
+ }
+
+ grouped[key]["prompt_tokens"] += _as_int(
+ event.get(
+ "prompt_tokens",
+ 0,
+ )
+ )
+ grouped[key]["completion_tokens"] += _as_int(
+ event.get(
+ "completion_tokens",
+ 0,
+ )
+ )
+ grouped[key]["total_tokens"] += _as_int(
+ event.get(
+ "total_tokens",
+ 0,
+ )
+ )
+ grouped[key]["context_tokens"] += _as_int(
+ event.get(
+ "context_tokens",
+ 0,
+ )
+ )
+
+ return list(
+ grouped.values()
+ )
+
+
+def format_token_usage_summary(
+ context,
+) -> str:
+
+ summary = summarize_token_usage(
+ context
+ )
+ breakdown = summarize_token_usage_by_role(
+ context
+ )
+
+ lines = [
+ "PROVIDER USAGE",
+ ]
+
+ for item in breakdown:
+ lines.append(
+ (
+ f"{item['role']}: "
+ f"{item['total_tokens']}"
+ f" (prompt={item['prompt_tokens']}, "
+ f"completion={item['completion_tokens']})"
+ )
+ )
+
+ lines.append(
+ f"total: {summary['total_tokens']}"
+ )
+
+ return "\n".join(
+ lines
+ )
diff --git a/utils/tokens.py b/utils/tokens.py
index e37494c5..38c1d44d 100644
--- a/utils/tokens.py
+++ b/utils/tokens.py
@@ -1,5 +1,6 @@
from settings.app_settings import settings
+
def estimate_tokens(
text: str,
) -> int:
@@ -24,55 +25,37 @@ def estimate_optional_tokens(
)
-def estimate_stream_tokens(
+def estimate_stream_input_tokens(
stream,
*,
prompt_text: str = "",
) -> int:
- visible_tokens = (
- estimate_optional_tokens(
- getattr(
- stream,
- "response",
- "",
- )
+ return estimate_optional_tokens(
+ prompt_text
+ )
+
+
+def estimate_stream_live_tokens(
+ stream,
+ *,
+ prompt_text: str = "",
+) -> int:
+
+ return (
+ estimate_stream_input_tokens(
+ stream,
+ prompt_text=prompt_text,
)
+ estimate_optional_tokens(
getattr(
stream,
- "reasoning",
+ "response",
"",
)
)
)
- prompt_tokens = (
- getattr(
- stream,
- "prompt_tokens",
- 0,
- )
- or estimate_optional_tokens(
- prompt_text
- )
- )
-
- provider_total = (
- getattr(
- stream,
- "total_tokens",
- 0,
- )
- or 0
- )
-
- return max(
- provider_total,
- prompt_tokens + visible_tokens,
- visible_tokens,
- )
-
def translation_token_limit(
text: str,
@@ -98,12 +81,16 @@ def estimate_runtime_tokens(
reasoning: str = "",
) -> int:
- total_text = (
- user_input
- + system_prompt
- + context_payload
- + response
- + reasoning
+ total_text = "\n".join(
+ value
+ for value in (
+ user_input,
+ system_prompt,
+ context_payload,
+ response,
+ reasoning,
+ )
+ if value
)
return estimate_tokens(
diff --git a/websocket.py b/websocket.py
index 7699f817..cb43932e 100644
--- a/websocket.py
+++ b/websocket.py
@@ -19,14 +19,39 @@
AgentRuntime,
)
+from clients.brain_client import (
+ build_brain_payload,
+ build_brain_system_prompt,
+)
+
from agents.agent_state import (
AgentState,
)
+from utils.brain import (
+ get_brain_runtime_config,
+)
+
+from utils.language import (
+ contains_cyrillic,
+)
+
+from utils.runtime_state_sync import (
+ refresh_runtime_state,
+)
+
+from utils.token_usage import (
+ format_token_usage_summary,
+)
+
from utils.telemetry import (
send_telemetry,
)
+from utils.tokens import (
+ estimate_runtime_tokens,
+)
+
from utils.ws_errors import (
handle_fatal_runtime_error,
handle_websocket_error,
@@ -103,6 +128,60 @@ async def receive_message(
# PROCESS MESSAGE
# ---------------------------------------------------------
+async def refresh_pending_brain_usage(
+ context,
+ user_text: str,
+):
+
+ if contains_cyrillic(
+ user_text
+ ):
+ return
+
+ brain_runtime = (
+ get_brain_runtime_config()
+ )
+
+ runtime_actions = (
+ brain_runtime.get(
+ "runtime_actions",
+ {},
+ )
+ )
+
+ system_prompt = (
+ build_brain_system_prompt(
+ context,
+ runtime_actions=runtime_actions,
+ )
+ )
+
+ brain_payload = (
+ build_brain_payload(
+ user_text,
+ context=context,
+ )
+ )
+
+ await refresh_runtime_state(
+ context,
+ runtime_id=(
+ brain_runtime["runtime_id"]
+ ),
+ used_tokens=(
+ estimate_runtime_tokens(
+ system_prompt=system_prompt,
+ user_input=brain_payload,
+ )
+ ),
+ max_tokens=(
+ brain_runtime["context_window"]
+ ),
+ last_error=None,
+ status="online",
+ )
+
+
async def process_message(
context,
message_data: dict,
@@ -123,6 +202,15 @@ async def process_message(
user_input=user_text
)
+ if hasattr(
+ context,
+ "runtime_usage_events",
+ ):
+ context.runtime_usage_events.clear()
+
+ else:
+ context.runtime_usage_events = []
+
runtime = AgentRuntime()
await logger.log_system(
@@ -142,6 +230,13 @@ async def process_message(
context,
)
+ await logger.log(
+ "[FLOW TELEMETRY]",
+ format_token_usage_summary(
+ context
+ ),
+ )
+
await logger.log_system(
"[WS] agent runtime end"
)
@@ -336,6 +431,11 @@ async def websocket_endpoint(
# START BACKGROUND TASK
# -------------------------------------------------
+ await refresh_pending_brain_usage(
+ context,
+ user_text,
+ )
+
current_task = (
asyncio.create_task(
process_message(