fix(claude_sdk,conversations,rotation): context size from the last api call, new-day text by rotation reason

This commit is contained in:
hh
2026-08-29 16:54:29 +02:00
parent 72639c4b90
commit 4a59c59bda
8 changed files with 120 additions and 16 deletions
+19 -3
View File
@@ -391,7 +391,7 @@ class ClaudeSdkBackend:
live.session_id = turn.result.session_id live.session_id = turn.result.session_id
if conversation_id is None: if conversation_id is None:
self._rekey(live.key, fingerprint([*history, *turn.synthesized])) 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 interrupted = live.interrupt_requested
live.interrupt_requested = False live.interrupt_requested = False
if capture is not None: if capture is not None:
@@ -455,6 +455,7 @@ class ClaudeSdkBackend:
event = message.event event = message.event
if event.get("type") == "message_start": if event.get("type") == "message_start":
offset = next_index offset = next_index
turn.context_tokens = _context_of(event.get("message"))
continue continue
index = event.get("index") index = event.get("index")
if isinstance(index, int): if isinstance(index, int):
@@ -740,6 +741,9 @@ class _SessionSpec:
class _Turn: class _Turn:
events: int = 0 events: int = 0
"""Wire events already yielded to the caller.""" """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) synthesized: list[dict[str, Any]] = field(default_factory=list)
result: ResultMessage | None = None result: ResultMessage | None = None
stop_reason: StopReason = "end_turn" stop_reason: StopReason = "end_turn"
@@ -895,11 +899,23 @@ def _block_to_dict(block: Any) -> dict[str, Any]:
raise TypeError(msg) 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: if result is None:
return TurnUsage() return TurnUsage(context_tokens=context_tokens)
usage = result.usage or {} usage = result.usage or {}
return TurnUsage( return TurnUsage(
context_tokens=context_tokens,
input_tokens=_int(usage.get("input_tokens")), input_tokens=_int(usage.get("input_tokens")),
output_tokens=_int(usage.get("output_tokens")), output_tokens=_int(usage.get("output_tokens")),
cache_read_tokens=_int(usage.get("cache_read_input_tokens")), cache_read_tokens=_int(usage.get("cache_read_input_tokens")),
+1
View File
@@ -428,6 +428,7 @@ async def _build_backends(
output_tokens=event.usage.output_tokens, output_tokens=event.usage.output_tokens,
cache_read_tokens=event.usage.cache_read_tokens, cache_read_tokens=event.usage.cache_read_tokens,
cache_creation_tokens=event.usage.cache_creation_tokens, cache_creation_tokens=event.usage.cache_creation_tokens,
context_tokens=event.usage.context_tokens,
cost_usd=event.usage.cost_usd, cost_usd=event.usage.cost_usd,
duration_ms=event.usage.duration_ms, duration_ms=event.usage.duration_ms,
num_turns=event.usage.num_turns, num_turns=event.usage.num_turns,
+45 -8
View File
@@ -23,7 +23,7 @@ import logging
import re import re
import uuid import uuid
from dataclasses import dataclass, field 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 typing import TYPE_CHECKING, Any, cast
from claude_agent_sdk import ( from claude_agent_sdk import (
@@ -134,6 +134,19 @@ class SeedContext:
title: str | None 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) @dataclass(frozen=True, slots=True)
class ConversationTexts: class ConversationTexts:
"""Texts the gateway cannot invent for a setup. """Texts the gateway cannot invent for a setup.
@@ -155,7 +168,10 @@ class ConversationTexts:
"Этот мастер закрывается ({reason}). Напиши хендаут за {day}: справку " "Этот мастер закрывается ({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 distill: Callable[[DistillContext], Awaitable[str] | str] | None = None
"""The distiller fork's first message (§8.4): which chat, what day, where """The distiller fork's first message (§8.4): which chat, what day, where
the digest goes; ``None`` uses a path-less default.""" the digest goes; ``None`` uses a path-less default."""
@@ -553,9 +569,7 @@ class Conversations:
.limit(1) .limit(1)
) )
).first() ).first()
if row is None: return context_of(row)
return 0
return row.input_tokens + row.cache_read_tokens + row.cache_creation_tokens
async def usage_tokens(self, since: datetime) -> int: async def usage_tokens(self, since: datetime) -> int:
async with self._db.session() as session: 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) _log.exception("%s could not mark %s", fe.name, conv.external_id)
return marked return marked
async def new_day(self, conv: Conversation, *, moved: int = 0) -> InjectQueueItem: async def new_day(
day = datetime.now(UTC).astimezone().date().isoformat() self, conv: Conversation, *, reason: str = "ночь", moved: int = 0
text = self._texts.new_day.format(day=day) ) -> 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: if moved:
text += f" Инжектов переехало из старого мастера: {moved}." text += f" Инжектов переехало из старого мастера: {moved}."
return await self.inject( return await self.inject(
@@ -1902,3 +1925,17 @@ def _usage_dict(capture: TurnCapture) -> dict[str, Any] | None:
"cost_usd": usage.cost_usd, "cost_usd": usage.cost_usd,
"duration_ms": usage.duration_ms, "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))
+1 -1
View File
@@ -115,7 +115,7 @@ class Rotation:
elif not branch.running_turn: elif not branch.running_turn:
await conversations.mark_closed(branch) await conversations.mark_closed(branch)
moved = await conversations.queue.move(old, new, priority="normal") 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( conversations.bus.publish(
"conversation.rotated", "conversation.rotated",
conversation_id=new.external_id, conversation_id=new.external_id,
+6
View File
@@ -27,6 +27,12 @@ class TurnUsage:
model_usage: dict[str, Any] | None = None model_usage: dict[str, Any] | None = None
"""``ResultMessage.model_usage`` verbatim: per-model tokens, cost, web searches.""" """``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 @dataclass
class TurnCapture: class TurnCapture:
+2
View File
@@ -299,6 +299,8 @@ class Usage(SQLModel, table=True):
output_tokens: int = 0 output_tokens: int = 0
cache_read_tokens: int = 0 cache_read_tokens: int = 0
cache_creation_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) cost_usd: float | None = Field(default=None)
duration_ms: int | None = Field(default=None) duration_ms: int | None = Field(default=None)
num_turns: int | None = Field(default=None) num_turns: int | None = Field(default=None)
+18 -2
View File
@@ -98,7 +98,12 @@ class FakeClient:
self.interrupted = True self.interrupted = True
async def receive_response(self): 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) yield StreamEvent(uuid="u", session_id="s", event=start)
for e in _stream(0, "calling "): for e in _stream(0, "calling "):
yield e yield e
@@ -119,7 +124,17 @@ class FakeClient:
yield UserMessage( yield UserMessage(
content=[ToolResultBlock(tool_use_id="toolu_1", content="42")] 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"): for e in _stream(0, "done"):
yield e yield e
yield AssistantMessage(content=[TextBlock(text="done")], model="m") 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.session_id == "fresh-session"
assert capture.usage is not None and capture.usage.cost_usd == 0.01 assert capture.usage is not None and capture.usage.cost_usd == 0.01
assert capture.usage.context_tokens == 23_007
assert capture.synthesized_messages == [ assert capture.synthesized_messages == [
{ {
"role": "assistant", "role": "assistant",
+28 -2
View File
@@ -100,7 +100,7 @@ async def test_rotation_order_handout_close_marks_moves_and_new_day(
return f"напиши хендаут за {ctx.day}" return f"напиши хендаут за {ctx.day}"
world.conversations._texts = ConversationTexts( # noqa: SLF001 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 world.conversations.spawn(kind="master", agent="a", seed="clean")
old = await age(world, old, datetime(2026, 8, 27, 9, 0, tzinfo=UTC)) 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] prompt = new_client.prompts[0]
assert prompt.startswith("[сид: morning] master") assert prompt.startswith("[сид: morning] master")
assert "[инжект: ротация" in prompt assert "[инжект: ротация" in prompt
assert "Новый день" in prompt assert "Новый день 20" in prompt and "(ночь)" in prompt
assert "переехало из старого мастера: 1" in prompt assert "переехало из старого мастера: 1" in prompt
moved = await world.conversations.queue.pending(new.id) moved = await world.conversations.queue.pending(new.id)
assert [(i.priority, i.text) for i in moved] == [("normal", "later")] 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)) rotation = Rotation(world.conversations, RotationPolicy(max_context_tokens=5))
(pair,) = await rotation.due() (pair,) = await rotation.due()
assert pair[0].id == conv.id and pair[1] == "транскрипт" 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