From 10eccf288ff8184ce4035f49795bc2462a080f74 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 19:51:26 -0700 Subject: [PATCH] refactor(gateway): share transcript-estimate helper across /context and /usage; checkpoint kwargs by name --- gateway/slash_commands.py | 9 ++---- gateway/slash_commands_status.py | 52 ++++++++++++++------------------ 2 files changed, 25 insertions(+), 36 deletions(-) diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index bff20e34f8..d7ea6be803 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -259,12 +259,9 @@ class GatewaySlashCommandsMixin( cp = _checkpoint_agent_kwargs(_load_gateway_config()) if not cp["checkpoints_enabled"]: return None - return CheckpointManager( - enabled=True, - max_snapshots=cp["checkpoint_max_snapshots"], - max_total_size_mb=cp["checkpoint_max_total_size_mb"], - max_file_size_mb=cp["checkpoint_max_file_size_mb"], - ) + # AIAgent kwargs are ``checkpoint_``; CheckpointManager takes the bare field names. + return CheckpointManager(enabled=True, **{k[len("checkpoint_"):]: v for k, v in cp.items() + if k.startswith("checkpoint_")}) def _write_approval_setter(self, section: str, event: MessageEvent): """``set_mode_fn`` for /memory and /skills: persist ``
.write_approval``. Raw read is diff --git a/gateway/slash_commands_status.py b/gateway/slash_commands_status.py index 5a8bcd7025..b110e2beab 100644 --- a/gateway/slash_commands_status.py +++ b/gateway/slash_commands_status.py @@ -51,8 +51,11 @@ def _clip(text: str, limit: int) -> str: return text[: limit - 3] + "..." if len(text) > limit else text -def _chat_msgs(history) -> list[dict]: - return [m for m in history if m.get("role") in {"user", "assistant"} and m.get("content")] +def _transcript_estimate(history) -> tuple[int, int]: + """``(approx_tokens, message_count)`` over the user/assistant messages of a transcript.""" + from agent.model_metadata import estimate_messages_tokens_rough + msgs = [m for m in history if m.get("role") in {"user", "assistant"} and m.get("content")] + return estimate_messages_tokens_rough(msgs), len(msgs) async def _quiet(call, default=None): @@ -317,14 +320,12 @@ class GatewayStatusCommandsMixin: if used > 0 and context_length > 0: pct = _pct(used, context_length) filled = int(round(pct / 100 * 24)) - bar = "█" * max(0, filled) + "░" * max(0, 24 - filled) lines = [ - t("gateway.context.header"), - "", + t("gateway.context.header"), "", t("gateway.context.model", model=model_name or "?"), t("gateway.context.window", total=_fmt(context_length)), t("gateway.context.in_use", used=_fmt(used), total=_fmt(context_length), pct=f"{pct:.0f}"), - t("gateway.context.bar", bar=bar), + t("gateway.context.bar", bar="█" * max(0, filled) + "░" * max(0, 24 - filled)), t("gateway.context.headroom", headroom=_fmt(max(0, context_length - used))), "", ] @@ -344,26 +345,21 @@ class GatewayStatusCommandsMixin: # Last resort: rough estimate from transcript history = await self.async_session_store.load_transcript(session_entry.session_id) - if history: - from agent.model_metadata import estimate_messages_tokens_rough - msgs = _chat_msgs(history) - return "\n".join([ - t("gateway.context.header"), - "", - t("gateway.context.estimated", count=_fmt(estimate_messages_tokens_rough(msgs)), - messages=len(msgs)), - t("gateway.context.detail_after_first"), - ]) - return t("gateway.context.no_data") + if not history: + return t("gateway.context.no_data") + approx, count = _transcript_estimate(history) + return "\n".join([ + t("gateway.context.header"), "", + t("gateway.context.estimated", count=_fmt(approx), messages=count), + t("gateway.context.detail_after_first"), + ]) async def _resolve_context_figures(self, agent, ctx, session_entry, source): """``(used, context_length, model_name)`` for /context: used = compressor -> SessionStore; model = agent -> SessionDB row; window = compressor -> gateway model route -> model metadata.""" - used = _n(ctx, "last_prompt_tokens") if ctx is not None else 0 - context_length = _n(ctx, "context_length") if ctx is not None else 0 + used = _n(ctx, "last_prompt_tokens") or _int_value(getattr(session_entry, "last_prompt_tokens", 0)) + context_length = _n(ctx, "context_length") model_name = _clean_str(getattr(agent, "model", "")) if agent is not None else "" - if not used: - used = _int_value(getattr(session_entry, "last_prompt_tokens", 0)) if not model_name and self._session_db: row = await _quiet(lambda: self._session_db.get_session(session_entry.session_id)) model_name = _clean_str(row.get("model", "")) if isinstance(row, dict) else "" @@ -460,11 +456,8 @@ class GatewayStatusCommandsMixin: if view.identity_line: lines += ["", view.identity_line] if view.topup_url: - lines += [ - "", - f"Manage billing on the portal: {view.topup_url}", - "Top up and manage billing in the browser — your balance updates here after.", - ] + lines += ["", f"Manage billing on the portal: {view.topup_url}", + "Top up and manage billing in the browser — your balance updates here after."] return "\n".join(lines) def _context_breakdown_block(self, agent, source, expanded: bool) -> list[str]: @@ -583,12 +576,11 @@ class GatewayStatusCommandsMixin: session_entry = await self.async_session_store.get_or_create_session(source) history = await self.async_session_store.load_transcript(session_entry.session_id) if history: - from agent.model_metadata import estimate_messages_tokens_rough - msgs = _chat_msgs(history) + approx, count = _transcript_estimate(history) return _with_account_blocks([ t("gateway.usage.header_session_info"), - t("gateway.usage.label_messages", count=len(msgs)), - t("gateway.usage.label_estimated_context", count=_fmt(estimate_messages_tokens_rough(msgs))), + t("gateway.usage.label_messages", count=count), + t("gateway.usage.label_estimated_context", count=_fmt(approx)), t("gateway.usage.detailed_after_first"), ]) if account_lines or credits_lines: