diff --git a/.env.example b/.env.example
index 7084c7d..2fafeb8 100644
--- a/.env.example
+++ b/.env.example
@@ -24,6 +24,8 @@ BASE_URL=192.168.1.5:8000
# If set, uses llama-server as the primary LLM backend with OpenAI as fallback
LLAMA_SERVER_URL=http://192.168.1.213:8080/v1
LLAMA_MODEL_NAME=llama-3.1-8b-instruct
+# Abandon a turn when an in-flight model call produces no output (prefill included)
+# for this long. Thinking tokens count as output, so long reasoning is fine.
TTFT_TIMEOUT_SECONDS=30
# Embedding Server Configuration
diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py
index 33434c7..e2c2752 100644
--- a/blueprints/conversation/__init__.py
+++ b/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_CLOSE = ""
@@ -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:
diff --git a/blueprints/conversation/agents.py b/blueprints/conversation/agents.py
index 16a00b3..c9079ac 100644
--- a/blueprints/conversation/agents.py
+++ b/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=()):
diff --git a/blueprints/conversation/context.py b/blueprints/conversation/context.py
new file mode 100644
index 0000000..1e0afb5
--- /dev/null
+++ b/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
diff --git a/tests/unit/test_context_trim.py b/tests/unit/test_context_trim.py
new file mode 100644
index 0000000..88fb7f8
--- /dev/null
+++ b/tests/unit/test_context_trim.py
@@ -0,0 +1,78 @@
+"""Tests for trimming the replayed history (blueprints/conversation/context.py)."""
+
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+
+from blueprints.conversation.context import (
+ EARLIER_MESSAGE_MAX_CHARS,
+ HISTORY_WINDOW,
+ IMAGE_PLACEHOLDER,
+ TRUNCATION_NOTE,
+ trim_history,
+)
+
+BIG = "x" * (EARLIER_MESSAGE_MAX_CHARS * 3)
+
+
+def _call(call_id: str, name: str = "simba_search") -> AIMessage:
+ return AIMessage(content="", tool_calls=[{"id": call_id, "name": name, "args": {}}])
+
+
+def _result(call_id: str, text: str = BIG) -> ToolMessage:
+ return ToolMessage(content=text, tool_call_id=call_id, name="simba_search")
+
+
+def test_earlier_tool_results_are_capped_current_turn_untouched():
+ messages = [
+ HumanMessage(content="first question"),
+ _call("c1"),
+ _result("c1"),
+ AIMessage(content="first answer"),
+ HumanMessage(content="second question"),
+ _call("c2"),
+ _result("c2"),
+ ]
+
+ out = trim_history(messages)
+
+ earlier, current = out[2], out[6]
+ assert earlier.content == BIG[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE
+ assert current.content == BIG
+ # The checkpointed originals are never mutated.
+ assert messages[2].content == BIG
+
+
+def test_earlier_attachments_and_images_are_shrunk():
+ image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
+ messages = [
+ HumanMessage(content=[{"type": "text", "text": BIG}, image]),
+ AIMessage(content="nice photo"),
+ HumanMessage(content=[{"type": "text", "text": "now this one"}, image]),
+ ]
+
+ out = trim_history(messages)
+
+ earlier = out[0].content
+ assert earlier[0]["text"].endswith(TRUNCATION_NOTE)
+ assert earlier[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
+ # The current turn keeps its image.
+ assert out[2].content[1] == image
+
+
+def test_window_never_starts_with_orphaned_tool_results():
+ # Many parallel tool results so the window cut lands between an AIMessage's
+ # tool_calls and their ToolMessages.
+ messages = [HumanMessage(content="q"), _call("c0")]
+ messages += [_result("c0", "r") for _ in range(HISTORY_WINDOW)]
+ messages += [AIMessage(content="answer"), HumanMessage(content="next")]
+
+ out = trim_history(messages)
+
+ assert len(out) <= HISTORY_WINDOW
+ assert not isinstance(out[0], ToolMessage)
+
+
+def test_short_history_is_returned_as_is():
+ messages = [HumanMessage(content="hi"), AIMessage(content="hello")]
+ out = trim_history(messages)
+ assert out == messages
+ assert out[0] is messages[0]
diff --git a/tests/unit/test_stream_stall.py b/tests/unit/test_stream_stall.py
new file mode 100644
index 0000000..629cf7f
--- /dev/null
+++ b/tests/unit/test_stream_stall.py
@@ -0,0 +1,203 @@
+"""Tests for the stream-query stall timeout and per-call model metrics."""
+
+import asyncio
+import json
+import logging
+
+import pytest
+from langchain_core.messages import AIMessage, AIMessageChunk
+from quart import Quart
+from quart_jwt_extended import JWTManager, create_access_token
+from tortoise import Tortoise
+
+import blueprints.conversation as conv
+from blueprints.conversation.models import Conversation
+from blueprints.users.models import User
+
+# Shortened stall limit for the tests (seconds).
+LIMIT = 0.3
+
+
+@pytest.fixture
+async def owner():
+ await Tortoise.init(
+ db_url="sqlite://:memory:",
+ modules={
+ "models": [
+ "blueprints.users.models",
+ "blueprints.conversation.models",
+ "blueprints.artifacts.models",
+ ]
+ },
+ )
+ await Tortoise.generate_schemas()
+ user = await User.create(username="owner", email="owner@example.com")
+ conversation = await Conversation.create(name="c", user=user)
+ yield user, conversation
+ await Tortoise.close_connections()
+
+
+class _TimedAgent:
+ """Yields (delay_seconds, event) pairs, sleeping before each event."""
+
+ def __init__(self, script):
+ self._script = script
+
+ async def astream_events(self, payload, version, config):
+ for delay, event in self._script:
+ await asyncio.sleep(delay)
+ yield event
+
+
+def _start(run_id):
+ return {"event": "on_chat_model_start", "run_id": run_id, "name": "llm", "data": {}}
+
+
+def _chunk(run_id, text):
+ return {
+ "event": "on_chat_model_stream",
+ "run_id": run_id,
+ "data": {"chunk": AIMessageChunk(content=text)},
+ }
+
+
+def _end(run_id, prompt_tokens=None):
+ usage = (
+ {
+ "input_tokens": prompt_tokens,
+ "output_tokens": 5,
+ "total_tokens": prompt_tokens + 5,
+ }
+ if prompt_tokens
+ else None
+ )
+ return {
+ "event": "on_chat_model_end",
+ "run_id": run_id,
+ "data": {"output": AIMessage(content="", usage_metadata=usage)},
+ }
+
+
+def _final(text):
+ return {
+ "event": "on_chain_end",
+ "name": "agent",
+ "data": {"output": {"messages": [AIMessage(content=text)]}},
+ }
+
+
+async def _run(owner, monkeypatch, script):
+ user, conversation = owner
+ monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", LIMIT)
+ monkeypatch.setattr(conv, "get_agent_for", lambda u: _TimedAgent(script))
+
+ async def _prompt(user_id):
+ return "system"
+
+ monkeypatch.setattr(conv, "build_system_prompt_with_memories", _prompt)
+
+ app = Quart(__name__)
+ app.config["JWT_SECRET_KEY"] = "test-secret"
+ JWTManager(app)
+ app.register_blueprint(conv.conversation_blueprint)
+ async with app.app_context():
+ token = create_access_token(identity=str(user.id))
+ resp = await app.test_client().post(
+ "/api/conversation/stream-query",
+ json={"query": "hi", "conversation_id": str(conversation.id)},
+ headers={"Authorization": f"Bearer {token}"},
+ )
+ body = (await resp.get_data()).decode()
+ return [
+ json.loads(line[len("data: ") :])
+ for line in body.splitlines()
+ if line.startswith("data: {")
+ ]
+
+
+def test_stall_timeout_only_while_a_model_call_is_in_flight(monkeypatch):
+ monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", 30.0)
+ assert conv.stall_timeout(None, 100.0) is None
+ assert conv.stall_timeout(100.0, 110.0) == 20.0
+ assert conv.stall_timeout(100.0, 200.0) == 1.0 # floor
+
+
+async def test_long_thinking_is_not_a_stall(owner, monkeypatch):
+ # The whole call takes ~4x the limit, but output never pauses for longer
+ # than half of it; the old time-to-first-content deadline killed this.
+ script = [(0, _start("m1"))]
+ script += [(LIMIT / 2, _chunk("m1", "hmm ")) for _ in range(8)]
+ script += [
+ (0, _chunk("m1", "Answer")),
+ (0, _end("m1")),
+ (0, _final("Answer")),
+ ]
+
+ frames = await _run(owner, monkeypatch, script)
+
+ types = [f["type"] for f in frames]
+ assert "error" not in types
+ assert frames[-1] == {"type": "response", "message": "Answer"}
+
+
+async def test_tool_time_between_model_calls_does_not_count(owner, monkeypatch):
+ tool = {"event": "on_tool_start", "name": "web_search", "run_id": "t1", "data": {}}
+ tool_end = {
+ "event": "on_tool_end",
+ "name": "web_search",
+ "run_id": "t1",
+ "data": {},
+ }
+ script = [
+ (0, _start("m1")),
+ (0, _chunk("m1", "")),
+ (0, _end("m1")),
+ (0, tool),
+ (LIMIT * 3, tool_end), # slow tool, no model call in flight
+ (0, _start("m2")),
+ (LIMIT / 2, _chunk("m2", "Done")),
+ (0, _end("m2")),
+ (0, _final("Done")),
+ ]
+
+ frames = await _run(owner, monkeypatch, script)
+
+ assert "error" not in [f["type"] for f in frames]
+ assert frames[-1]["message"] == "Done"
+
+
+async def test_silent_model_call_is_a_stall(owner, monkeypatch):
+ script = [
+ (0, _start("m1")),
+ (0, _chunk("m1", "start")),
+ (LIMIT * 4, _chunk("m1", "never arrives")),
+ ]
+
+ frames = await _run(owner, monkeypatch, script)
+
+ assert frames[-1]["type"] == "error"
+ assert frames[-1]["stalled"] is True
+
+
+async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
+ script = [
+ (0, _start("m1")),
+ (0, _chunk("m1", "Hi")),
+ (0, _end("m1", prompt_tokens=27262)),
+ (0, _final("Hi")),
+ ]
+
+ with caplog.at_level(logging.INFO):
+ await _run(owner, monkeypatch, script)
+
+ summaries = [
+ json.loads(r.getMessage())
+ for r in caplog.records
+ if r.getMessage().startswith('{"event": "stream_request"')
+ ]
+ assert summaries, "stream_request summary not logged"
+ summary = summaries[-1]
+ assert summary["max_prompt_tokens"] == 27262
+ (call,) = summary["model_calls"]
+ assert call["prompt_tokens"] == 27262
+ assert call["prefill_ms"] is not None and call["ms"] is not None