"""The streaming turn loop: model rounds, tool dispatch and the final intercept."""
import asyncio
import inspect
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING
from collections.abc import AsyncGenerator
import structlog
from navi.config import settings
from navi.exceptions import (
ContextTooLargeError,
MaxIterationsReached,
SessionNotFound,
)
from navi.llm.base import Message
from navi.session_files import save_images as save_session_images
from navi.tools._internal.base import (
ToolContext,
current_plan_runner,
current_stop_event,
current_user_info,
current_user_role,
)
from navi.tools.plan import PlanRunner
from ..agent_run_context import AgentTurnContext, StreamState
from ..anti_stall import AntiStallMonitor
from ..events import (
AgentEvent,
ModelInfo,
StreamEnd,
StreamStopped,
TextDelta,
ThinkingDelta,
ThinkingEnd,
)
from ..image_tokens import context_image_tokens, estimate_image_tokens
from ..stream_guard import _iter_stream_guarded
log = structlog.get_logger()
from .helpers import (
_FINAL_INTERCEPT_NUDGES,
_PROTOCOL_FINAL_NUDGE,
_inline_images_notice,
_protocol_progress_message,
_todo_progress_message,
)
class AgentStreamMixin:
async def run_stream(
self,
session_id: str,
user_message: str,
images: list[str] | None = None,
display_message: str | None = None,
files: list[dict] | None = None,
is_recall: bool = False,
cwd: str | None = None,
hidden: bool = False,
) -> AsyncGenerator[AgentEvent, None]:
"""
Streaming variant. Yields AgentEvent objects:
- ThinkingDelta / ThinkingEnd: reasoning chunks
- ToolEvent: tool call + result
- TextDelta / StreamEnd: final streamed response
- ContextCompressed: emitted by workers after compression
Args:
user_message: The text sent to the LLM (may contain injected hints).
display_message: Optional text shown to the user in the chat UI.
When omitted, user_message is used for both LLM context and display.
hidden: If True, create a single user message with is_display=False
and is_context=True. Used for form submissions and other invisible
user input that must reach the LLM but should never appear in the UI.
"""
session = await self._sessions.get(session_id)
if session is None:
raise SessionNotFound(session_id)
profile = self._profiles.get(session.profile_id)
_scope_token, _keyless_token = await self._apply_mcp_scope(session)
tools = self._tool_list(profile.get_agent_tools())
tool_schemas = [t.schema() for t in tools]
llm = self._get_backend(profile.llm_backend)
# The profile these four are bound to. The post-turn reload must compare
# against this, not against session.profile_id: after a switch_profile the
# two can be equal again (the store rewrote the row, or mutated this very
# object) while the tool set above is still the old profile's.
bound_profile_id = profile.id
mem = await self._ctx_builder._memory_msg(user_id=session.user_id)
# Expose session_id, model and cwd to tools via ContextVar
from navi.tools._internal.base import (
current_allowed_mcp_servers as _mcp_scope_var,
)
from navi.tools._internal.base import (
current_keyless_mcp_servers as _keyless_var,
)
from navi.tools._internal.base import (
current_model as _model_var,
)
from navi.tools._internal.base import (
current_profile_id as _profile_var,
)
from navi.tools._internal.base import (
current_route as _route_var,
)
from navi.tools._internal.base import (
current_session_id as _sid_var,
)
from navi.tools._internal.base import (
current_working_directory as _cwd_var,
)
_sid_token = _sid_var.set(session_id)
_model_var.set(profile.model)
_route_var.set(getattr(profile, "llm_route", None))
_profile_var.set(profile.id)
_cwd_token = _cwd_var.set(cwd)
# Start the turn on the route's primary entry when it is healthy
# again: the background probe is what notices a recovered host, and
# this is where that knowledge becomes visible to the user — the next
# turn, never a running one. A turn that starts while the cloud is
# down begins on the local entry and stays there until it ends.
begin_turn = getattr(llm, "begin_turn", None)
if callable(begin_turn):
begin_turn(getattr(profile, "llm_route", None))
# Pre-turn compression: if the last turn filled the context past the
# threshold, compress NOW before calling the LLM. This prevents the
# model from seeing an over-full context and generating gibberish
# (e.g. a "summary of the conversation" instead of a real answer).
# The incoming message is priced here because the gate runs before it
# is appended — an attachment is the one part of the next prompt this
# context cannot yet account for.
pending = Message(role="user", content=user_message, images=images or None)
pending_tokens = self._compressor.estimate_tokens_delta([pending], profile.model)
async for _ev in self._compression_events_preturn(
session, llm, profile, session_id, pending_tokens=pending_tokens
):
yield _ev
if hidden:
# Hidden user input (e.g. form submission): one message, context-only,
# never shown in the UI. No image budgeting is applied because hidden
# messages are expected to carry plain JSON payloads.
user_msg = Message(
role="user",
content=user_message,
images=images or None,
files=files or None,
created_at=datetime.now(timezone.utc),
is_recall=is_recall,
is_display=False,
is_context=True,
)
session.messages.append(user_msg)
session.context.append(user_msg)
await self._sessions.save(session)
else:
display_text = display_message if display_message is not None else user_message
user_msg_display = Message(role="user", content=display_text, images=images or None,
files=files or None, created_at=datetime.now(timezone.utc),
is_recall=is_recall, is_context=False)
# Inline chat images land in the session directory too, and the paths go into
# the context next to the images themselves. A pasted image used to exist in
# the transcript only, so an agent told the session directory exists went
# looking there for the "attached images", found nothing, and spent turns on
# fruitless `find` calls before giving up.
saved_paths: list[str] = []
if images:
try:
names = save_session_images(session_id, images)
except OSError:
log.exception("agent.image_save_failed", session_id=session_id)
names = []
saved_paths = [
str(Path(settings.session_files_dir) / session_id / name) if name else ""
for name in names
]
# Image token budgeting: fit as many images as possible into the LLM context.
# Images that do not fit stay reachable through the session directory, so the
# notice has to name them. Each image is priced by its own size: a flat 500
# admitted twice as many screenshots as the model can take (a 1024x1024 image
# is ~1026 tokens on qwen3.6 — see navi/core/image_tokens.py).
images_for_context = images
context_content = user_message
if images:
current_tokens = self._compressor.estimate_tokens(
session.context, profile.model
)
available_tokens = int(self._window_for(profile) * 0.8) - current_tokens
fitted: list[str] = []
for img in images:
cost = estimate_image_tokens(img)
if cost > available_tokens:
break
fitted.append(img)
available_tokens -= cost
if len(fitted) < len(images):
images_for_context = fitted
context_content += "\n\n" + _inline_images_notice(
len(images), len(images_for_context or []), saved_paths
)
user_msg_context = Message(role="user", content=context_content, images=images_for_context or None,
files=files or None, created_at=datetime.now(timezone.utc),
is_recall=is_recall, is_display=False)
session.messages.append(user_msg_display)
session.messages.append(user_msg_context)
session.context.append(user_msg_context)
# Persist user message immediately so it survives a client disconnect
# before the assistant reply is ready.
await self._sessions.save(session)
stop_event = current_stop_event.get()
turn_ctx = AgentTurnContext(turn_start=time.monotonic())
# Final-turn intercept nudge pending injection into the next iteration's
# built_ctx. Lives in the generator's local scope so it survives across
# iterations of the for-loop (injections are ephemeral — built_ctx is
# rebuilt each iteration — so this is the channel that carries the
# nudge from the intercept site to the next turn's context).
pending_final_nudge: Message | None = None
# First assistant row this turn wrote. Its global display rank is what
# the client id (`h_<rank>`) and the feedback key are built from, so it
# is reported in stream_end / stream_stopped.
turn_first_assistant: Message | None = None
ctx_task = asyncio.create_task(self._ctx_builder._collect_context_injections(profile))
mem_facts_task = asyncio.create_task(
self._ctx_builder._memory_facts_msg(
user_message,
user_id=session.user_id,
injected_ids=turn_ctx.injected_fact_ids,
scope_boundary_enabled=profile.scope_boundary_enabled,
session_cwd=(session.session_metadata or {}).get("cwd"),
)
)
# Raw excerpts compression dropped. Deduped across turns through the
# session's ephemeral metadata: the same chunk should not ride along in
# every prompt of a long session just because the topic persists.
if session.session_metadata is None:
session.session_metadata = {}
# A list, not a set: session_metadata is JSON-encoded into the
# sessions row on every save, and a set is not serializable.
injected_chunks = session.session_metadata.setdefault(
"injected_chunk_ids", []
)
mem_recall_task = asyncio.create_task(
self._ctx_builder._memory_recall_msg(
session,
user_message,
user_id=session.user_id,
injected_ids=injected_chunks,
)
)
ctx_injections = await ctx_task
mem_facts = await mem_facts_task
if mem_facts:
ctx_injections.append(mem_facts)
log.debug(
"agent.memory_facts_injected",
session_id=session_id,
facts_msg_length=len(mem_facts.content or ""),
)
else:
log.debug("agent.memory_facts_none", session_id=session_id)
mem_recall = await mem_recall_task
if mem_recall:
ctx_injections.append(mem_recall)
# Drain background-task completion notes into the session context
# (persisted so compression keeps them; not displayed — the client
# already saw the live task_update event).
await self._drain_task_notes(session_id, session)
anti_stall = AntiStallMonitor(profile)
if profile.anti_stall_enabled:
await anti_stall.init(session_id)
# Tool-calling loop — uses stream_complete() for every turn so thinking
# is captured in real-time via ThinkingDelta/ThinkingEnd events.
for iteration in range(profile.max_iterations):
# Cooperative stop: check before each LLM call
if stop_event and stop_event.is_set():
await self._sessions.save(session)
yield StreamStopped()
return
async for _ev in self._compression_events_midturn(
session, llm, profile, session_id, iteration, ctx_injections, mem
):
yield _ev
state = StreamState()
built_ctx = self._ctx_builder.build(
session.context,
profile,
mem,
iteration=iteration,
max_iterations=profile.max_iterations,
extra_system=ctx_injections,
session_id=session_id,
session_metadata=session.session_metadata,
)
if (
profile.goal_anchoring_enabled
and iteration > 0
and iteration % profile.goal_anchoring_interval == 0
):
built_ctx.append(
await self._ctx_builder._build_goal_anchor(session_id, user_message)
)
todo_msg = await _todo_progress_message(session_id, first_iteration=(iteration == 0))
if todo_msg:
built_ctx.append(todo_msg)
protocol_msg = await _protocol_progress_message(session_id)
if protocol_msg:
built_ctx.append(protocol_msg)
# Anti-stall / adaptive re-plan: inject intervention message if needed
_stall_msg = await anti_stall.pre_turn(session_id, iteration)
if _stall_msg:
built_ctx.append(_stall_msg)
# Final-turn intercept nudge carried over from a previous bare-text
# turn: the model stopped to narrate while the todo had open steps,
# so we inject the recovery pointer and continue the loop.
if pending_final_nudge is not None:
built_ctx.append(pending_final_nudge)
pending_final_nudge = None
try:
self._compressor.check_context_size(
built_ctx,
session_context=session.context,
model=profile.model,
window=self._window_for(profile),
)
except ContextTooLargeError as e:
# Surface the error as a Navi response (not a raw system error) so the
# user sees a coherent message and the exchange is saved to history.
error_text = str(e)
assistant_msg = Message(role="assistant", content=error_text)
session.context.append(assistant_msg)
session.messages.append(assistant_msg)
if turn_first_assistant is None:
turn_first_assistant = assistant_msg
await self._sessions.save(session)
yield TextDelta(delta=error_text)
# This path was never exercised (the old chars//3 check never
# fired), so it kept an invalid StreamEnd(content=...) and died
# with a TypeError instead of reporting the overflow.
yield StreamEnd(
full_content=error_text,
max_context_tokens=self._window_for(profile),
message_index=await self._turn_message_index(
session_id, turn_first_assistant
),
)
return
async for _ev in self._consume_stream(
_iter_stream_guarded(
llm.stream_complete(
built_ctx,
tools=tool_schemas if tools else None,
temperature=profile.temperature,
model=profile.model,
think=profile.think_enabled,
top_k=profile.top_k,
top_p=profile.top_p,
num_thread=profile.num_thread,
# The window the compressor accounts against must be the
# window the local model is actually loaded with.
num_ctx=self._window_for(profile),
),
stop_event=stop_event,
first_chunk_timeout=settings.llm_stream_first_chunk_timeout,
chunk_timeout=settings.llm_stream_chunk_timeout,
),
stop_event,
turn_ctx,
state,
):
yield _ev
# Record the real prompt-token count for the context we just sent so
# the next iteration's midturn/check estimates use real tokens for
# the bulk of the context (not the chars//3 heuristic, which
# undercounts code-heavy tool output and fires compression too late).
# len(session.context) here is exactly what was sent — the new
# assistant/tool messages are appended only below.
# `sent_tokens` lets the compressor notice a prompt the backend cut
# down instead of caching the cut size as the truth (Ollama keeps
# roughly half the window and reports only that).
self._compressor.record_real_baseline(
len(session.context),
state.context_tokens,
chars=self._compressor.context_chars(built_ctx),
model=turn_ctx.resolved_model or profile.model,
sent_tokens=self._compressor.real_baseline_estimate(
session.context, built_ctx, turn_ctx.resolved_model or profile.model
),
image_tokens=context_image_tokens(built_ctx),
)
# Stopped mid-stream — save partial response and exit
if stop_event and stop_event.is_set():
if state.accumulated_text:
stopped_msg = Message(
role="assistant",
content=state.accumulated_text,
created_at=datetime.now(timezone.utc),
is_context=False,
)
session.messages.append(stopped_msg)
if turn_first_assistant is None:
turn_first_assistant = stopped_msg
await self._sessions.save(session)
yield StreamStopped(
message_index=await self._turn_message_index(
session_id, turn_first_assistant
)
)
return
turn_tool_calls = state.turn_tool_calls
if not turn_tool_calls:
# Final response — text already streamed above. The base message
# is built without metrics so it is shared by the intercept and the
# genuine-final paths; the model must SEE its own bare text next
# iteration when we intercept, so it goes into session.context now.
_elapsed = round(time.monotonic() - turn_ctx.turn_start, 1)
# Net tokens = accumulated across all iterations + subagent tokens
_net_tokens = turn_ctx.turn_tokens + turn_ctx.subagent_tokens
assistant_msg = Message(
role="assistant",
content=state.accumulated_text or None,
thinking=state.accumulated_thinking or None,
created_at=datetime.now(timezone.utc),
)
session.messages.append(assistant_msg)
session.context.append(assistant_msg)
if turn_first_assistant is None:
turn_first_assistant = assistant_msg
if self._compressor.context_truncated:
# A truncated prompt is reported as what the backend
# evaluated — roughly half the window — while the context
# that produced it was already over the window. Showing
# that number would drop the bar to ~50% at the exact
# moment the session is full; keep the honest estimate.
session.context_token_count = max(
state.context_tokens or 0,
self._compressor.real_baseline_estimate(
session.context, session.context, profile.model
),
)
else:
session.context_token_count = state.context_tokens or 0
# Final-turn intercept: the model stopped with bare text, but the
# anti-stall/adaptive nudges never fired (they live inside the
# tool-loop; a bare-text turn closes the run before post_turn).
# If the todo still has open steps — or a protocol run is open
# even though its steps are all closed — and we haven't hit the
# limit, queue a recovery nudge and continue the loop instead of
# closing.
from navi.tools.protocol import has_active_run
from navi.tools.todo import has_open_steps
may_intercept = (
profile.final_intercept_enabled
and turn_ctx.final_interceptions < profile.final_intercept_limit
)
todo_open = await has_open_steps(session_id) if may_intercept else False
protocol_open = await has_active_run(session_id) if may_intercept else False
if may_intercept and (todo_open or protocol_open):
turn_ctx.final_interceptions += 1
if protocol_open and not todo_open:
# A run whose steps are all closed needs a different
# closing move than "there is still work to do".
nudge = _PROTOCOL_FINAL_NUDGE
else:
nudge = _FINAL_INTERCEPT_NUDGES[
min(turn_ctx.final_interceptions, len(_FINAL_INTERCEPT_NUDGES)) - 1
]
pending_final_nudge = Message(role="system", content=nudge)
await self._sessions.save(session)
continue # No StreamEnd, no workers — the run continues.
# Genuine final: stamp metrics on the already-appended message
# (so history carries them) and close the run.
assistant_msg.elapsed_seconds = _elapsed
assistant_msg.tool_call_count = turn_ctx.tool_call_count if turn_ctx.tool_call_count else None
assistant_msg.token_count = _net_tokens if _net_tokens else None
await self._sessions.save(session)
yield StreamEnd(
full_content=state.accumulated_text,
context_tokens=state.context_tokens,
max_context_tokens=self._window_for(profile),
elapsed_seconds=_elapsed,
tool_call_count=turn_ctx.tool_call_count,
token_count=_net_tokens if _net_tokens else None,
message_index=await self._turn_message_index(
session_id, turn_first_assistant
),
)
for event in await self._run_workers(
session, llm, profile.model, state.context_tokens, profile=profile
):
yield event
return
# Tool calls turn — record to session and execute
assistant_msg = Message(
role="assistant",
content=state.accumulated_text or None,
thinking=state.accumulated_thinking or None,
tool_calls=turn_tool_calls,
)
session.messages.append(assistant_msg)
session.context.append(assistant_msg)
if turn_first_assistant is None:
turn_first_assistant = assistant_msg
# Persist the assistant's tool-call decision before executing tools, so
# a crash / CancelledError mid-tool-execution doesn't lose it.
await self._sessions.save(session)
tool_ctx = ToolContext(
session_id=session_id,
profile_id=profile.id,
event_sink=None, # set per-tool inside _execute_tools_with_sink
stop_event=stop_event,
model=profile.model,
user_id=session.user_id,
user_role=current_user_role.get(),
user_info=current_user_info.get(),
cwd=_cwd_var.get(),
)
# Parallel batch gate: profile override wins, else the global flag.
if turn_ctx.parallel_tool_calls is False:
if profile.parallel_tool_calls is None:
from navi.config import settings as _settings
turn_ctx.parallel_tool_calls = _settings.parallel_tool_calls
else:
turn_ctx.parallel_tool_calls = profile.parallel_tool_calls
# Expose a per-iteration PlanRunner so the `plan` tool can run the
# planner over the live session. Constructed inside the loop (not
# once before) using this iteration's profile/llm/tool_schemas so a
# post-switch_profile plan call uses the reloaded profile. reset on
# exit so the runner never leaks beyond the tool-execution window.
_plan_runner = PlanRunner(
self._planning, session, profile, llm, mem, tool_schemas
)
_plan_token = current_plan_runner.set(_plan_runner)
try:
async for _ev in self._execute_tools_with_sink(
turn_tool_calls, tools, turn_ctx, session, stop_event, tool_ctx
):
yield _ev
finally:
current_plan_runner.reset(_plan_token)
# The todo tool may have been called this turn (set/update/clear).
# Re-read and push the latest state to the UI's side panel.
yield await self._todo_updated(session)
# 6. Cooperative stop: check after tool execution before next LLM call
if stop_event and stop_event.is_set():
await self._sessions.save(session)
# Reset context vars before returning so they don't leak.
_cwd_var.reset(_cwd_token)
_sid_var.reset(_sid_token)
_mcp_scope_var.reset(_scope_token)
_keyless_var.reset(_keyless_token)
yield StreamStopped(
message_index=await self._turn_message_index(
session_id, turn_first_assistant
)
)
return
# Update anti-stall counters
await anti_stall.post_turn(session_id, turn_tool_calls)
# 7. If switch_profile was called this iteration, reload profile + tools.
# Compare against bound_profile_id, never against session.profile_id:
# set_profile() either rewrites the row (and this run's own save()
# then writes its stale value back over it) or mutates the very
# Session object the run holds — in both cases session.profile_id
# and the stored row agree again, the check stayed false and the
# turn kept calling the old profile's tools. The bindings at the
# top of the run are the only truth about what tools are live.
fresh = await self._sessions.get(session_id)
if fresh and fresh.profile_id != bound_profile_id:
session.profile_id = fresh.profile_id
profile = self._profiles.get(session.profile_id)
bound_profile_id = profile.id
self._set_active_profile(profile)
# The rest of this run executes as the new profile — tools that
# resolve a default profile must not keep reporting the old one.
_profile_var.set(profile.id)
tools = self._tool_list(profile.get_agent_tools())
tool_schemas = [t.schema() for t in tools]
llm = self._get_backend(profile.llm_backend)
log.info(
"agent.profile_reloaded",
session_id=session_id,
new_profile=session.profile_id,
)
await self._sessions.save(session)
# Reset cwd ContextVar so it does not leak into other calls.
_cwd_var.reset(_cwd_token)
_sid_var.reset(_sid_token)
_mcp_scope_var.reset(_scope_token)
_keyless_var.reset(_keyless_token)
raise MaxIterationsReached(profile.max_iterations)
async def _turn_message_index(
self, session_id: str, msg: Message | None
) -> int | None:
"""Global display rank of the turn's first assistant row, or None.
The web client keys its `h_<n>` ids and the rating feedback on this
rank, and it is what a reload recomputes from ``metadata.display_index``
— so a live turn and the same turn after a reload agree. Cosmetic
metadata: a store that cannot answer, a fake store in a test, or a
failed query must never break the turn; the client falls back to its own
approximation when this is None.
"""
if msg is None or msg.sequence_number is None or msg.sequence_number < 0:
return None
fn = getattr(self._sessions, "display_index_of", None)
if fn is None:
return None
try:
index = fn(session_id, msg.sequence_number)
if inspect.isawaitable(index):
index = await index
except Exception: # metadata, never fatal to the turn
log.warning(
"agent.message_index_failed",
session_id=session_id,
sequence_number=msg.sequence_number,
exc_info=True,
)
return None
# A MagicMock store or an unexpected type counts as "cannot answer".
return index if isinstance(index, int) and not isinstance(index, bool) else None
async def _consume_stream(
self, stream_gen, stop_event, turn_ctx: AgentTurnContext, state: StreamState
):
async for chunk in stream_gen:
if stop_event and stop_event.is_set():
if state.thinking_active:
yield ThinkingEnd()
break
if chunk.prompt_tokens is not None:
state.context_tokens = chunk.prompt_tokens
if chunk.completion_tokens is not None:
turn_ctx.turn_tokens += chunk.completion_tokens
# Surface the model that actually served this turn to the UI. The
# backend stamps it on the first successful chunk (fallback may pick
# a different model than the profile's first priority). Re-emit only
# when it changes across iterations within the same turn.
if chunk.model and chunk.model != turn_ctx.resolved_model:
turn_ctx.resolved_model = chunk.model
yield ModelInfo(
model=chunk.model,
target=getattr(chunk, "target", None),
fallback=bool(getattr(chunk, "is_fallback", False)),
)
if chunk.thinking:
state.accumulated_thinking += chunk.thinking
if not state.thinking_active:
state.thinking_active = True
yield ThinkingDelta(delta=chunk.thinking)
if chunk.delta:
if state.thinking_active:
state.thinking_active = False
yield ThinkingEnd()
state.accumulated_text += chunk.delta
yield TextDelta(delta=chunk.delta)
if chunk.tool_calls:
state.turn_tool_calls = chunk.tool_calls
if chunk.finish_reason and state.thinking_active:
state.thinking_active = False
yield ThinkingEnd()
