[aidan] ui/ux: hide noisy memory/mode from chat title, fix token bug

Drop the inline 'Memory N% full' pill from AgentChat and the mode label from AgentCard so the chat header stays quiet until tools are connected. Backend now emits agent:context_update after compact and clear so the session token counters reflect the trimmed history instead of stale pre-compact totals, with a new _estimate_post_compact_input helper and matching invariants in test_v2_invariants.
This commit is contained in:
Aidan
2026-06-13 00:40:19 -07:00
committed by abccodes
parent 08526b6149
commit 7b88be7003
6 changed files with 148 additions and 34 deletions
+37
View File
@@ -53,6 +53,7 @@ from backend.apps.agents.manager.prompt.tool_catalog import (
from backend.apps.agents.core.aux_llm import _safe_resp_text, clean_short_label, aux_max_tokens_for
from backend.apps.agents.manager.session.history_compaction import (
_build_history_prefix,
_estimate_post_compact_input,
_get_branch_messages,
_truncate_large_tool_result,
)
@@ -377,6 +378,35 @@ class AgentManager:
session.compacted_through_msg_id = last_id
return True
async def _emit_context_update(
self,
session_id: str,
session: AgentSession,
*,
input_tokens: int | None = None,
output_tokens: int | None = None,
cache_read_tokens: int = 0,
cache_read_pct: float = 0.0,
) -> None:
if input_tokens is None:
input_tokens = int(session.tokens.get("input", 0) or 0)
if output_tokens is None:
output_tokens = int(session.tokens.get("output", 0) or 0)
session.tokens["input"] = input_tokens
session.tokens["output"] = output_tokens
ctx_window = max(1, getattr(session, "context_window", 0) or 200_000)
await ws_manager.send_to_session(session_id, "agent:context_update", {
"session_id": session_id,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"cache_read_tokens": cache_read_tokens,
"cache_read_pct": cache_read_pct,
"ctx_used_pct": round(input_tokens / ctx_window, 4) if input_tokens else 0.0,
"context_window": ctx_window,
"framework_overhead_tokens": session.framework_overhead_tokens,
"active_mcps": list(session.active_mcps),
})
def _build_prompt_content(self, prompt: str, images: list | None = None, context_paths: list | None = None, forced_tools: list[str] | None = None, attached_skills: list | None = None, api_type: str = "anthropic", model: str = ""):
return _build_prompt_content(prompt, images, context_paths, forced_tools, attached_skills, api_type, model)
@@ -1813,11 +1843,18 @@ class AgentManager:
# zero latency on the user's turn.
try:
if self._maybe_compact(session):
new_input = _estimate_post_compact_input(session)
await ws_manager.send_to_session(session_id, "agent:context_status", {
"session_id": session_id,
"reason": "compacted",
"compacted_through_msg_id": session.compacted_through_msg_id,
})
await self._emit_context_update(
session_id,
session,
input_tokens=new_input,
output_tokens=session.tokens.get("output", 0),
)
except Exception:
logger.exception("compaction failed; proceeding without it")
+13 -1
View File
@@ -2,6 +2,7 @@ from backend.config.Apps import SubApp
from backend.apps.agents.agent_manager import agent_manager
from backend.apps.agents.core.ws_manager import ws_manager
from backend.apps.agents.core.models import AgentConfig, ApprovalResponse
from backend.apps.agents.manager.session.history_compaction import _estimate_post_compact_input
from contextlib import asynccontextmanager
from fastapi import WebSocket, WebSocketDisconnect, HTTPException
from fastapi.responses import JSONResponse
@@ -304,6 +305,12 @@ async def compact_session(session_id: str):
"reason": "compacted",
"compacted_through_msg_id": session.compacted_through_msg_id,
})
await agent_manager._emit_context_update(
session_id,
session,
input_tokens=_estimate_post_compact_input(session),
output_tokens=session.tokens.get("output", 0),
)
except Exception:
pass
return {"ok": True, "compacted": fired}
@@ -329,6 +336,12 @@ async def clear_session(session_id: str):
"status": session.status,
"session": session.model_dump(mode="json"),
})
await agent_manager._emit_context_update(
session_id,
session,
input_tokens=0,
output_tokens=0,
)
except Exception:
pass
return {"ok": True}
@@ -817,4 +830,3 @@ async def subscriptions_disconnect(body: dict):
return {"ok": False, "error": "Connection not found"}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@@ -71,6 +71,39 @@ def _build_history_prefix(messages, cutoff_msg_id: str | None = None) -> str:
return "<prior_conversation>\n" + "\n".join(lines) + "\n</prior_conversation>"
def _estimate_post_compact_input(session) -> int:
"""Return a conservative token estimate after compaction trims history."""
try:
messages = _get_branch_messages(session)
cutoff_msg_id = getattr(session, "compacted_through_msg_id", None)
if cutoff_msg_id:
skip_idx = next(
(i for i, m in enumerate(messages) if m.id == cutoff_msg_id),
-1,
)
if skip_idx >= 0:
messages = messages[skip_idx + 1:]
surviving_chars = 0
for message in messages:
if getattr(message, "hidden", False):
continue
content = getattr(message, "content", "")
if isinstance(content, str):
serialized = content
else:
try:
serialized = json.dumps(content, ensure_ascii=False)
except Exception:
serialized = str(content)
surviving_chars += len(serialized)
framework_overhead = int(getattr(session, "framework_overhead_tokens", 0) or 0)
summary_overhead = 200 if cutoff_msg_id else 0
return max(0, framework_overhead + summary_overhead + (surviving_chars // 4))
except Exception:
logger.debug("post-compact token estimate failed", exc_info=True)
return max(0, int(getattr(session, "framework_overhead_tokens", 0) or 0))
def _truncate_large_tool_result(content: object, session_id: str, msg_id: str, max_bytes: int = 50_000) -> tuple[object, str | None]:
"""Spill a large tool_result body to disk, return a truncated
inline replacement plus the on-disk path (or None if untouched).
+60
View File
@@ -784,6 +784,66 @@ def test_compact_threshold_default():
assert s.context_window == 200_000
def test_post_compact_estimate_excludes_compacted_messages():
from backend.apps.agents.core.models import AgentSession, Message
from backend.apps.agents.manager.session.history_compaction import (
_estimate_post_compact_input,
)
messages = [
Message(id=f"m{i}", role="user", content=("old" * 1000 if i < 6 else "keep"))
for i in range(8)
]
s = AgentSession(id="x", name="t", model="sonnet", mode="agent")
s.messages = messages
s.compacted_through_msg_id = "m5"
s.framework_overhead_tokens = 100
assert _estimate_post_compact_input(s) == 100 + 200 + (len("keepkeep") // 4)
@pytest.mark.asyncio
async def test_context_update_emitter_refreshes_session_tokens(monkeypatch):
import backend.apps.agents.agent_manager as agent_manager_module
from backend.apps.agents.agent_manager import AgentManager
from backend.apps.agents.core.models import AgentSession
sent = []
async def fake_send_to_session(session_id, event, payload):
sent.append((session_id, event, payload))
monkeypatch.setattr(
agent_manager_module.ws_manager,
"send_to_session",
fake_send_to_session,
)
s = AgentSession(id="x", name="t", model="sonnet", mode="agent")
s.context_window = 1_000
s.tokens = {"input": 900, "output": 7}
s.framework_overhead_tokens = 42
s.active_mcps = ["github"]
await AgentManager()._emit_context_update("x", s, input_tokens=250)
assert s.tokens == {"input": 250, "output": 7}
assert sent == [(
"x",
"agent:context_update",
{
"session_id": "x",
"input_tokens": 250,
"output_tokens": 7,
"cache_read_tokens": 0,
"cache_read_pct": 0.0,
"ctx_used_pct": 0.25,
"context_window": 1_000,
"framework_overhead_tokens": 42,
"active_mcps": ["github"],
},
)]
# ===========================================================================
# Group L, Sentence-case display (the parseMcpToolName fix)
# ===========================================================================
+5 -30
View File
@@ -1041,42 +1041,17 @@ const AgentChat: React.FC<AgentChatProps> = ({ sessionId: sessionIdProp, onClose
);
})()}
{(() => {
const liveWindow = session.context_window || contextEstimate.limit || 200_000;
const liveInput = session.tokens?.input ?? 0;
const pct = liveInput > 0
? Math.min(1, liveInput / Math.max(1, liveWindow))
: (contextEstimate.used / Math.max(1, liveWindow));
const mcpCount = session.active_mcps?.length ?? 0;
// Quiet until it matters: hide the memory meter until the chat is filling up,
// and hide the tool count unless tools are actually connected.
const showMemory = pct >= 0.60;
const showTools = mcpCount > 0;
if (!showMemory && !showTools) return null;
const pctTxt = `${Math.round(pct * 100)}%`;
const memColor = pct >= 0.85 ? '#ef4444' : '#f59e0b';
const tip = [
showMemory ? `Memory ${pctTxt} full. As the chat fills up, the oldest messages start dropping out.` : null,
showTools ? `${mcpCount} tool${mcpCount === 1 ? '' : 's'} connected.` : null,
].filter(Boolean).join('\n');
if (mcpCount === 0) return null;
return (
<Typography
variant="caption"
sx={{ display: 'inline-flex', alignItems: 'center', gap: 0.5, fontVariantNumeric: 'tabular-nums' }}
title={tip}
title={`${mcpCount} tool${mcpCount === 1 ? '' : 's'} connected.`}
>
{showMemory && (
<Box component="span" sx={{ color: memColor, fontWeight: 500 }}>
Memory {pctTxt} full
</Box>
)}
{showMemory && showTools && (
<Box component="span" sx={{ color: c.text.ghost }}>·</Box>
)}
{showTools && (
<Box component="span" sx={{ color: c.text.tertiary }}>
{mcpCount} tool{mcpCount === 1 ? '' : 's'}
</Box>
)}
<Box component="span" sx={{ color: c.text.tertiary }}>
{mcpCount} tool{mcpCount === 1 ? '' : 's'}
</Box>
</Typography>
);
})()}
@@ -818,9 +818,6 @@ const AgentCard: React.FC<Props> = ({
<Typography variant="caption" sx={{ color: c.text.tertiary }}>
{friendlyModelLabel}
</Typography>
<Typography variant="caption" sx={{ color: c.text.tertiary }}>
{session.mode}
</Typography>
<Typography variant="caption" sx={{ color: c.text.tertiary }}>
<ElapsedTimer messages={session.messages} status={session.status} />
</Typography>