Treat thinking as progress in stream timeout; trim replayed history

The 30s time-to-first-token guard measured from the turn's first model call
to the first answer token, so thinking, tool execution and the second call's
prefill all counted against it. Long reasoning on a long conversation was
cancelled mid-generation (llama-server: stop: cancel task at n_gen=1287).

- stream_query: the timeout now fires only when an in-flight model call goes
  silent for TTFT_TIMEOUT_SECONDS. Any chunk (thinking or empty) resets the
  clock, each model call starts a fresh one, and time spent in tools between
  calls doesn't count.
- Trim replayed history (blueprints/conversation/context.py): tool results
  and user messages from earlier turns are capped at 4000 chars (attachments
  included) and their images replaced by a placeholder. The current turn is
  untouched. The window no longer starts with orphaned ToolMessages.
- Log per-call prefill time, duration and prompt/output tokens
  (stream_usage=True) plus max_prompt_tokens in the stream_request line.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

blueprints/conversation/__init__.py

@@ -48,7 +48,35 @@ from .models import (
 from .tool_events import tool_end_payload, tool_name_from_message, tool_start_payload
 from .tool_labels import tool_label
 
-TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30"))
+# Longest an in-flight model call may go without producing any output (prompt
+# prefill included) before the turn is abandoned as stalled. Thinking tokens and
+# empty chunks count as output; time spent running tools does not count at all.
+MODEL_STALL_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30"))
+
+
+def stall_timeout(last_model_activity: float | None, now: float) -> float | None:
+    """Seconds left before the in-flight model call counts as stalled.
+
+    None when no model call is in flight (e.g. while tools run): the httpx read
+    timeout and the tools' own timeouts cover those phases.
+    """
+    if last_model_activity is None:
+        return None
+    return max(MODEL_STALL_TIMEOUT - (now - last_model_activity), 1.0)
+
+
+def _model_call_summary(call: dict) -> dict:
+    """Per-LLM-call timing and token counts for the stream_request log line."""
+    start = call["start"]
+    first_chunk = call.get("first_chunk")
+    end = call.get("end")
+    return {
+        "prefill_ms": round((first_chunk - start) * 1000) if first_chunk else None,
+        "ms": round((end - start) * 1000) if end else None,
+        "prompt_tokens": call.get("prompt_tokens"),
+        "output_tokens": call.get("output_tokens"),
+    }
+
 
 _THINK_OPEN = "<think>"
 _THINK_CLOSE = "</think>"
@@ -381,6 +409,9 @@ async def stream_query():
         first_content_at: float | None = None
         tool_timers: dict[str, float] = {}  # run_id -> start time
         tool_spans: list[dict] = []
+        # Last output from the in-flight model call; None between calls.
+        last_model_activity: float | None = None
+        model_calls: dict[str, dict] = {}  # run_id -> timing / token usage
         n_tokens = 0
         model_name: str | None = None
         cancelled = False
@@ -403,6 +434,8 @@ async def stream_query():
             ttft_ms = (
                 round((first_content_at - t0) * 1000) if first_content_at else None
             )
+            calls = [_model_call_summary(c) for c in model_calls.values()]
+            prompt_tokens = [c["prompt_tokens"] for c in calls if c["prompt_tokens"]]
             logging.info(
                 json.dumps(
                     {
@@ -414,6 +447,8 @@ async def stream_query():
                         "ttft_ms": ttft_ms,
                         "tokens": n_tokens,
                         "tools": tool_spans,
+                        "model_calls": calls,
+                        "max_prompt_tokens": max(prompt_tokens, default=None),
                         "model": model_name,
                         "cancelled": cancelled,
                         **({"error": error} if error else {}),
@@ -426,15 +461,10 @@ async def stream_query():
                 payload, version="v2", config=agent_config
             ).__aiter__()
             while True:
-                # Apply a short TTFT timeout only during the window between
-                # on_chat_model_start and the first content token. Before the
-                # LLM call starts (tool phase) and after the first token, let
-                # the httpx read timeout handle stalls instead.
-                if llm_start_at is not None and first_content_at is None:
-                    ttft_wait = TTFT_TIMEOUT - (time.monotonic() - llm_start_at)
-                    timeout = max(ttft_wait, 1.0)
-                else:
-                    timeout = None
+                # Only a silent model call counts as a stall. A model that is
+                # thinking (or streaming empty chunks) keeps resetting the clock,
+                # and each new call starts a fresh one.
+                timeout = stall_timeout(last_model_activity, time.monotonic())
 
                 try:
                     if timeout is not None:
@@ -446,7 +476,9 @@ async def stream_query():
                 except StopAsyncIteration:
                     break
                 except asyncio.TimeoutError:
-                    _log_summary(error=f"TTFT timeout after {TTFT_TIMEOUT:.0f}s")
+                    _log_summary(
+                        error=f"model stalled: no output for {MODEL_STALL_TIMEOUT:.0f}s"
+                    )
                     yield f"data: {json.dumps({'type': 'error', 'message': 'Simba took too long to respond. Please try again.', 'stalled': True})}\n\n"
                     return
 
@@ -498,14 +530,34 @@ async def stream_query():
 
                 elif event_type == "on_chat_model_start":
                     turn_has_content = False
+                    now = time.monotonic()
+                    last_model_activity = now
+                    model_calls[event["run_id"]] = {"start": now}
                     if llm_start_at is None:
-                        llm_start_at = time.monotonic()
+                        llm_start_at = now
                     if model_name is None:
                         model_name = event.get("metadata", {}).get(
                             "ls_model_name"
                         ) or event.get("name")
 
+                elif event_type == "on_chat_model_end":
+                    last_model_activity = None
+                    call = model_calls.get(event["run_id"])
+                    if call is not None:
+                        call["end"] = time.monotonic()
+                        usage = getattr(
+                            event.get("data", {}).get("output"), "usage_metadata", None
+                        )
+                        if usage:
+                            call["prompt_tokens"] = usage.get("input_tokens")
+                            call["output_tokens"] = usage.get("output_tokens")
+
                 elif event_type == "on_chat_model_stream":
+                    now = time.monotonic()
+                    last_model_activity = now
+                    call = model_calls.get(event.get("run_id"))
+                    if call is not None and "first_chunk" not in call:
+                        call["first_chunk"] = now
                     chunk = event.get("data", {}).get("chunk")
                     raw = getattr(chunk, "content", None)
                     if isinstance(raw, str) and raw:

blueprints/conversation/agents.py

@@ -24,6 +24,7 @@ from blueprints.conversation.artifacts_util import (
     truncate,
     unique,
 )
+from blueprints.conversation.context import trim_history
 from blueprints.conversation.memory import save_memory
 from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
 from blueprints.rag.logic import query_vector_store
@@ -54,6 +55,8 @@ if llama_url:
         timeout=httpx.Timeout(LLAMA_READ_TIMEOUT, connect=10.0),
         max_retries=1,
         model_kwargs={"reasoning_effort": "low"},
+        # Ask llama-server for token usage on the final chunk (prompt size logging).
+        stream_usage=True,
     )
 else:
     llama_chat = None
@@ -1106,10 +1109,8 @@ def _build_prompt(state, config: RunnableConfig = None) -> list:
     # LangGraph passes the full state dict {"messages": [...]}, not just the list
     messages = state["messages"] if isinstance(state, dict) else state
     non_system = [m for m in messages if not isinstance(m, SystemMessage)]
-    # Keep only the most recent messages to stay within the 32k context window.
-    # Tool call/result pairs each count as 2 messages, so 12 ≈ 4-6 conversational turns.
-    non_system = non_system[-12:]
-    return [SystemMessage(content=system_prompt)] + non_system
+    # Window + trim earlier turns to stay within the 32k context window.
+    return [SystemMessage(content=system_prompt)] + trim_history(non_system)
 
 
 def build_agent(checkpointer=None, extra_tools=()):

blueprints/conversation/context.py

@@ -0,0 +1,67 @@
+"""Shape the message history replayed to the LLM so it fits the context window.
+
+The LangGraph checkpoint keeps every message at full size: whole search hits,
+full notes, attached files, images. Replaying all of that each turn pushes the
+prompt toward the 32k-token llama-server slot, which makes every prefill slow.
+The current turn keeps everything; earlier turns are trimmed.
+"""
+
+from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage
+
+# Matches TOOL_RESULT_MAX_CHARS in logic.py (what gets persisted to Postgres).
+EARLIER_MESSAGE_MAX_CHARS = 4000
+TRUNCATION_NOTE = "\n[truncated: earlier context]"
+IMAGE_PLACEHOLDER = "[image from an earlier message]"
+# 12 ≈ 4-6 conversational turns (tool call/result pairs count as 2 messages).
+HISTORY_WINDOW = 12
+
+
+def _cap(text: str) -> str:
+    if len(text) <= EARLIER_MESSAGE_MAX_CHARS:
+        return text
+    return text[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE
+
+
+def _shrink(message: BaseMessage) -> BaseMessage:
+    """Cap text and drop images from a message that belongs to an earlier turn."""
+    content = message.content
+    if isinstance(content, str):
+        new_content = _cap(content)
+    else:
+        new_content = []
+        for part in content:
+            if isinstance(part, dict) and part.get("type") in ("image_url", "image"):
+                new_content.append({"type": "text", "text": IMAGE_PLACEHOLDER})
+            elif isinstance(part, dict) and isinstance(part.get("text"), str):
+                new_content.append({**part, "text": _cap(part["text"])})
+            elif isinstance(part, str):
+                new_content.append(_cap(part))
+            else:
+                new_content.append(part)
+    if new_content == content:
+        return message
+    return message.model_copy(update={"content": new_content})
+
+
+def trim_history(messages: list[BaseMessage]) -> list[BaseMessage]:
+    """Window the history and trim tool results/user messages from earlier turns.
+
+    The current turn starts at the last HumanMessage; it and everything after it
+    (this turn's tool calls and results) are left untouched.
+    """
+    current_turn = max(
+        (i for i, m in enumerate(messages) if isinstance(m, HumanMessage)),
+        default=len(messages),
+    )
+    trimmed = [
+        _shrink(m)
+        if i < current_turn and isinstance(m, (ToolMessage, HumanMessage))
+        else m
+        for i, m in enumerate(messages)
+    ]
+    window = trimmed[-HISTORY_WINDOW:]
+    # A window that starts mid tool-exchange would begin with ToolMessages whose
+    # tool_calls were cut off; OpenAI-style APIs reject those.
+    while window and isinstance(window[0], ToolMessage):
+        window = window[1:]
+    return window