From 4a59c59bda88e1e993479284e11e3da3c2ced516 Mon Sep 17 00:00:00 2001 From: h Date: Sat, 29 Aug 2026 16:54:29 +0200 Subject: [PATCH] fix(claude_sdk,conversations,rotation): context size from the last api call, new-day text by rotation reason --- src/beaver_gateway/backends/claude_sdk.py | 22 ++++++++-- src/beaver_gateway/cli.py | 1 + src/beaver_gateway/core/conversations.py | 53 +++++++++++++++++++---- src/beaver_gateway/core/rotation.py | 2 +- src/beaver_gateway/core/turn_capture.py | 6 +++ src/beaver_gateway/storage/models.py | 2 + tests/test_claude_sdk_backend.py | 20 ++++++++- tests/test_rotation.py | 30 ++++++++++++- 8 files changed, 120 insertions(+), 16 deletions(-) diff --git a/src/beaver_gateway/backends/claude_sdk.py b/src/beaver_gateway/backends/claude_sdk.py index 3809301..ee89989 100644 --- a/src/beaver_gateway/backends/claude_sdk.py +++ b/src/beaver_gateway/backends/claude_sdk.py @@ -391,7 +391,7 @@ class ClaudeSdkBackend: live.session_id = turn.result.session_id if conversation_id is None: self._rekey(live.key, fingerprint([*history, *turn.synthesized])) - usage = _usage_of(turn.result) + usage = _usage_of(turn.result, context_tokens=turn.context_tokens) interrupted = live.interrupt_requested live.interrupt_requested = False if capture is not None: @@ -455,6 +455,7 @@ class ClaudeSdkBackend: event = message.event if event.get("type") == "message_start": offset = next_index + turn.context_tokens = _context_of(event.get("message")) continue index = event.get("index") if isinstance(index, int): @@ -740,6 +741,9 @@ class _SessionSpec: class _Turn: events: int = 0 """Wire events already yielded to the caller.""" + + context_tokens: int = 0 + """Input size of the latest API call (``message_start`` usage).""" synthesized: list[dict[str, Any]] = field(default_factory=list) result: ResultMessage | None = None stop_reason: StopReason = "end_turn" @@ -895,11 +899,23 @@ def _block_to_dict(block: Any) -> dict[str, Any]: raise TypeError(msg) -def _usage_of(result: ResultMessage | None) -> TurnUsage: +def _context_of(message: Any) -> int: + usage = message.get("usage") if isinstance(message, dict) else None + if not isinstance(usage, dict): + return 0 + return ( + _int(usage.get("input_tokens")) + + _int(usage.get("cache_read_input_tokens")) + + _int(usage.get("cache_creation_input_tokens")) + ) + + +def _usage_of(result: ResultMessage | None, *, context_tokens: int = 0) -> TurnUsage: if result is None: - return TurnUsage() + return TurnUsage(context_tokens=context_tokens) usage = result.usage or {} return TurnUsage( + context_tokens=context_tokens, input_tokens=_int(usage.get("input_tokens")), output_tokens=_int(usage.get("output_tokens")), cache_read_tokens=_int(usage.get("cache_read_input_tokens")), diff --git a/src/beaver_gateway/cli.py b/src/beaver_gateway/cli.py index c0e0b38..765be78 100644 --- a/src/beaver_gateway/cli.py +++ b/src/beaver_gateway/cli.py @@ -428,6 +428,7 @@ async def _build_backends( output_tokens=event.usage.output_tokens, cache_read_tokens=event.usage.cache_read_tokens, cache_creation_tokens=event.usage.cache_creation_tokens, + context_tokens=event.usage.context_tokens, cost_usd=event.usage.cost_usd, duration_ms=event.usage.duration_ms, num_turns=event.usage.num_turns, diff --git a/src/beaver_gateway/core/conversations.py b/src/beaver_gateway/core/conversations.py index e2c6d2a..31351eb 100644 --- a/src/beaver_gateway/core/conversations.py +++ b/src/beaver_gateway/core/conversations.py @@ -23,7 +23,7 @@ import logging import re import uuid from dataclasses import dataclass, field -from datetime import UTC, datetime, timedelta +from datetime import UTC, date, datetime, timedelta from typing import TYPE_CHECKING, Any, cast from claude_agent_sdk import ( @@ -134,6 +134,19 @@ class SeedContext: title: str | None +@dataclass(frozen=True, slots=True) +class NewDayContext: + """What the new master hears first. + + ``reason`` is ``ночь`` / ``возраст`` / ``транскрипт`` - only the first + one is actually a new day. + """ + + day: date + reason: str + moved: int + + @dataclass(frozen=True, slots=True) class ConversationTexts: """Texts the gateway cannot invent for a setup. @@ -155,7 +168,10 @@ class ConversationTexts: "Этот мастер закрывается ({reason}). Напиши хендаут за {day}: справку " "на утро, не задание - прошедшее время, без повелительного наклонения." ) - new_day: str = "Новый день: мастер сменился, хендаут за {day} записан." + new_day: Callable[[NewDayContext], Awaitable[str] | str] | str = ( + "Мастер сменился ({reason}), хендаут за {day} записан." + ) + """First inject of the new master; a callable sees the rotation reason.""" distill: Callable[[DistillContext], Awaitable[str] | str] | None = None """The distiller fork's first message (§8.4): which chat, what day, where the digest goes; ``None`` uses a path-less default.""" @@ -553,9 +569,7 @@ class Conversations: .limit(1) ) ).first() - if row is None: - return 0 - return row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens + return context_of(row) async def usage_tokens(self, since: datetime) -> int: async with self._db.session() as session: @@ -1213,9 +1227,18 @@ class Conversations: _log.exception("%s could not mark %s", fe.name, conv.external_id) return marked - async def new_day(self, conv: Conversation, *, moved: int = 0) -> InjectQueueItem: - day = datetime.now(UTC).astimezone().date().isoformat() - text = self._texts.new_day.format(day=day) + async def new_day( + self, conv: Conversation, *, reason: str = "ночь", moved: int = 0 + ) -> InjectQueueItem: + ctx = NewDayContext( + day=datetime.now(UTC).astimezone().date(), reason=reason, moved=moved + ) + source = self._texts.new_day + if isinstance(source, str): + text = source.format(day=ctx.day.isoformat(), reason=ctx.reason) + else: + produced: Any = source(ctx) + text = await produced if inspect.isawaitable(produced) else produced if moved: text += f" Инжектов переехало из старого мастера: {moved}." return await self.inject( @@ -1902,3 +1925,17 @@ def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None: "cost_usd": usage.cost_usd, "duration_ms": usage.duration_ms, } + + +def context_of(row: Usage | None) -> int: + """Context size of a turn. + + The last API call's input, or, for rows written before it was recorded, + the per-call average of the turn's input sums. + """ + if row is None: + return 0 + if row.context_tokens: + return row.context_tokens + total = row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens + return round(total / max(row.num_turns or 1, 1)) diff --git a/src/beaver_gateway/core/rotation.py b/src/beaver_gateway/core/rotation.py index 9d3abfe..0507039 100644 --- a/src/beaver_gateway/core/rotation.py +++ b/src/beaver_gateway/core/rotation.py @@ -115,7 +115,7 @@ class Rotation: elif not branch.running_turn: await conversations.mark_closed(branch) moved = await conversations.queue.move(old, new, priority="normal") - await conversations.new_day(new, moved=moved) + await conversations.new_day(new, reason=reason, moved=moved) conversations.bus.publish( "conversation.rotated", conversation_id=new.external_id, diff --git a/src/beaver_gateway/core/turn_capture.py b/src/beaver_gateway/core/turn_capture.py index cfeacf2..34bf1d8 100644 --- a/src/beaver_gateway/core/turn_capture.py +++ b/src/beaver_gateway/core/turn_capture.py @@ -27,6 +27,12 @@ class TurnUsage: model_usage: dict[str, Any] | None = None """``ResultMessage.model_usage`` verbatim: per-model tokens, cost, web searches.""" + context_tokens: int = 0 + """Input of the last API call in the turn (fresh + cached + written to + cache) - the context size the model actually ran with, what Claude Code + shows as the context. The token fields above are sums over every API + call of the turn and grow with the number of tool calls.""" + @dataclass class TurnCapture: diff --git a/src/beaver_gateway/storage/models.py b/src/beaver_gateway/storage/models.py index 9c646bc..1ec10a4 100644 --- a/src/beaver_gateway/storage/models.py +++ b/src/beaver_gateway/storage/models.py @@ -299,6 +299,8 @@ class Usage(SQLModel, table=True): output_tokens: int = 0 cache_read_tokens: int = 0 cache_creation_tokens: int = 0 + context_tokens: int | None = Field(default=None) + """Input of the turn's last API call - the real context size.""" cost_usd: float | None = Field(default=None) duration_ms: int | None = Field(default=None) num_turns: int | None = Field(default=None) diff --git a/tests/test_claude_sdk_backend.py b/tests/test_claude_sdk_backend.py index 3ac451c..57bca9a 100644 --- a/tests/test_claude_sdk_backend.py +++ b/tests/test_claude_sdk_backend.py @@ -98,7 +98,12 @@ class FakeClient: self.interrupted = True async def receive_response(self): - start = {"type": "message_start", "message": {}} + start = { + "type": "message_start", + "message": { + "usage": {"input_tokens": 5, "cache_read_input_tokens": 20_000} + }, + } yield StreamEvent(uuid="u", session_id="s", event=start) for e in _stream(0, "calling "): yield e @@ -119,7 +124,17 @@ class FakeClient: yield UserMessage( content=[ToolResultBlock(tool_use_id="toolu_1", content="42")] ) - yield StreamEvent(uuid="u", session_id="s", event=start) + second = { + "type": "message_start", + "message": { + "usage": { + "input_tokens": 7, + "cache_read_input_tokens": 20_000, + "cache_creation_input_tokens": 3_000, + } + }, + } + yield StreamEvent(uuid="u", session_id="s", event=second) for e in _stream(0, "done"): yield e yield AssistantMessage(content=[TextBlock(text="done")], model="m") @@ -194,6 +209,7 @@ async def test_stream_envelope_and_index_rebase(cwd: Path) -> None: assert capture.session_id == "fresh-session" assert capture.usage is not None and capture.usage.cost_usd == 0.01 + assert capture.usage.context_tokens == 23_007 assert capture.synthesized_messages == [ { "role": "assistant", diff --git a/tests/test_rotation.py b/tests/test_rotation.py index 9809fcd..6bd29f9 100644 --- a/tests/test_rotation.py +++ b/tests/test_rotation.py @@ -100,7 +100,7 @@ async def test_rotation_order_handout_close_marks_moves_and_new_day( return f"напиши хендаут за {ctx.day}" world.conversations._texts = ConversationTexts( # noqa: SLF001 - handout=handout, new_day="Новый день {day}." + handout=handout, new_day=lambda ctx: f"Новый день {ctx.day} ({ctx.reason})." ) old = await world.conversations.spawn(kind="master", agent="a", seed="clean") old = await age(world, old, datetime(2026, 8, 27, 9, 0, tzinfo=UTC)) @@ -140,7 +140,7 @@ async def test_rotation_order_handout_close_marks_moves_and_new_day( prompt = new_client.prompts[0] assert prompt.startswith("[сид: morning] master") assert "[инжект: ротация" in prompt - assert "Новый день" in prompt + assert "Новый день 20" in prompt and "(ночь)" in prompt assert "переехало из старого мастера: 1" in prompt moved = await world.conversations.queue.pending(new.id) assert [(i.priority, i.text) for i in moved] == [("normal", "later")] @@ -172,3 +172,29 @@ async def test_due_uses_last_usage_row_for_context_size(world: World) -> None: rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=5)) (pair,) = await rotation.due() assert pair[0].id == conv.id and pair[1] == "транскрипт" + + +def test_context_of_prefers_last_call_and_averages_old_rows() -> None: + from beaver_gateway.core.conversations import context_of + from beaver_gateway.storage.models import Usage + + fresh = Usage( + agent_name="a", + model="m", + input_tokens=12, + cache_read_tokens=142_984, + cache_creation_tokens=29_584, + num_turns=6, + context_tokens=29_000, + ) + assert context_of(fresh) == 29_000 + old = Usage( + agent_name="a", + model="m", + input_tokens=12, + cache_read_tokens=142_984, + cache_creation_tokens=29_584, + num_turns=6, + ) + assert context_of(old) == 28_763 + assert context_of(None) == 0