diff --git a/raven/agent/loop/_shared.py b/raven/agent/loop/_shared.py index 2b88fb942..58c8a036f 100644 --- a/raven/agent/loop/_shared.py +++ b/raven/agent/loop/_shared.py @@ -164,7 +164,11 @@ def _stamp_reasoning_ms(messages: list[dict[str, Any]], response: Any) -> None: # How full a model call left the window, on the assistant entry that stands for -# it: the reply it produced, or the text a hook put in that reply's place. +# it: the reply it produced, or the text a hook put in that reply's place. A +# reply that is not filed leaves it on the newest row the turn keeps instead +# (``turn_path._stamp_context_used_on_kept``), and a turn that fails or is +# stopped files its newest reading on the marker it ends on +# (``_save_broken_turn``). # Same underscore-then-rename convention as the spans above; ``_save_turn`` # writes it as ``context_used``. A resume has no other way to know what the # model was sent: the context engine compacts on its own and leaves the stored @@ -178,8 +182,9 @@ def _window_reading(usage: dict[str, Any] | None) -> int: Not what the prompt was billed as. A provider that reports cache reads apart from ``prompt_tokens`` (``prompt_tokens_include_cache`` false, the Anthropic Messages shape) bills a fraction of a warm prompt while all of it still - occupies the window, so the cached counts are added back. Same add-back, for - the same reason, as ``transport_failure.never_processed``. + occupies the window, so the cached counts are added back -- and only then: + where ``prompt_tokens`` already holds them, adding them again would count the + warm prompt twice. """ if not usage: return 0 @@ -191,13 +196,31 @@ def _window_reading(usage: dict[str, Any] | None) -> int: return used +def _reads_no_window(response: Any) -> bool: + """Whether a call's usage says nothing about how full the window is. + + A response the provider flagged as an upstream transport failure + (``upstream_transport_failure``) carries the upstream's account of a call + that failed, and that account can be the very evidence that flagged it -- a + prompt billed at a size the request cannot have been. The call may never + have been read, so its usage is no reading of the window. + """ + classification = getattr(response, "error_classification", None) + return getattr(classification, "category", None) == "upstream_transport_failure" + + +def _call_window(response: Any) -> int: + """How full a model call left the window, or 0 when it gives no reading.""" + return 0 if _reads_no_window(response) else _window_reading(getattr(response, "usage", None)) + + def _stamp_context_used(messages: list[dict[str, Any]], response: Any) -> None: """Put a model call's window reading on the assistant entry that stands for it. Absent when the provider reported no usage, and absent means unknown: a resume reads the newest entry that has one. """ - used = _window_reading(getattr(response, "usage", None)) + used = _call_window(response) if used > 0 and messages: messages[-1][_CONTEXT_USED_KEY] = used diff --git a/raven/agent/loop/turn_path.py b/raven/agent/loop/turn_path.py index a9172ab83..3f1b8e5bc 100644 --- a/raven/agent/loop/turn_path.py +++ b/raven/agent/loop/turn_path.py @@ -49,6 +49,7 @@ Session, TurnSynthesisPolicy, _appended_by_hook, + _call_window, _display_label, _file_change_payload, _file_composed_reply, @@ -60,7 +61,6 @@ _stamp_reasoning_ms, _stamp_replied_via_tool, _strip_inline_images, - _window_reading, append_hook_note, asyncio, autofill_resolver, @@ -358,6 +358,28 @@ def _stamp_turn_observers(messages: list[dict[str, Any]], metadata: dict[str, An _TURN_TRANSIENT_KEYS = ("_recovery_synthetic", ATTACHED_IMAGE_KEY) +def _stamp_context_used_on_kept(messages: list[dict[str, Any]], response: Any) -> None: + """Put the reading of a call whose reply is not filed on the newest row that is. + + An empty reply never becomes a row of its own in the session: recovery + stands scaffolding in for it, the loop gives up on it, or it is filed empty + and dropped on save. Nor does a provider's error, which is never filed. The + live gauge still showed that call's window, and the window it measured ends + at the newest row the turn keeps, so the reading goes there. When that row + is the turn's question it is already on disk and the stamp never reaches + it: the turn either files a row of its own later, which takes a newer + reading, or ends on the marker ``_save_broken_turn`` files, which carries + it. + """ + used = _call_window(response) + if used <= 0: + return + for message in reversed(messages): + if not any(message.get(key) for key in _TURN_TRANSIENT_KEYS): + message[_CONTEXT_USED_KEY] = used + return + + # How a turn reports the model saying nothing with no recovery behind it. The # texts recovery fails a turn with say what it tried; nothing was tried here, # and this says why. @@ -910,6 +932,7 @@ async def _run_agent_loop( # noqa: C901 (cc 100: pre-existing, above the ceilin turn_log: TurnLog | None = None, rungs_used: set[str] | None = None, reply_awaited: bool = False, + window_reading: dict[str, int] | None = None, ) -> tuple[str | None, list[str], list[dict], LoopOutcome]: """Run the agent iteration loop. @@ -969,6 +992,11 @@ async def _run_agent_loop( # noqa: C901 (cc 100: pre-existing, above the ceilin thinking-only reply down to, each asked for once. A caller that may run the loop again for one turn hands every run the same set; omitted, this run keeps a set of its own. + + ``window_reading``, when given, is kept at ``context_used`` with the + newest window a call of this loop measured. It outlives the loop the way + ``spend`` does: a turn that fails or is stopped files it on its marker, + which can be the only row the turn adds. """ spend = spend if spend is not None else TurnSpend(session_key) messages = initial_messages @@ -1599,11 +1627,23 @@ async def _end_on_the_partial(last: LLMResponse) -> str: # keeps it for a resume (``_stamp_context_used``). They differ # in their thresholds, not in what they measure -- and what # they measure is how full the window is, which is not what - # the prompt was billed as (see ``_window_reading``). - context_used = _window_reading(response.usage) + # the prompt was billed as (see ``_window_reading``), and a + # call can give none (see ``_call_window``). + context_used = _call_window(response) if context_used > 0: window.last_context_used = context_used - if usage_sink is not None and response.usage: + if window_reading is not None: + window_reading["context_used"] = context_used + if usage_sink is not None and response.usage and context_used <= 0 and usage_sink.get("context_used"): + # No reading -- a block of zeros, or an account ``_call_window`` + # does not take as one -- so the gauge stays on the call before + # it, as the reading the turn files does (``window_reading`` + # above moves on the same terms). The cost is the turn's, a + # flow, and moves on regardless. With no reading to keep, the + # branch below fills the sink: an answered turn sends it whole. + usage_sink["cost_usd"] = spend.cost_usd + usage_sink["cost_missing_calls"] = spend.cost_missing_calls + elif usage_sink is not None and response.usage: # An explicitly configured window always wins over the live # table -- that is what setting it means. Otherwise the live # window from the model's provider table (e.g. OpenRouter, @@ -2193,6 +2233,7 @@ async def _end_on_the_partial(last: LLMResponse) -> str: status = "error" error_detail = _llm_failure_detail(clean, verdict) error_full = response_failure_text(clean, response.call_record) + _stamp_context_used_on_kept(messages, response) break if recovery.parts: @@ -2208,6 +2249,9 @@ async def _end_on_the_partial(last: LLMResponse) -> str: recovery.parts.clear() carried_gate = None + if not (clean or "").strip(): + _stamp_context_used_on_kept(messages, response) + # Empty-response recovery. A reply with no tool call is read for # what it is (``reply_kind``) and ``decide`` names the move; an # empty assistant turn would otherwise end here as a turn with no @@ -3152,6 +3196,12 @@ async def _tap_episode(_index: int) -> None: # turn: a set per attempt would pay for ``high`` and ``low`` again. rungs_used: set[str] = set() + # The newest window a call of this turn measured, held across attempts. + # A turn that fails or is stopped after calls that each came back empty + # or as an error has filed no row of its own, and the question it would + # otherwise end on is already on disk, so its marker carries the reading. + reading: dict[str, int] = {} + async def _attempt(seed: list[dict], attempt: int): pending.update(rerun=False, reasons=[]) # The record this attempt persists from, for the rescue paths below as @@ -3185,6 +3235,7 @@ async def _attempt(seed: list[dict], attempt: int): turn_log=live["log"], rungs_used=rungs_used, reply_awaited=reply_awaited, + window_reading=reading, ) # Taken BEFORE the first attempt, because the loop appends to the list it is @@ -3346,6 +3397,7 @@ def _rerun_pending(final_content: str | None, msgs: list[dict], status: str) -> None, streamed, status="cancelled", + context_used=reading.get("context_used"), ) raise except AnswerlessTurnError as exc: @@ -3361,6 +3413,7 @@ def _rerun_pending(final_content: str | None, msgs: list[dict], status: str) -> status="failed", reason=str(exc), full_detail=attach_full_detail(exc), + context_used=reading.get("context_used"), ) raise except Exception as exc: @@ -3376,6 +3429,7 @@ def _rerun_pending(final_content: str | None, msgs: list[dict], status: str) -> # and a model both read back. reason=bound_failure_text(str(exc)), full_detail=attach_full_detail(exc), + context_used=reading.get("context_used"), ) raise self._stash_recovery(key, outcome) @@ -3519,6 +3573,7 @@ def _save_broken_turn( reason: str | None = None, full_detail: str | None = None, inbound_original: str | None = None, + context_used: int | None = None, ) -> None: """Persist what a cancelled or failed turn got as far as producing. @@ -3539,6 +3594,9 @@ def _save_broken_turn( ``_marker_failure``). ``turn_ended.full_detail`` is the long account beside the bounded ``reason``, filed only when there is more to it. + ``context_used`` is the newest window the turn measured, filed on the + marker because it is the newest row: a resume reads the window from + there, and a turn whose calls left no row of their own has no other. ``received_at`` and ``inbound_original`` describe the turn's question, which ``_process_message`` files before the attempt starts and hands @@ -3603,6 +3661,8 @@ def _save_broken_turn( **({"full_detail": full_detail} if full_detail else {}), }, } + if context_used: + marker[_CONTEXT_USED_KEY] = context_used tail.append(marker) self._save_turn(session, tail, 0, received_at=received_at, inbound_original=inbound_original) self.sessions.save(session) diff --git a/raven/cli/tui_commands.py b/raven/cli/tui_commands.py index a28d61fc5..4d08b06bb 100644 --- a/raven/cli/tui_commands.py +++ b/raven/cli/tui_commands.py @@ -23,7 +23,7 @@ import sys import threading from pathlib import Path -from typing import Callable, Optional, Tuple +from typing import Any, Callable, Optional, Tuple import typer @@ -447,6 +447,9 @@ def _agent_loop_factory(): # turn.send binds a turn's addressee into it and the spine's outlet/sink read # it back to tag that turn's events (see build_rpc_spine). direct_targets: dict[str, dict[str, str]] = {} + # Owned here for the same reason: the spine keeps what each running turn has + # measured in it, and turn.cancel reads a stopped turn's reading back. + live_usages: dict[str, dict[str, Any]] = {} turn_teardown = None if agent_loop is not None: from raven.core.cron_stack import make_on_cron_job @@ -462,6 +465,7 @@ def _agent_loop_factory(): direct_targets=direct_targets, readback_texts=cron_readback, approval_responder=approval_broker, + live_usages=live_usages, ) # Subagent result re-injection submits a SUBAGENT-origin turn. agent_loop.subagents.set_submit(turn_scheduler.submit) @@ -517,6 +521,7 @@ def _agent_loop_factory(): turn_ids=turn_ids, direct_targets=direct_targets, build_error=build_error, + live_usages=live_usages, ) serve_task = asyncio.create_task(server.serve_forever()) diff --git a/raven/rpc/bootstrap.py b/raven/rpc/bootstrap.py index bea39dfb9..b989a9e31 100644 --- a/raven/rpc/bootstrap.py +++ b/raven/rpc/bootstrap.py @@ -437,6 +437,9 @@ def _agent_loop_factory(): # turn.send binds a turn's addressee into it and the spine's outlet/sink read # it back to tag that turn's events (see build_rpc_spine). direct_targets: dict[str, dict[str, str]] = {} + # Owned here for the same reason: the spine keeps what each running turn has + # measured in it, and turn.cancel reads a stopped turn's reading back. + live_usages: dict[str, dict[str, Any]] = {} turn_teardown = None if agent_loop is not None: from raven.core.cron_stack import make_on_cron_job, make_on_session_wake @@ -459,6 +462,7 @@ def _agent_loop_factory(): # is not necessarily removing that method. approval_responder=approval_responder or approval_broker, credential_asker=credential_broker if credential_cards else None, + live_usages=live_usages, ) if owns_loop: agent_loop.subagents.set_submit(turn_scheduler.submit) @@ -507,6 +511,7 @@ def _agent_loop_factory(): send_frame=send_frame, default_channel=channel, ensure_stack=ensure_stack, + live_usages=live_usages, ) if agent_loop is not None: _lend_settings_writers(agent_loop, dispatcher) diff --git a/raven/rpc/methods/__init__.py b/raven/rpc/methods/__init__.py index 3ccc1b2d2..e12f4cb5a 100644 --- a/raven/rpc/methods/__init__.py +++ b/raven/rpc/methods/__init__.py @@ -84,6 +84,7 @@ def register_aligned_methods( build_error: "RpcError | None" = None, send_frame: "Any" = None, default_channel: str = "tui", + live_usages: "dict[str, dict[str, Any]] | None" = None, ) -> None: """Register every aligned RPC handler on a dispatcher. @@ -116,6 +117,7 @@ def register_aligned_methods( build_error=build_error, send_frame=send_frame, default_channel=default_channel, + live_usages=live_usages, ) @@ -135,6 +137,7 @@ def register_aligned_methods_except_system( send_frame: "Any" = None, default_channel: str = "tui", ensure_stack: "Callable[[], Awaitable[bool]] | None" = None, + live_usages: "dict[str, dict[str, Any]] | None" = None, ) -> None: """Register every aligned RPC handler EXCEPT system.* on a dispatcher. @@ -223,6 +226,7 @@ def register_aligned_methods_except_system( # Must equal the channel ``build_rpc_spine`` registered its outlet # under; ``turn.py`` documents that a mismatch drops the reply. default_channel=default_channel, + live_usages=live_usages, ) # confirm.respond — needs a ConfirmBroker to resolve the pending # confirm future. Gated like turn.*: when no broker is supplied (demo diff --git a/raven/rpc/methods/session.py b/raven/rpc/methods/session.py index b62486a2b..40ec25f8a 100644 --- a/raven/rpc/methods/session.py +++ b/raven/rpc/methods/session.py @@ -684,7 +684,7 @@ async def session_resume( def _fill_resumed_context(usage: dict[str, Any], session: Any) -> None: """How full the context window is, into a caller's usage mapping. - Read off the newest stored reply that carries the window its call measured + Read off the newest stored row that carries the window a call measured (``context_used``): that is the number the meter showed when the turn ended. The stored log cannot stand in for it. The context engine compacts on its own and leaves the log whole, so ``last_consolidated`` never moves and the @@ -692,8 +692,8 @@ def _fill_resumed_context(usage: dict[str, Any], session: Any) -> None: times the window its last call measured, and a short one at a few tokens, the system prompt and the tool list being no part of it. - A session with no reading -- written before replies carried one, or never - answered -- falls back to a tiktoken estimate of + A session with no reading -- written before rows carried one, or whose + calls never reported their usage -- falls back to a tiktoken estimate of ``get_history(max_messages=0)``, marked ``context_estimated``. Failures leave the zero baseline alone rather than guessing -- a wrong denominator is worse than an unknown one. @@ -707,12 +707,13 @@ def _fill_resumed_context(usage: dict[str, Any], session: Any) -> None: usage["context_estimated"] = False return try: - # max_messages=0 is "all of it", which is what both places that build or - # measure the real prompt pass (loop/main.py and the consolidator's own - # estimator). The default caps at the last 500, and the tail is bounded - # in tokens rather than in messages -- consolidation fires on the window - # -- so a chatty session on a large-window model sits well past 500 and - # the meter would go quiet exactly as it started to matter. + # The estimate counts what an engine that leaves compaction to the + # consolidator would be handed: max_messages=0 is "all of it", the slice + # ``DefaultMemory.candidate_messages`` and the consolidator's own + # estimator both take. The default caps at the last 500, and that tail is + # bounded in tokens rather than in messages -- consolidation fires on the + # window -- so a chatty session on a large-window model sits well past + # 500 and the meter would go quiet exactly as it started to matter. messages = session.get_history(max_messages=0) except Exception: logger.exception("session.*: could not read the session history") @@ -729,11 +730,17 @@ def _fill_resumed_context(usage: dict[str, Any], session: Any) -> None: usage["context_estimated"] = True +# No context window holds a billion tokens, so a stored reading past this was +# not written by a call; one past float range would also raise on the way to a +# percentage. Either is read as no reading. +_READING_CEILING = 1_000_000_000 + + def _last_window_reading(messages: list[Any]) -> int: - """The window the newest reading-carrying reply measured, or 0 for none.""" + """The window the newest reading-carrying row records, or 0 for none.""" for message in reversed(messages): used = message.get("context_used") if isinstance(message, dict) else None - if isinstance(used, int) and not isinstance(used, bool) and used > 0: + if isinstance(used, int) and not isinstance(used, bool) and 0 < used <= _READING_CEILING: return used return 0 diff --git a/raven/rpc/methods/turn.py b/raven/rpc/methods/turn.py index 930e80d1d..e004a3678 100644 --- a/raven/rpc/methods/turn.py +++ b/raven/rpc/methods/turn.py @@ -749,7 +749,7 @@ async def turn_unsubscribe( return {"unsubscribed": unsubscribed} -def _cancel_payload(turn_id: str) -> dict[str, Any]: +def _cancel_payload(turn_id: str, usage: dict[str, Any] | None = None) -> dict[str, Any]: """The cancelled-turn error's payload, carrying the cancelled turn's id when there is one. The id is whatever identified the turn being cancelled: the one a named @@ -758,6 +758,10 @@ def _cancel_payload(turn_id: str) -> dict[str, Any]: key is omitted rather than sent empty, so that "absent" has one representation on the wire -- a consumer that correlates a request to a turn treats both the same way, as not its own. + + ``usage`` is what the stopped turn measured, in ``message.complete``'s + shape: the record the turn leaves carries that reading, and this event is + the only one its client is sent. Omitted when the turn read no window. """ payload: dict[str, Any] = { "code": _TURN_FAILED_CODE, @@ -766,9 +770,21 @@ def _cancel_payload(turn_id: str) -> dict[str, Any]: } if turn_id: payload["turn_id"] = turn_id + if usage: + payload["usage"] = usage return payload +def _stopped_usage(live_usages: dict[str, dict[str, Any]] | None, turn_id: str) -> dict[str, Any] | None: + """The usage sink of the turn about to be stopped, while it is still there. + + Read before the cancel: the turn's end takes its entry back. What it holds + is the reading the stopped turn files, because the loop moves this sink's + reading on the same terms as the one it stores. + """ + return (live_usages or {}).get(turn_id) if turn_id else None + + def _refused(lane: str) -> dict[str, Any]: """``turn.cancel``'s answer when it reached nothing on ``lane``. @@ -787,6 +803,7 @@ async def turn_cancel( emitter: SubscriptionEmitter | None = None, direct_targets: dict[str, dict[str, str]] | None = None, turn_ids: dict[str, str] | None = None, + live_usages: dict[str, dict[str, Any]] | None = None, ) -> dict[str, Any]: """``turn.cancel`` — cancel the in-flight turn + notify subscribers. @@ -827,7 +844,14 @@ async def turn_cancel( turn: those are registered under ``direct_lane(...)``, never under the bare session key. The emit still goes to the session, which is what the client subscribes to; the addressee rides the payload. + + ``live_usages`` is what each running turn has measured, by turn id (see + ``RpcTurnRunner``); a stopped turn's error carries its reading from there. """ + # Function-level: the rule lives with the sink that applies it to a failed + # turn, and that module wires the agent loop, which this one must not import. + from raven.rpc.spine import reported_usage + parsed = TurnCancelParams.model_validate(params) lane = ( @@ -844,13 +868,17 @@ async def turn_cancel( if running is None or running[0] != parsed.turn_id: return _refused(lane) handle = running[1] + measured = _stopped_usage(live_usages, parsed.turn_id) await handle.cancel() if emitter is not None: await emitter.emit( parsed.session_key, { "type": "error", - "payload": _tag(_cancel_payload(parsed.turn_id), (direct_targets or {}).get(lane)), + "payload": _tag( + _cancel_payload(parsed.turn_id, reported_usage(measured)), + (direct_targets or {}).get(lane), + ), }, ) await handle.result() @@ -860,6 +888,7 @@ async def turn_cancel( if reach is None: return _refused(lane) handle, inject_id = reach + measured = _stopped_usage(live_usages, inject_id or (turn_ids or {}).get(lane, "")) await handle.cancel() @@ -874,7 +903,7 @@ async def turn_cancel( { "type": "error", "payload": _tag( - _cancel_payload(inject_id or (turn_ids or {}).get(lane, "")), + _cancel_payload(inject_id or (turn_ids or {}).get(lane, ""), reported_usage(measured)), (direct_targets or {}).get(lane), ), }, @@ -937,6 +966,7 @@ def register_turn_methods( default_channel: str = "tui", agent_loop_factory: "AgentLoopFactory | None" = None, question_broker: "QuestionBroker | None" = None, + live_usages: dict[str, dict[str, Any]] | None = None, ) -> None: """Register ``turn.{send,subscribe,unsubscribe,cancel}`` on a dispatcher. @@ -949,6 +979,9 @@ def register_turn_methods( client omits one — and it MUST match the channel the outlet was registered under (the hub routes deliverables by ``source.channel``), or the reply is dropped. Defaults to ``"tui"``. + + ``live_usages`` must be the dict ``build_rpc_spine`` was given, or a stopped + turn's cancel has no reading to report. """ bind_scheduler(scheduler) @@ -973,7 +1006,9 @@ async def _unsubscribe(params: dict[str, Any]) -> dict[str, Any]: return await turn_unsubscribe(params, emitter=emitter) async def _cancel(params: dict[str, Any]) -> dict[str, Any]: - return await turn_cancel(params, emitter=emitter, direct_targets=direct_targets, turn_ids=turn_ids) + return await turn_cancel( + params, emitter=emitter, direct_targets=direct_targets, turn_ids=turn_ids, live_usages=live_usages + ) dispatcher.register("turn.send", _send) dispatcher.register("turn.subscribe", _subscribe) diff --git a/raven/rpc/models.py b/raven/rpc/models.py index b6b1841d5..64ce7df8b 100644 --- a/raven/rpc/models.py +++ b/raven/rpc/models.py @@ -844,6 +844,14 @@ class ErrorEventPayload(_Strict): "bounded text already is the whole account." ), ) + usage: TurnUsage | None = Field( + None, + description=( + "What a failed or stopped turn measured, in message.complete's shape: its last reported call, " + "so a client's context gauge moves on with the turn's stored record. Absent when that call " + "read no window." + ), + ) target: DirectTarget | None = None diff --git a/raven/rpc/spine.py b/raven/rpc/spine.py index 42e708618..09d10c64c 100644 --- a/raven/rpc/spine.py +++ b/raven/rpc/spine.py @@ -161,6 +161,20 @@ class RpcTurnRunner(AgentTurnRunner): - it fires the synthetic tool.complete when the turn replied via the message tool (the loop's general tool path skips the message tool), so the UI records that the agent acted. + + ``live_usages`` holds each streamed turn's usage sink under its turn id from + the moment the turn starts until the sink sees it end. A turn that ends + without an answer reports what it measured from there: a failed one on its + ``error`` event, a stopped one on the error ``turn.cancel`` sends while it + unwinds. Keyed by turn rather than by lane, because a queued turn shares + its lane with the turn that is running. + + It sits beside ``usages``, the per-lane copy an answered turn's + ``message.complete`` reads. Both hold the loop's usage sink, whose window + reading moves only on a call that read one -- the terms the turn's stored + reading moves on -- so either end reports the reading its record keeps. + ``reported_usage`` adds one rule for an error: a turn that read no window + sends no usage. """ def __init__( @@ -171,6 +185,7 @@ def __init__( readback_texts: dict[str, str], approval_responder: ApprovalResponder | None = None, credential_asker: CredentialAsker | None = None, + live_usages: dict[str, dict[str, Any]] | None = None, ) -> None: super().__init__(agent_loop, stream=True) self._emitter = emitter @@ -178,6 +193,7 @@ def __init__( self._readback_texts = readback_texts self._approval_responder = approval_responder self._credential_asker = credential_asker + self._live_usages = live_usages if live_usages is not None else {} async def run(self, req: TurnRequest, emit: Emit, drain: Drain) -> TurnOutcome: cid = conversation_id(req) @@ -252,6 +268,8 @@ async def _on_review(phase: str, tool_name: str) -> None: self._readback_texts[req.conversation] = text return outcome usage_sink: dict[str, Any] = {} + if req.turn_id: + self._live_usages[req.turn_id] = usage_sink outcome = await self._loop.run_turn(req, emit, drain, stream=True, usage_sink=usage_sink) # A synthetic tool.complete when the message tool fired (the loop @@ -528,6 +546,7 @@ async def emit_error( target: dict[str, str] | None = None, turn_id: str = "", full_detail: str = "", + usage: dict[str, Any] | None = None, ) -> None: """A turn's failure, tagged with the turn it belongs to when known. @@ -541,6 +560,9 @@ async def emit_error( ``full_detail`` is the failure's long account beside the bounded ``detail``, sent only when the runner had more to say. + + ``usage`` is what the failed turn measured, in ``message.complete``'s + shape, sent only when it read a window (``reported_usage``). """ payload: dict[str, Any] = {"code": code, "message": message, "reason": reason} if turn_id: @@ -549,11 +571,26 @@ async def emit_error( payload["detail"] = detail if full_detail: payload["full_detail"] = full_detail + if usage: + payload["usage"] = usage if target is not None: payload["target"] = target await self._emitter.emit(self._subscription(conversation_id, target), {"type": "error", "payload": payload}) +def reported_usage(measured: dict[str, Any] | None) -> dict[str, Any] | None: + """What a turn that ended without an answer measured, for the error that says so. + + Only a usage that read the window is sent. A last call that reported zeros, or + a sink that holds no reading at all, has nothing to tell the context ring, and + sent it would draw the ring empty. + """ + if not measured: + return None + used = measured.get("context_used") + return dict(measured) if isinstance(used, int) and used > 0 else None + + def _make_rpc_sink( hub: DeliveryHub, outlet: RpcOutlet, @@ -563,6 +600,7 @@ def _make_rpc_sink( direct_targets: dict[str, dict[str, str]], on_turn_end: Callable[[str], None] | None, on_turn_start: Callable[[str, str], str | None] | None = None, + live_usages: dict[str, dict[str, Any]] | None = None, ) -> Callable[[TurnEvent], Awaitable[None]]: """Adapt the hub into the scheduler's EventSink for the TUI. Deliverables route through the hub; a turn's end fires message.complete / error after the @@ -574,7 +612,11 @@ def _make_rpc_sink( injected text -- so this sink can bind the lane to it and open it on the wire; ``turn.send`` could not, since it answered before the fallback was decided. This sink is build_rpc_spine's alone — the CLI keeps its own - lifecycle-dropping sink.""" + lifecycle-dropping sink. + + ``live_usages`` is the runner's map of what each turn has measured (see + ``RpcTurnRunner``); every turn's end takes its entry back.""" + live = live_usages if live_usages is not None else {} async def _finish(conversation_id: str) -> None: # close_stream clears the hub's per-stream state (so the next turn on this @@ -616,6 +658,7 @@ def _drop(conversation_id: str, *, owns: bool) -> dict[str, str] | None: async def sink(event: TurnEvent) -> None: if isinstance(event, TurnEnded): + live.pop(event.turn_id, None) await _finish(event.conversation_id) usage = usages.get(event.conversation_id) or { "prompt_tokens": 0, @@ -635,6 +678,7 @@ async def sink(event: TurnEvent) -> None: await outlet.emit_complete(event.conversation_id, event.turn_id, usage, target, round(event.latency_ms)) return if isinstance(event, TurnFailed): + measured = live.pop(event.turn_id, None) await _finish(event.conversation_id) owns = _owns_lane(event.conversation_id, event.turn_id) target = _drop(event.conversation_id, owns=owns) @@ -654,6 +698,7 @@ async def sink(event: TurnEvent) -> None: # started yet. turn_id=event.turn_id, full_detail=event.full_detail or "", + usage=reported_usage(measured), ) return if isinstance(event, TurnStarted): @@ -698,6 +743,7 @@ def build_rpc_spine( readback_texts: dict[str, str] | None = None, approval_responder: ApprovalResponder | None = None, credential_asker: CredentialAsker | None = None, + live_usages: dict[str, dict[str, Any]] | None = None, user_pool: int = 1, system_pool: int = 1, direct_pool: int = 8, @@ -729,7 +775,11 @@ def build_rpc_spine( permission. The runner binds it per turn on the same gate the asker uses: a USER turn always, a SUBAGENT relay when a surface is watching its conversation, and no other origin -- a CRON or otherwise unattended turn is - refused at the ask tier. ``credential_asker`` is bound on the same terms.""" + refused at the ask tier. ``credential_asker`` is bound on the same terms. + + ``live_usages`` is what each running turn has measured, by turn id (see + ``RpcTurnRunner``). Pass the same dict ``register_turn_methods`` is given, + so a stopped turn's cancel can report it; defaults to a private map.""" hub = DeliveryHub() if direct_targets is None: direct_targets = {} @@ -737,6 +787,8 @@ def build_rpc_spine( hub.register(outlet) turn_ids: dict[str, str] = {} usages: dict[str, dict[str, Any]] = {} + if live_usages is None: + live_usages = {} if readback_texts is None: readback_texts = {} scheduler = Scheduler( @@ -747,9 +799,12 @@ def build_rpc_spine( readback_texts, approval_responder=approval_responder, credential_asker=credential_asker, + live_usages=live_usages, ), OriginPools(user=user_pool, system=system_pool, direct=direct_pool), - _make_rpc_sink(hub, outlet, channel, turn_ids, usages, direct_targets, on_turn_end, on_turn_start), + _make_rpc_sink( + hub, outlet, channel, turn_ids, usages, direct_targets, on_turn_end, on_turn_start, live_usages=live_usages + ), ) async def teardown() -> None: diff --git a/rpc-schema/openrpc.json b/rpc-schema/openrpc.json index cbed12fc2..4d2c8054a 100644 --- a/rpc-schema/openrpc.json +++ b/rpc-schema/openrpc.json @@ -15325,6 +15325,10 @@ "type": "string", "description": "The failure's long account beside the bounded detail, for a reader who expands the row: the runner's own text, credentials scrubbed, capped at 8000 characters. Absent when the bounded text already is the whole account." }, + "usage": { + "$ref": "#/components/schemas/TurnUsage", + "description": "What a failed or stopped turn measured, in message.complete's shape: its last reported call, so a client's context gauge moves on with the turn's stored record. Absent when that call read no window." + }, "target": { "$ref": "#/components/schemas/DirectTarget" } diff --git a/tests/test_agent_loop_empty_recovery.py b/tests/test_agent_loop_empty_recovery.py index 38ff07345..c86b64b4c 100644 --- a/tests/test_agent_loop_empty_recovery.py +++ b/tests/test_agent_loop_empty_recovery.py @@ -2053,6 +2053,12 @@ def get_default_model(self) -> str: return "stub" +def _as_sent(messages: list[dict]) -> list[dict]: + """A request as the wire carries it: a provider sends only its allowed keys, so the loop's own + underscore bookkeeping -- a window reading stamped on the question, say -- is not what was asked.""" + return [{key: value for key, value in message.items() if not key.startswith("_")} for message in messages] + + @pytest.mark.asyncio async def test_a_switch_after_a_streamed_error_asks_the_fallback_the_original_question(workspace, monkeypatch): """The prefill is the thought of the model being left. A streamed error switches models in the loop @@ -2068,7 +2074,7 @@ async def test_a_switch_after_a_streamed_error_asks_the_fallback_the_original_qu assert final == "fb's answer" and outcome.status == "completed" asked = [(model, messages[-1]["role"]) for model, messages in provider.requests] assert asked == [("stub", "user"), ("stub", "assistant"), ("fb", "user")] - assert provider.requests[2][1] == provider.requests[0][1] + assert _as_sent(provider.requests[2][1]) == _as_sent(provider.requests[0][1]) class _StreamErrorsThenFallbackReplies(LLMProvider): diff --git a/tests/test_agent_loop_usage_sink.py b/tests/test_agent_loop_usage_sink.py index da4326f80..2632df0d2 100644 --- a/tests/test_agent_loop_usage_sink.py +++ b/tests/test_agent_loop_usage_sink.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio import json import tempfile import threading @@ -24,17 +25,20 @@ from raven.agent.hook import AgentHook, HookDecision from raven.agent.loop import AgentLoop from raven.agent.loop.bundles import EngineWiring, HostWiring, ToolWiring, TurnPolicy +from raven.agent.loop.recovery import RecoveryLimits from raven.config.loader import load_config from raven.config.raven import ContextConfig from raven.contracts.llm_provider import ToolCallRequest from raven.contracts.token_strategy import UsageSnapshot from raven.contracts.tool import Tool from raven.providers import rates, usage_record -from raven.providers.base import LLMProvider, LLMResponse +from raven.providers.base import ErrorClassification, LLMProvider, LLMResponse from raven.providers.binding import ModelBinding, use_binding +from raven.rpc.models import TurnUsage +from raven.rpc.spine import reported_usage from raven.session.manager import SessionManager from raven.spine.message import ChatType, Source -from raven.spine.turn import Origin, TurnRequest +from raven.spine.turn import AnswerlessTurnError, Origin, TurnRequest from raven.token_wise import usage_context from raven.token_wise.registry import StrategyRegistry from raven.token_wise.usage_tracker import UsageTracker @@ -582,7 +586,13 @@ def _costing_agent( ) -async def _turn(agent: AgentLoop, sink: dict, session_key: str = "s1") -> None: +async def _turn(agent: AgentLoop, sink: dict, session_key: str = "s1", *, streamed: bool = False) -> None: + """One user turn. ``streamed`` drives it the way the RPC runner does every + turn a gauge is shown for: through the streaming exit.""" + + async def _watch(_text: str) -> None: + return None + await agent._process_message( TurnRequest( origin=Origin.USER, @@ -591,6 +601,7 @@ async def _turn(agent: AgentLoop, sink: dict, session_key: str = "s1") -> None: ), session_key=session_key, usage_sink=sink, + on_token_delta=_watch if streamed else None, ) @@ -937,3 +948,317 @@ async def test_a_reply_a_hook_put_in_place_of_the_calls_keeps_that_calls_reading assert sink["context_used"] == 8000 assert (usage["context_used"], usage["context_percent"]) == (sink["context_used"], sink["context_percent"]) assert usage["context_estimated"] is False + + +_ANSWERS_THROUGH_THE_MESSAGE_TOOL = LLMResponse( + content="", + tool_calls=[ToolCallRequest(id="c1", name="message", arguments={"content": "here it is"})], + usage=_measured(6000, 200), +) + + +def _says_nothing(usage: dict) -> LLMResponse: + return LLMResponse(content="", finish_reason="stop", usage=usage) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limits", "empty_replies", "live"), + [ + ( + RecoveryLimits(llm_error_retry_delays=()), + [_says_nothing(_measured(7000, 0)), _says_nothing(_measured(7400, 0))], + 7400, + ), + (RecoveryLimits(enabled=False), [_says_nothing(_measured(7000, 0))], 7000), + (RecoveryLimits(llm_error_retry_delays=()), [_says_nothing(_measured(7000, 0)), _says_nothing({})], 7000), + ], + ids=["recovery_spent", "recovery_off", "last_empty_reply_unreported"], +) +async def test_a_turn_that_ends_on_an_empty_reply_keeps_that_calls_reading( + workspace, monkeypatch, limits, empty_replies, live +): + """The model answered through the message tool and then said nothing more, + so the turn completes on a call whose reply is not filed. The live meter + showed the last window reported, which ends at the newest row the turn + kept, so a resume must show it too: recovery's scaffolding is dropped, and + a call that reported nothing leaves the reading before it standing.""" + agent = AgentLoop( + provider=ScriptedProvider("stub", [_ANSWERS_THROUGH_THE_MESSAGE_TOOL, *empty_replies]), + workspace=workspace, + model="stub", + engine=EngineWiring(context_window_tokens=40000, context_config=ContextConfig(drop_segments=["curator"])), + policy=TurnPolicy(max_iterations=6, empty_recovery=limits), + tools=ToolWiring(restrict_to_workspace=True), + ) + + async def _delivered(_content: str, _media: list[str]) -> None: + return None + + agent.tools.get("message").set_send_callback(_delivered) + sink: dict = {} + await _turn(agent, sink, session_key="tui:c1") + + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == live + assert (usage["context_used"], usage["context_percent"]) == (sink["context_used"], sink["context_percent"]) + assert usage["context_estimated"] is False + + +def _fails_to_answer(usage: dict) -> LLMResponse: + return LLMResponse( + content="Error calling LLM (network@stub): connection reset", + finish_reason="error", + usage=usage, + error_classification=ErrorClassification("network", retryable=False), + ) + + +@pytest.mark.asyncio +async def test_an_answered_turn_that_ends_on_a_provider_error_keeps_that_calls_reading(workspace, monkeypatch): + """The answer went out through the message tool, so the provider error + that follows ends the turn without failing it, and its reply is not filed + either. The live meter showed that call's window, which ends at the newest + row the turn kept.""" + agent = AgentLoop( + provider=ScriptedProvider("stub", [_ANSWERS_THROUGH_THE_MESSAGE_TOOL, _fails_to_answer(_measured(7100, 0))]), + workspace=workspace, + model="stub", + engine=EngineWiring(context_window_tokens=40000, context_config=ContextConfig(drop_segments=["curator"])), + policy=TurnPolicy(max_iterations=6, empty_recovery=RecoveryLimits(llm_error_retry_delays=())), + tools=ToolWiring(restrict_to_workspace=True), + ) + + async def _delivered(_content: str, _media: list[str]) -> None: + return None + + agent.tools.get("message").set_send_callback(_delivered) + sink: dict = {} + await _turn(agent, sink, session_key="tui:c1") + + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == 7100 + assert (usage["context_used"], usage["context_percent"]) == (sink["context_used"], sink["context_percent"]) + assert usage["context_estimated"] is False + + +def _failing_agent(workspace: Path, limits: RecoveryLimits, script: list[LLMResponse]) -> AgentLoop: + return AgentLoop( + provider=ScriptedProvider("stub", script), + workspace=workspace, + model="stub", + engine=EngineWiring(context_window_tokens=40000, context_config=ContextConfig(drop_segments=["curator"])), + policy=TurnPolicy(max_iterations=6, empty_recovery=limits), + tools=ToolWiring(restrict_to_workspace=True), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("limits", "script", "live"), + [ + (RecoveryLimits(enabled=False), [_says_nothing(_measured(7000, 0))], 7000), + ( + RecoveryLimits(llm_error_retry_delays=()), + [_says_nothing(_measured(7000, 0)), _says_nothing(_measured(7400, 0))], + 7400, + ), + (RecoveryLimits(llm_error_retry_delays=()), [_says_nothing(_measured(7000, 0)), _says_nothing({})], 7000), + (RecoveryLimits(llm_error_retry_delays=()), [_fails_to_answer(_measured(7100, 0))], 7100), + ], + ids=["first_reply_empty", "every_retry_empty", "last_retry_unreported", "provider_error"], +) +async def test_a_failed_turn_that_filed_no_reply_keeps_its_reading_on_the_failure( + workspace, monkeypatch, limits, script, live +): + """No call of the turn left a row behind: each reply was empty or an error, + and the question was filed before the first call. The failure the turn + files is then the only row it adds, so that is where the live meter's last + reading has to go for a resume to show it.""" + agent = _failing_agent(workspace, limits, script) + sink: dict = {} + + with pytest.raises(AnswerlessTurnError): + await _turn(agent, sink, session_key="tui:c1") + + failure = _stored_replies(workspace, "tui:c1")[-1] + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == live + assert failure["turn_ended"]["status"] == "failed" + assert failure["context_used"] == live + assert (usage["context_used"], usage["context_percent"]) == (sink["context_used"], sink["context_percent"]) + assert usage["context_estimated"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("cut", "status"), + [(asyncio.CancelledError, "cancelled"), (RuntimeError, "failed")], + ids=["cancelled", "crashed"], +) +async def test_a_turn_cut_short_after_a_call_keeps_that_calls_reading_on_its_marker( + workspace, monkeypatch, cut, status +): + """A stop or a crash files the same marker an answerless turn does, and it + can come right after a call whose empty reply left no row behind.""" + agent = _failing_agent(workspace, RecoveryLimits(llm_error_retry_delays=()), [_says_nothing(_measured(7000, 0))]) + strip = agent._strip_think + seen: list[str | None] = [] + + def _cut_on_the_retry(text): + seen.append(text) + if len(seen) > 1: + raise cut("the turn was cut short") + return strip(text) + + monkeypatch.setattr(agent, "_strip_think", _cut_on_the_retry) + + with pytest.raises(cut): + await _turn(agent, {}, session_key="tui:c1") + + marker = _stored_replies(workspace, "tui:c1")[-1] + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert marker["turn_ended"]["status"] == status + assert marker["context_used"] == 7000 + assert (usage["context_used"], usage["context_estimated"]) == (7000, False) + + +@pytest.mark.asyncio +async def test_a_failed_turn_that_measured_nothing_files_no_reading(workspace): + """Nothing measured means nothing to file, so a resume falls back to its estimate.""" + agent = _failing_agent(workspace, RecoveryLimits(llm_error_retry_delays=()), [_fails_to_answer({})]) + + with pytest.raises(AnswerlessTurnError): + await _turn(agent, {}, session_key="tui:c1") + + assert "context_used" not in _stored_replies(workspace, "tui:c1")[-1] + + +def _fails_in_transport(usage: dict) -> LLMResponse: + return LLMResponse( + content="The upstream reported a failed call instead of an answer.", + finish_reason="error", + usage=usage, + error_classification=ErrorClassification("upstream_transport_failure", retryable=False), + ) + + +_LISTS_THE_WORKSPACE = LLMResponse( + content="", + tool_calls=[ToolCallRequest(id="c1", name="list_dir", arguments={"path": "."})], + usage=_measured(6000, 200), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answered", [False, True], ids=["failed", "answered"]) +async def test_a_call_that_failed_in_transport_reads_no_window(workspace, monkeypatch, answered): + """The upstream's account of a call it failed is no reading of the window: + it may be the very accounting that flagged the call, billed at a size the + request cannot have been. Neither the live meter nor the stored record takes + it, so both stay on the call before.""" + first = _ANSWERS_THROUGH_THE_MESSAGE_TOOL if answered else _LISTS_THE_WORKSPACE + agent = _failing_agent( + workspace, + RecoveryLimits(llm_error_retry_delays=()), + [first, _fails_in_transport(_measured(5, 0, cost_usd=0.002))], + ) + + async def _delivered(_content: str, _media: list[str]) -> None: + return None + + agent.tools.get("message").set_send_callback(_delivered) + sink: dict = {} + if answered: + await _turn(agent, sink, session_key="tui:c1") + else: + with pytest.raises(AnswerlessTurnError): + await _turn(agent, sink, session_key="tui:c1") + + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == 6200 + # The cost is the turn's, a flow, and the failed call was still billed. + assert sink["cost_usd"] == pytest.approx(0.002) + assert (usage["context_used"], usage["context_estimated"]) == (6200, False) + + +# What the incident's upstream returned: an empty reply to a prompt of thousands +# of tokens, billed as one. Through the streaming exit it reaches the transport +# verdict on its accounting alone. +_A_CALL_THE_UPSTREAM_NEVER_READ = LLMResponse( + content="", finish_reason="stop", usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answered", [False, True], ids=["failed", "answered"]) +async def test_a_streamed_call_that_failed_in_transport_reads_no_window(workspace, monkeypatch, answered): + """The pairing production makes: a streamed turn with a usage sink. The + streaming exit returns the call the upstream failed with the verdict the + other exit gives it, and its one-token account must not become the window + on the live meter, the kept row or the failure marker.""" + first = _ANSWERS_THROUGH_THE_MESSAGE_TOOL if answered else _LISTS_THE_WORKSPACE + agent = _failing_agent( + workspace, RecoveryLimits(enabled=False, llm_error_retry_delays=()), [first, _A_CALL_THE_UPSTREAM_NEVER_READ] + ) + + async def _delivered(_content: str, _media: list[str]) -> None: + return None + + agent.tools.get("message").set_send_callback(_delivered) + sink: dict = {} + if answered: + await _turn(agent, sink, session_key="tui:c1", streamed=True) + else: + with pytest.raises(AnswerlessTurnError): + await _turn(agent, sink, session_key="tui:c1", streamed=True) + + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == 6200 + assert (usage["context_used"], usage["context_estimated"]) == (6200, False) + + +@pytest.mark.asyncio +async def test_a_last_call_that_read_no_window_leaves_the_live_reading_where_the_record_is(workspace, monkeypatch): + """A usage block of zeros is no reading, so the record keeps the call + before it. The live meter has to keep it too: emptied, the failure's error + carries no usage at all and the ring stays a turn behind the record.""" + agent = _failing_agent( + workspace, RecoveryLimits(enabled=False), [_LISTS_THE_WORKSPACE, _says_nothing(_measured(0, 0))] + ) + sink: dict = {} + + with pytest.raises(AnswerlessTurnError): + await _turn(agent, sink, session_key="tui:c1", streamed=True) + + failure = _stored_replies(workspace, "tui:c1")[-1] + usage = await _resumed_usage(workspace, monkeypatch, "tui:c1") + + assert sink["context_used"] == 6200 + assert (reported_usage(sink) or {}).get("context_used") == 6200 + assert failure["context_used"] == 6200 + assert (usage["context_used"], usage["context_estimated"]) == (6200, False) + + +@pytest.mark.asyncio +async def test_a_turn_whose_calls_read_no_window_still_reports_a_whole_usage(workspace): + """With no earlier reading to keep, a call that reads none still fills the + sink: an answered turn's ``message.complete`` sends it as it is, and the + counts are what make it a ``TurnUsage``.""" + agent = _failing_agent( + workspace, + RecoveryLimits(enabled=False), + [LLMResponse(content="done", finish_reason="stop", usage=_measured(0, 0))], + ) + sink: dict = {} + + await _turn(agent, sink, session_key="tui:c1", streamed=True) + + assert TurnUsage.model_validate(sink).context_used == 0 + assert reported_usage(sink) is None diff --git a/tests/test_cli_tui_commands.py b/tests/test_cli_tui_commands.py index f4f2f0344..ad9322a98 100644 --- a/tests/test_cli_tui_commands.py +++ b/tests/test_cli_tui_commands.py @@ -512,6 +512,23 @@ async def test_the_acp_snapshot_backfill_is_scheduled_on_the_tui_server_path( ctx["schedule_backfill"].assert_called_once_with(ctx["agent_loop"].subagents) +async def test_the_tui_spine_and_turn_cancel_share_one_map_of_what_running_turns_measured( + monkeypatch: pytest.MonkeyPatch, rpc_server_deps +) -> None: + """Hand-wired like the backfill above: a stopped turn's cancel reads what + the turn measured from the map the spine's runner fills, so the two have to + be handed the same one.""" + ctx = rpc_server_deps + register = MagicMock() + monkeypatch.setattr("raven.rpc.methods.register_aligned_methods_except_system", register) + + await _run_until_done_with_immediate_proc_done(monkeypatch, ctx) + + spine_map = ctx["fake_build_rpc_spine"].call_args.kwargs["live_usages"] + assert isinstance(spine_map, dict) + assert register.call_args.kwargs["live_usages"] is spine_map + + async def _run_until_done_with_immediate_proc_done(monkeypatch, ctx): """Helper: drive ``_run_rpc_server_until_done`` with proc_done set immediately so the function exits as fast as possible (handshake timeout path — still diff --git a/tests/test_rpc_bootstrap.py b/tests/test_rpc_bootstrap.py index da5391b38..5df355df7 100644 --- a/tests/test_rpc_bootstrap.py +++ b/tests/test_rpc_bootstrap.py @@ -171,6 +171,37 @@ def spy(*args, **kwargs): await stack.teardown() +async def test_the_spine_and_turn_cancel_share_one_map_of_what_running_turns_measured(monkeypatch) -> None: + """A stopped turn's cancel reports what the turn measured, read from the map + the spine's runner fills: two maps would leave every stop with nothing to + report.""" + loop = _FakeLoop(_FakeCron()) + import raven.rpc.methods as methods_module + import raven.rpc.spine as spine_module + + seen: dict = {} + real_spine = spine_module.build_rpc_spine + real_register = methods_module.register_turn_methods + + def spy_spine(*args, **kwargs): + seen["spine"] = kwargs.get("live_usages") + return real_spine(*args, **kwargs) + + def spy_register(*args, **kwargs): + seen["cancel"] = kwargs.get("live_usages") + return real_register(*args, **kwargs) + + monkeypatch.setattr(spine_module, "build_rpc_spine", spy_spine) + # The name the aligned registration calls, so the map is followed through it. + monkeypatch.setattr(methods_module, "register_turn_methods", spy_register) + stack = await bootstrap.build_rpc_stack(_sink, agent_loop=loop) + try: + assert isinstance(seen["spine"], dict) + assert seen["cancel"] is seen["spine"] + finally: + await stack.teardown() + + def test_the_turn_pools_come_from_the_gateway_section_and_zero_is_allowed(tmp_path, monkeypatch) -> None: path = tmp_path / "config.json" monkeypatch.setattr("raven.config.loader.get_config_path", lambda: path) diff --git a/tests/test_rpc_session.py b/tests/test_rpc_session.py index b55edc9cf..e982a0ede 100644 --- a/tests/test_rpc_session.py +++ b/tests/test_rpc_session.py @@ -2095,11 +2095,17 @@ async def test_session_resume_reports_the_window_the_last_call_measured( assert usage["context_estimated"] is False -@pytest.mark.parametrize("junk", [0, -5, True, "50000", 1.5], ids=["zero", "negative", "bool", "string", "float"]) +@pytest.mark.parametrize( + "junk", + [0, -5, True, "50000", 1.5, 10**12, 10**400], + ids=["zero", "negative", "bool", "string", "float", "past_any_window", "past_float_range"], +) async def test_session_resume_reads_past_a_stored_reading_that_is_no_token_count( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, junk: object ) -> None: - """Only a positive integer is a reading -- and ``True`` is an ``int`` to Python.""" + """Only a positive integer a window could hold is a reading -- ``True`` is an + ``int`` to Python, and a count past float range would raise on the way to a + percentage, turning the reader's conversation into a fresh empty one.""" cfg = load_config() cfg.agents.defaults.workspace = str(tmp_path) cfg.agents.defaults.context_window_tokens = 200_000 diff --git a/tests/test_rpc_spine.py b/tests/test_rpc_spine.py index 5610f1900..0c6899dfc 100644 --- a/tests/test_rpc_spine.py +++ b/tests/test_rpc_spine.py @@ -17,6 +17,7 @@ from raven.permissions.builtin import BuiltinRulings from raven.permissions.gate import PermissionGate from raven.permissions.turn import current_turn +from raven.rpc.models import ErrorEvent from raven.rpc.spine import ( RpcOutlet, RpcTurnRunner, @@ -2138,3 +2139,132 @@ async def run_turn(self, req, emit, drain, *, stream, usage_sink=None) -> TurnOu payload = emitter.emitted[-1][1]["payload"] assert payload["detail"] == "RuntimeError: boom" assert payload["full_detail"] == "RuntimeError: boom, at length" + + +# --- a failed turn's usage reaches the page on its error event --- + + +_A_FAILED_TURNS_LAST_CALL = { + "prompt_tokens": 7000, + "completion_tokens": 0, + "total_tokens": 7000, + "context_max": 40000, + "context_used": 7000, + "context_percent": 18, +} + + +class _MeasuresThenFails(_RunTurnLoop): + """Fills the usage sink the way a loop that made calls does, then gives up.""" + + async def run_turn(self, req, emit, drain, *, stream, usage_sink=None, text_sink=None) -> TurnOutcome: + if usage_sink is not None and self._usage: + usage_sink.update(self._usage) + raise AnswerlessTurnError("The model returned no content on 3 attempt(s).") + + +async def test_runner_keeps_what_a_turn_measured_under_its_own_id(): + """A turn that fails still measured the calls it made, and its stored record + carries that reading, so its end has to report it as well. Kept under the + turn's id, never in the per-lane register a queued turn shares.""" + usages: dict[str, dict] = {} + live: dict[str, dict] = {} + runner = RpcTurnRunner( + _MeasuresThenFails(usage=_A_FAILED_TURNS_LAST_CALL), FakeEmitter(), usages, {}, live_usages=live + ) + req = TurnRequest(origin=Origin.USER, source=_src(), text="hi", conversation="tui:c1", turn_id="t-1") + _events, emit = _collect() + + with pytest.raises(AnswerlessTurnError): + await runner.run(req, emit, lambda: []) + + assert live == {"t-1": _A_FAILED_TURNS_LAST_CALL} + assert "tui:c1" not in usages + + +_A_LAST_CALL_THAT_READ_NO_WINDOW = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "context_max": 40000, + "context_used": 0, + "context_percent": 0, +} + + +@pytest.mark.parametrize( + ("measured", "carried"), + [ + (_A_FAILED_TURNS_LAST_CALL, _A_FAILED_TURNS_LAST_CALL), + (None, None), + (_A_LAST_CALL_THAT_READ_NO_WINDOW, None), + ({"cost_usd": 0.01, "cost_missing_calls": 0}, None), + ], + ids=["measured_a_call", "measured_nothing", "last_call_reported_zeros", "no_window_reading"], +) +async def test_a_failed_turn_hands_its_measured_usage_to_its_error(measured, carried): + """Without it the context ring stays on the turn before while a resume reads + the failed turn's own record, so the two disagree until the next turn. A + usage that read no window is not sent: drawn, it would empty the ring.""" + emitter = FakeEmitter() + scheduler, hub, turn_ids, teardown = build_rpc_spine(_MeasuresThenFails(usage=measured), emitter) + try: + handle = scheduler.submit(TurnRequest(origin=Origin.USER, source=_src(), text="hi", conversation="tui:c1")) + await handle.result() + finally: + await teardown() + + assert emitter.types() == ["error"] + event = emitter.emitted[-1][1] + assert event["payload"].get("usage") == carried + ErrorEvent.model_validate(event) + + +def _failure_sink(usages: dict, live: dict): + from raven.rpc.spine import RpcOutlet, _make_rpc_sink + from raven.spine.delivery import DeliveryHub + + hub = DeliveryHub() + emitter = FakeEmitter() + outlet = RpcOutlet("tui", emitter, {}) + hub.register(outlet) + return emitter, _make_rpc_sink(hub, outlet, "tui", {"tui:c1": "owner"}, usages, {}, None, live_usages=live) + + +async def test_a_non_owning_turns_failure_does_not_report_the_owners_usage(): + """The register is keyed by lane, and a queued turn shares the key with the + turn that owns it: its failure must not hand on the owner's reading.""" + from raven.spine.events import TurnFailed + + live = {"owner": dict(_A_FAILED_TURNS_LAST_CALL)} + emitter, sink = _failure_sink({"tui:c1": dict(_A_FAILED_TURNS_LAST_CALL)}, live) + + await sink(TurnFailed(error="boom", cancelled=False, conversation_id="tui:c1", turn_id="queued-1")) + + errors = [e for _k, e in emitter.emitted if e["type"] == "error"] + assert errors and "usage" not in errors[-1]["payload"] + assert live == {"owner": _A_FAILED_TURNS_LAST_CALL} + + +@pytest.mark.parametrize( + "end", + ["completed", "failed", "cancelled"], +) +async def test_a_turns_end_lets_go_of_what_it_measured(end): + """Every way a turn ends takes its reading back, or the map grows by a turn + for every one a process runs.""" + from raven.spine import TurnEnded, Usage + from raven.spine.events import TurnFailed + + live = {"t-1": dict(_A_FAILED_TURNS_LAST_CALL)} + _emitter, sink = _failure_sink({}, live) + if end == "completed": + event = TurnEnded( + usage=Usage(0, 0, 0), latency_ms=1.0, explicit_reply=False, conversation_id="tui:c1", turn_id="t-1" + ) + else: + event = TurnFailed(error="boom", cancelled=end == "cancelled", conversation_id="tui:c1", turn_id="t-1") + + await sink(event) + + assert live == {} diff --git a/tests/test_rpc_turn_cancel.py b/tests/test_rpc_turn_cancel.py index d7a2d66d7..9f6fcf15f 100644 --- a/tests/test_rpc_turn_cancel.py +++ b/tests/test_rpc_turn_cancel.py @@ -1243,3 +1243,87 @@ async def test_a_waiter_cancelled_mid_unwind_does_not_wedge_the_lane( unwind.set() _turn_mod.bind_scheduler(None) await teardown() + + +# --- a stopped turn's cancel reports what the turn measured --- + + +_A_STOPPED_TURNS_LAST_CALL = { + "prompt_tokens": 6000, + "completion_tokens": 200, + "total_tokens": 6200, + "context_max": 40000, + "context_used": 6200, + "context_percent": 16, +} + + +@pytest.mark.parametrize( + ("named", "measured", "carried"), + [ + (False, _A_STOPPED_TURNS_LAST_CALL, _A_STOPPED_TURNS_LAST_CALL), + (True, _A_STOPPED_TURNS_LAST_CALL, _A_STOPPED_TURNS_LAST_CALL), + (False, None, None), + (False, {**_A_STOPPED_TURNS_LAST_CALL, "context_used": 0, "context_percent": 0}, None), + ], + ids=["unnamed_cancel", "named_cancel", "measured_nothing", "last_call_read_no_window"], +) +async def test_a_stopped_turns_cancel_carries_what_it_measured( + emitter: SubscriptionEmitter, send_frame_capture: AsyncMock, named, measured, carried +) -> None: + """A stopped turn files what it measured with the record it leaves, so a + resume shows that reading; the cancel is the one event its client gets, and + without the reading on it the ring stays on the turn before.""" + from raven.rpc.methods import turn as _turn_mod + from raven.rpc.models import ErrorEvent + from raven.rpc.spine import build_rpc_spine + + class _MeasuresThenHangs: + tools: dict = {} + + def __init__(self) -> None: + self.measured = asyncio.Event() + + async def run_turn(self, req, emit, drain, *, stream=True, usage_sink=None, **_kwargs): + if usage_sink is not None and measured: + usage_sink.update(measured) + self.measured.set() + await asyncio.Event().wait() + + loop = _MeasuresThenHangs() + live: dict[str, dict] = {} + scheduler, _hub, turn_ids, teardown = build_rpc_spine( + loop, emitter, on_turn_end=_turn_mod.clear_active, live_usages=live + ) + # Through the registered handlers, so the map reaching turn.cancel is the + # one the wiring hands it, not one this test passes by hand. + d = Dispatcher() + register_turn_methods(d, emitter=emitter, scheduler=scheduler, turn_ids=turn_ids, live_usages=live) + try: + await turn_subscribe({"session_key": "tui:default"}, emitter=emitter) + sent = await d.dispatch( + { + "jsonrpc": "2.0", + "id": 1, + "method": "turn.send", + "params": {"session_key": "tui:default", "content": "hi"}, + } + ) + await asyncio.wait_for(loop.measured.wait(), 5) + params = {"session_key": "tui:default", **({"turn_id": sent["result"]["turn_id"]} if named else {})} + + cancel = await d.dispatch({"jsonrpc": "2.0", "id": 2, "method": "turn.cancel", "params": params}) + assert cancel["result"] == {"cancelled": True} + await asyncio.sleep(0.05) + finally: + _turn_mod.bind_scheduler(None) + await teardown() + + (cancelled,) = [ + e + for e in _collect_events(send_frame_capture) + if e.get("type") == "error" and e["payload"].get("reason") == "cancelled_by_client" + ] + assert cancelled["payload"].get("usage") == carried + ErrorEvent.model_validate(cancelled) + assert live == {} diff --git a/ui-tui/src/__tests__/chatStream.test.ts b/ui-tui/src/__tests__/chatStream.test.ts index 0bb461125..d2af49042 100644 --- a/ui-tui/src/__tests__/chatStream.test.ts +++ b/ui-tui/src/__tests__/chatStream.test.ts @@ -1398,3 +1398,61 @@ describe('createChatStream — cron.missed startup notice', () => { expect(sysCalls[0]).not.toMatch(/scheduled \d{2}-\d{2} /) }) }) + +describe('createChatStream — the context bar after a turn that ended without an answer', () => { + beforeEach(() => { + resetTurnState() + resetUiState() + turnController.fullReset() + }) + + const usage = { + completion_tokens: 0, + context_max: 40000, + context_used: 7000, + prompt_tokens: 7000, + total_tokens: 7000 + } + + it.each([ + ['failed', { code: -32099, message: 'turn_failed', reason: 'internal' as const }], + ['stopped', { code: -32099, message: 'turn_cancelled', reason: 'cancelled_by_client' as const }] + ])('moves to the window a %s turn measured', async (_end, payload) => { + const fake = makeFakeRpc() + const stream = createChatStream({ + appendMessage: () => {}, + rpcClient: fake, + sessionKey: 'tui:default', + sys: () => {} + }) + await stream.attach() + patchUiState({ busy: true, usage: { ...getUiState().usage, context_max: 40000, context_used: 3100 } }) + await stream.send('question') + + fake.__pushEvent({ type: 'message.start', payload: { turn_id: 'turn-1' } }) + fake.__pushEvent({ type: 'error', payload: { ...payload, turn_id: 'turn-1', usage } }) + + expect(getUiState().usage.context_used).toBe(7000) + }) + + it('stays where it was when the turn measured nothing', async () => { + const fake = makeFakeRpc() + const stream = createChatStream({ + appendMessage: () => {}, + rpcClient: fake, + sessionKey: 'tui:default', + sys: () => {} + }) + await stream.attach() + patchUiState({ busy: true, usage: { ...getUiState().usage, context_max: 40000, context_used: 3100 } }) + await stream.send('question') + + fake.__pushEvent({ type: 'message.start', payload: { turn_id: 'turn-1' } }) + fake.__pushEvent({ + type: 'error', + payload: { code: -32099, message: 'turn_failed', reason: 'internal', turn_id: 'turn-1' } + }) + + expect(getUiState().usage.context_used).toBe(3100) + }) +}) diff --git a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts index 7402466d6..484a42535 100644 --- a/ui-tui/src/__tests__/createGatewayEventHandler.test.ts +++ b/ui-tui/src/__tests__/createGatewayEventHandler.test.ts @@ -11,7 +11,7 @@ import { createGatewayEventHandler } from '../app/createGatewayEventHandler.js' import { getOverlayState, resetOverlayState } from '../app/overlayStore.js' import { turnController } from '../app/turnController.js' import { getTurnState, resetTurnState } from '../app/turnStore.js' -import { patchUiState, resetUiState } from '../app/uiStore.js' +import { getUiState, patchUiState, resetUiState } from '../app/uiStore.js' import { estimateTokensRough } from '../lib/text.js' const ref = (current: T) => ({ current }) @@ -583,6 +583,44 @@ describe('createGatewayEventHandler', () => { expect(appended[1]?.tools ?? []).toEqual([]) }) + it('moves the context bar to the window a failed turn measured', () => { + const onEvent = createGatewayEventHandler(buildCtx([])) + + patchUiState(state => ({ + ...state, + usage: { ...state.usage, context_max: 40000, context_percent: 8, context_used: 3100 } + })) + onEvent({ + payload: { + message: 'turn_failed', + reason: 'internal', + usage: { + completion_tokens: 0, + context_max: 40000, + context_percent: 18, + context_used: 7000, + prompt_tokens: 7000, + total_tokens: 7000 + } + }, + type: 'error' + } as any) + + expect(getUiState().usage).toMatchObject({ context_max: 40000, context_percent: 18, context_used: 7000 }) + }) + + it('leaves the context bar alone when the failure measured nothing', () => { + const onEvent = createGatewayEventHandler(buildCtx([])) + + patchUiState(state => ({ + ...state, + usage: { ...state.usage, context_max: 40000, context_percent: 8, context_used: 3100 } + })) + onEvent({ payload: { message: 'turn_failed', reason: 'internal' }, type: 'error' } as any) + + expect(getUiState().usage).toMatchObject({ context_max: 40000, context_percent: 8, context_used: 3100 }) + }) + it('shows setup panel for missing provider startup error', () => { const appended: Msg[] = [] const onEvent = createGatewayEventHandler(buildCtx(appended)) diff --git a/ui-tui/src/app/chatStream.ts b/ui-tui/src/app/chatStream.ts index d8a36c57f..55ccbdf72 100644 --- a/ui-tui/src/app/chatStream.ts +++ b/ui-tui/src/app/chatStream.ts @@ -625,6 +625,12 @@ const onError = ( const { reason, message } = ev.payload state.turns.delete(MAIN_VIEW_KEY) clearRunning(null) + // A turn that ended without an answer, failed or stopped, filed the reading + // its calls last measured with its record, so the bar moves on with it as + // `onMessageComplete` moves it for an answered one. + if (ev.payload.usage) { + patchUiState(s => ({ ...s, usage: { ...s.usage, ...ev.payload.usage } })) + } if (reason === 'cancelled_by_client') { restoreInputPrompt(appendMessage, sys) appendArtifacts(state, appendMessage) diff --git a/ui-tui/src/app/createGatewayEventHandler.ts b/ui-tui/src/app/createGatewayEventHandler.ts index e45ddb53c..959d77e68 100644 --- a/ui-tui/src/app/createGatewayEventHandler.ts +++ b/ui-tui/src/app/createGatewayEventHandler.ts @@ -747,6 +747,12 @@ export function createGatewayEventHandler(ctx: GatewayEventHandlerContext): (ev: case 'error': turnController.recordError() + // A failed turn still measured its calls and filed the reading, so the + // bar moves on with it rather than staying on the turn before. + if (ev.payload?.usage) { + patchUiState(state => ({ ...state, usage: { ...state.usage, ...ev.payload!.usage } })) + } + { const message = String(ev.payload?.message || 'unknown error') diff --git a/ui-tui/src/gatewayTypes.ts b/ui-tui/src/gatewayTypes.ts index f3b0263ea..7ee6d4a0b 100644 --- a/ui-tui/src/gatewayTypes.ts +++ b/ui-tui/src/gatewayTypes.ts @@ -698,4 +698,4 @@ export type GatewayEvent = session_id?: string type: 'message.complete' } - | { payload?: { message?: string }; session_id?: string; type: 'error' } + | { payload?: { message?: string; usage?: Usage }; session_id?: string; type: 'error' } diff --git a/ui-tui/src/rpc/generated.ts b/ui-tui/src/rpc/generated.ts index 2c512f948..60ad89be3 100644 --- a/ui-tui/src/rpc/generated.ts +++ b/ui-tui/src/rpc/generated.ts @@ -1859,9 +1859,23 @@ export interface ErrorEvent { * The failure's long account beside the bounded detail, for a reader who expands the row: the runner's own text, credentials scrubbed, capped at 8000 characters. Absent when the bounded text already is the whole account. */ full_detail?: string; + usage?: TurnUsage1; target?: DirectTarget; }; } +/** + * What a failed or stopped turn measured, in message.complete's shape: its last reported call, so a client's context gauge moves on with the turn's stored record. Absent when that call read no window. + */ +export interface TurnUsage1 { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + cost_usd?: number | null; + context_used?: number; + context_max?: number; + context_percent?: number; + cost_missing_calls?: number; +} /** * This interface was referenced by `RavenRpcRoot`'s JSON-Schema * via the `definition` "CronDeliveredEvent". diff --git a/ui-web/src/rpc/generated.ts b/ui-web/src/rpc/generated.ts index ee0a9d89a..a738b42fd 100644 --- a/ui-web/src/rpc/generated.ts +++ b/ui-web/src/rpc/generated.ts @@ -1509,9 +1509,23 @@ export interface ErrorEvent { * The failure's long account beside the bounded detail, for a reader who expands the row: the runner's own text, credentials scrubbed, capped at 8000 characters. Absent when the bounded text already is the whole account. */ full_detail?: string; + usage?: TurnUsage1; target?: DirectTarget; }; } +/** + * What a failed or stopped turn measured, in message.complete's shape: its last reported call, so a client's context gauge moves on with the turn's stored record. Absent when that call read no window. + */ +export interface TurnUsage1 { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + cost_usd?: number | null; + context_used?: number; + context_max?: number; + context_percent?: number; + cost_missing_calls?: number; +} export interface CronDeliveredEvent { type: 'cron.delivered'; payload: { diff --git a/ui-web/src/state/ctxChip.ts b/ui-web/src/state/ctxChip.ts index 792262aca..b9355c665 100644 --- a/ui-web/src/state/ctxChip.ts +++ b/ui-web/src/state/ctxChip.ts @@ -98,7 +98,7 @@ export function draw(): void { bearing for as long as nobody greps for its readers. Marking an estimate is worth doing; it is a change to what the page says, so it belongs in a change about that and not in this one. */ -export function set(nextUsed: number, nextMax: number): void { +export function set(nextUsed: number | undefined, nextMax: number | undefined): void { if (typeof nextMax === 'number' && nextMax > 0) max = nextMax if (typeof nextUsed === 'number' && nextUsed >= 0) used = nextUsed draw() diff --git a/ui-web/src/state/session/ctxReport.test.ts b/ui-web/src/state/session/ctxReport.test.ts new file mode 100644 index 000000000..846250e01 --- /dev/null +++ b/ui-web/src/state/session/ctxReport.test.ts @@ -0,0 +1,43 @@ +// @vitest-environment happy-dom +import { afterEach, beforeEach, describe, expect, it } from 'vitest' + +import { resetTranslator, setCode } from '../../i18n/t' +import { _resetForTests, get } from '../ctxChip' +import { applyUsage } from './ctxReport' + +beforeEach(() => { + resetTranslator() + setCode('en') + _resetForTests() +}) + +afterEach(() => { + _resetForTests() + resetTranslator() +}) + +describe('a usage report on the context ring', () => { + it('draws the window the report measured', () => { + applyUsage({ context_used: 7000, context_max: 40000 }) + expect(get().tip).toBe('Context 7.0k/40k 18%') + }) + + it('stands the fallback in only for a report that measured no window', () => { + applyUsage({ context_used: 0, context_max: 40000 }, 3100) + expect(get().tip).toBe('Context 3.1k/40k 8%') + applyUsage({ context_used: 7000, context_max: 40000 }, 3100) + expect(get().tip).toBe('Context 7.0k/40k 18%') + }) + + it("keeps a resumed session's zero, which is an empty session and not a gap", () => { + applyUsage({ context_used: 7000, context_max: 40000 }) + applyUsage({ context_used: 0, context_max: 40000 }) + expect(get().tip).toBe('Context 0/40k 0%') + }) + + it('moves nothing without a report', () => { + applyUsage({ context_used: 7000, context_max: 40000 }) + applyUsage(undefined) + expect(get().tip).toBe('Context 7.0k/40k 18%') + }) +}) diff --git a/ui-web/src/state/session/ctxReport.ts b/ui-web/src/state/session/ctxReport.ts new file mode 100644 index 000000000..13a7a4052 --- /dev/null +++ b/ui-web/src/state/session/ctxReport.ts @@ -0,0 +1,23 @@ +/* A usage report onto the composer's context ring. + * + * Three reports carry how full the window is: a finished turn's + * message.complete, the error a turn that ended without an answer sends, and a + * resumed session's usage. They all reach the ring through here, so the fields + * are read one way whichever of them arrives. + */ + +import { set } from '../ctxChip' + +/* The two fields of the wire's usage shape the ring is drawn from. */ +export interface UsageReport { + context_used?: number + context_max?: number +} + +/* `fallbackUsed` stands in for a report that measured no window, and only a + finished turn has one to give: its prompt count. A resumed session's zero is + a reading -- an empty session -- so without a fallback it is drawn as zero. */ +export function applyUsage(u: UsageReport | undefined, fallbackUsed?: number): void { + if (!u) return + set(fallbackUsed !== undefined && !u.context_used ? fallbackUsed : u.context_used, u.context_max) +} diff --git a/ui-web/src/state/session/pipeline.events.test.ts b/ui-web/src/state/session/pipeline.events.test.ts index af9a4c649..fd763a8be 100644 --- a/ui-web/src/state/session/pipeline.events.test.ts +++ b/ui-web/src/state/session/pipeline.events.test.ts @@ -875,6 +875,14 @@ describe('tool.complete', () => { }) describe('message.complete', () => { + it('draws the prompt count when the turn measured no window', async () => { + const h = await harness() + + h.dispatch({ type: 'message.complete', payload: { usage: { prompt_tokens: 1200, context_max: 40000 } } }) + + expect(h.did('setCtx')).toEqual([['setCtx', 1200, 40000]]) + }) + it('folds the turn (050-turn.js:192)', async () => { const h = await harness({ rows: [{ id: 's1', title: 'a deck' }] }) h.dispatch({ type: 'token.delta', payload: { text: 'the answer' } }) @@ -939,6 +947,58 @@ describe('error', () => { expect(h.did('noteRow')).toEqual([['noteRow', 'FAILED_LABEL', 'short', ['full', 'retry']]]) }) + + it('moves the context ring to the window the failed turn measured', async () => { + const h = await harness() + + h.dispatch({ + type: 'error', + payload: { + message: 'turn_failed', + reason: 'internal', + detail: 'no content', + usage: { prompt_tokens: 7000, completion_tokens: 0, total_tokens: 7000, context_used: 7000, context_max: 40000 }, + }, + }) + + expect(h.did('setCtx')).toEqual([['setCtx', 7000, 40000]]) + }) + + it('leaves the context ring alone when the failure measured nothing', async () => { + const h = await harness() + + h.dispatch({ type: 'error', payload: { message: 'turn_failed', reason: 'internal', detail: 'boom' } }) + + expect(h.did('setCtx')).toEqual([]) + }) + + it("draws a failed turn's reading the way a completion draws the same one", async () => { + const usage = { prompt_tokens: 7000, completion_tokens: 0, total_tokens: 7000, context_used: 7000, context_max: 40000 } + const done = await harness() + done.dispatch({ type: 'message.complete', payload: { usage } }) + const failed = await harness() + + failed.dispatch({ type: 'error', payload: { message: 'turn_failed', reason: 'internal', detail: 'boom', usage } }) + + expect(failed.did('setCtx')).toEqual(done.did('setCtx')) + }) + + /* `cancelling` is this window's own stop, which the arm otherwise leaves + early; `stream` is a stop another window on the session asked for. */ + it.each(['cancelling', 'stream'])('moves the ring to what a stopped turn measured (%s)', async (phase) => { + const h = await harness({ phase, busy: true }) + + h.dispatch({ + type: 'error', + payload: { + message: 'turn_cancelled', + reason: 'cancelled_by_client', + usage: { prompt_tokens: 6000, completion_tokens: 200, total_tokens: 6200, context_used: 6200, context_max: 40000 }, + }, + }) + + expect(h.did('setCtx')).toEqual([['setCtx', 6200, 40000]]) + }) }) describe('cron.delivered', () => { diff --git a/ui-web/src/state/session/registry.ts b/ui-web/src/state/session/registry.ts index 0ea6cc8cf..2ff634756 100644 --- a/ui-web/src/state/session/registry.ts +++ b/ui-web/src/state/session/registry.ts @@ -31,13 +31,13 @@ import { isImportedKey } from '../../lib/imported' import { current as sessionCurrent, setCurrent as sessionSet } from '../../lib/session' import { gateway } from '../../rpc/gateway' import { draw as drawBanner } from '../banner' -import { set as setCtx } from '../ctxChip' import { adopt as adoptBases, clearStaged as clearStagedBases } from '../mentions' import { load as loadTier } from '../tier' import { show as toast } from '../toast' import { clearStaged as clearStagedWorkdir } from '../workdir' import { reset as wsReset, setOpen as setWs } from '../ws' import { noteRow, pitch, readOnly, unpitch } from './conversation' +import { applyUsage } from './ctxReport' import { dropAll as dropHeldHosts } from './hosts' import { park, resume as resumeInPlace } from './residency' import { refreshDag, resume as resumeConversation } from './resume' @@ -372,7 +372,7 @@ export async function switchTo(s: SessRow, turnOver = false): Promise { /* context_estimated rides along in this payload and is not passed on: the ring has nowhere to say an estimate, so the writer takes two numbers. See state/ctxChip.ts. */ - setCtx(u.context_used, u.context_max) + applyUsage(u) /* The turn this conversation is in the middle of, which nothing else on a page that has just loaded can know: its frames went to a socket this page did not have. Put the machine back into the state `message.start` would diff --git a/ui-web/src/state/session/runtime.ts b/ui-web/src/state/session/runtime.ts index 902bd7505..7cd5d3713 100644 --- a/ui-web/src/state/session/runtime.ts +++ b/ui-web/src/state/session/runtime.ts @@ -38,13 +38,13 @@ import { current as sessionCurrent, setCurrent as sessionSet } from '../../lib/s import { hasNamingFlag, hasTurnDuration } from '../../rpc/capabilities' import { gateway } from '../../rpc/gateway' import { ask as confirmAsk } from '../confirm' -import { set as setCtx } from '../ctxChip' import { adopt as adoptBases, staged as stagedBases } from '../mentions' import { ds, sources } from '../sources' import { load as loadTier } from '../tier' import { show as toast } from '../toast' import { clearStaged as clearStagedWorkdir, staged as stagedWorkdir } from '../workdir' import { ask, noteRow, noteSay, pitch, splitAtts } from './conversation' +import { applyUsage } from './ctxReport' import { generation } from './generation' import { beginNaming, namingDeclined } from './naming' import { get, isActiveRuntime, isDraft as registryIsDraft, mint, subscribe, viewRuntime } from './registry' @@ -262,7 +262,7 @@ export function finishTurn(payload: unknown, rt: SessionRuntime = viewRuntime()) turn.dispatch({ type: 'idle' }) const inTok = usage.input_tokens || usage.prompt_tokens || 0 /* The window fill is the turn's prompt, not the running total. */ - setCtx(usage.context_used || inTok, usage.context_max as number) + applyUsage(usage, inTok) const s = sess(sessionCurrent()) /* A silent turn delivered nothing: it writes no preview, does not move its row up the rail (the server's list skips it for its time too) and raises no diff --git a/ui-web/src/state/session/stages.ts b/ui-web/src/state/session/stages.ts index c6107638b..164deb707 100644 --- a/ui-web/src/state/session/stages.ts +++ b/ui-web/src/state/session/stages.ts @@ -33,6 +33,7 @@ import { session as sheetSession } from '../sheetRack' import { ds, sources } from '../sources' import { show as toast } from '../toast' import { ask, noteRow, unask } from './conversation' +import { applyUsage } from './ctxReport' import { namingEnded, settleNaming } from './naming' import { rereadChips, viewRuntime } from './registry' import { @@ -293,6 +294,12 @@ export const STAGES: readonly Stage[] = [ notice, and dropping it keeps the queue held for the turn that is starting. Any silent mark that turn set is dropped with it. */ if (rt.heldTurn !== null && p.turn_id === rt.heldTurn) { rt.silent = false; return } + /* A turn that ended without an answer, failed or stopped, still measured + its calls, and the record it filed carries the last reading, so the ring + moves on with it rather than staying on the turn before -- otherwise a + reload reads a number the ring never showed. Ahead of the stop's own + early exits: a stopped turn's reading rides only this event. */ + applyUsage(p.usage) /* A cancelled turn is the one "error" a person asked for; the event still matters when the cancel came from ANOTHER client on the same session. */ if (p.reason === 'cancelled_by_client') {