fix(agent): recover broken threads and lazily create scenario runs

Two recurring failures from live logs (conversations 69651ca1 / a8c0dff8):

1. A checkpoint whose AI messages carry tool_calls without ToolMessages
   (run crashed after the LLM emitted a call) makes every send raise
   INVALID_CHAT_HISTORY with no recovery. The send path now repairs the
   thread via _repair_pending_tool_calls: pending calls are answered with
   synthetic error ToolMessages (THREAD_REPAIRED) so the user can retry.

2. Scenario tools scheduled from plain chat (no build_dashboard_test_scenario
   UI intent) had no durable AgentRun, so the resume fallback refused with
   SCENARIO_RUN_REQUIRED. _ensure_scenario_run now lazily creates the run
   from the tool args (dashboard_context from scenario_json for
   validate/resolve), mirroring the UIContextV2 scenario contract.
This commit is contained in:
2026-08-19 20:02:12 +03:00
parent 614f675f64
commit 7cc3297f1e
3 changed files with 292 additions and 1 deletions

View File

@@ -83,6 +83,71 @@ async def _resolve_scenario_run_id(conversation_id: str) -> str:
# #endregion AgentChat.Confirmation.ResolveScenarioRunId
# #region AgentChat.Confirmation.EnsureScenarioRun [C:2] [TYPE Function] [SEMANTICS agent-chat,resume,agent-run,lazy-create]
# @ingroup AgentChat
# @BRIEF Create a durable AgentRun on the fly when a scenario tool was scheduled without one.
# @POST Returns the new run_id, or "" when the tool args carry no dashboard identity or creation fails.
# @RATIONALE The send path only creates the durable run when the UI context intent is
# build_dashboard_test_scenario. When the LLM schedules a scenario tool from plain chat
# (intent missing), the resume fallback previously refused with SCENARIO_RUN_REQUIRED,
# killing the flow. Creating the run lazily from the tool args (dashboard_context from
# scenario_json for validate/resolve) restores durability. Mirrors the UIContextV2
# contract: scenario intent requires contextVersion=2 and objectType=dashboard.
async def _ensure_scenario_run(
tool_args: dict[str, Any],
conversation_id: str,
) -> str:
dashboard_id = tool_args.get("dashboard_id")
env_id = str(tool_args.get("environment_id") or "")
dashboard_name = str(tool_args.get("dashboard_name") or "")
if dashboard_id is None:
scenario_json = tool_args.get("scenario_json")
if isinstance(scenario_json, str) and scenario_json.strip():
try:
dashboard_context = (json.loads(scenario_json).get("dashboard_context") or {})
dashboard_id = dashboard_context.get("dashboard_id")
env_id = env_id or str(dashboard_context.get("environment_id") or "")
dashboard_name = dashboard_name or str(dashboard_context.get("dashboard_name") or "")
except (ValueError, TypeError):
pass
if dashboard_id is None:
return ""
try:
import os
from ss_tools.agent._config import FASTAPI_URL
from ss_tools.agent._run_tracker import RunTracker
from ss_tools.agent.context import get_user_jwt
from ss_tools.shared.logger import logger
backend_url = os.environ.get("BACKEND_URL") or FASTAPI_URL or "http://localhost:8000"
tracker = RunTracker(backend_url, get_user_jwt() or "")
context: dict[str, Any] = {
"objectType": "dashboard",
"objectId": str(dashboard_id),
"objectName": dashboard_name or None,
"envId": env_id or "default",
"route": f"/dashboards/{dashboard_id}",
"contextVersion": 2,
"intent": "build_dashboard_test_scenario",
}
run_id = await tracker.create(context, conversation_id=conversation_id)
logger.reason(
"Durable run created lazily in resume fallback",
payload={"run_id": run_id, "conv_id": conversation_id},
extra={"src": "AgentChat.Confirmation"},
)
return run_id
except Exception as exc:
logger.explore(
"Failed to lazily create durable run in resume fallback",
payload={"conv_id": conversation_id}, error=str(exc),
extra={"src": "AgentChat.Confirmation"},
)
return ""
# #endregion AgentChat.Confirmation.EnsureScenarioRun
# #region AgentChat.Confirmation.ToolAcceptsRunId [C:2] [TYPE Function] [SEMANTICS agent-chat,resume,agent-run,tools]
# @ingroup AgentChat
# @BRIEF Return True when a tool's args schema accepts an agent_run_id field.
@@ -589,6 +654,20 @@ async def handle_resume( # noqa: C901
run_id = get_agent_run_id()
if run_id:
tool_args = {**tool_args, "agent_run_id": run_id}
if _tool_accepts_agent_run_id(tool_obj) and not tool_args.get("agent_run_id"):
# No run bound (plain-chat scenario scheduling): create
# the durable run lazily from the tool args instead of
# refusing the whole flow with SCENARIO_RUN_REQUIRED.
lazy_run_id = await _ensure_scenario_run(tool_args, conversation_id)
if lazy_run_id:
reset_agent_run_id(agent_run_token)
agent_run_token = set_agent_run_id(lazy_run_id)
tool_args = {**tool_args, "agent_run_id": lazy_run_id}
logger.reason(
"Scenario tool bound to lazily created run",
payload={"tool": tool_name, "run_id": lazy_run_id, "conv_id": conversation_id},
extra={"src": "AgentChat.Confirmation"},
)
if _tool_accepts_agent_run_id(tool_obj) and not tool_args.get("agent_run_id"):
err = ("agent_run_id is required for scenario operations; "
"no durable run is bound to this conversation")

View File

@@ -32,17 +32,19 @@ import gradio as gr
import httpx
from jose import JWTError
from langchain_core.exceptions import OutputParserException
from langchain_core.messages import HumanMessage
from langchain_core.messages import HumanMessage, ToolMessage
from langchain_openai import ChatOpenAI
from openai import APIConnectionError, APITimeoutError, AuthenticationError, RateLimitError
from ss_tools.agent._config import GRADIO_ROOT_PATH, GRADIO_SERVER_NAME, GRADIO_SERVER_PORT, STORAGE_ROOT as _STORAGE_ROOT
from ss_tools.agent._confirmation import (
_await_if_async,
_pending_confirmations,
confirmation_payload,
handle_resume,
permission_denied_payload,
)
from ss_tools.agent._tool_resolver import pending_tool_calls_from_state
from ss_tools.agent._jwt_decoder import decode_token
from ss_tools.shared._llm_health import (
_LLM_CHECK_CACHE_TTL,
@@ -167,6 +169,40 @@ async def _build_agent_context(env_id: str | None) -> str:
# #endregion AgentChat.GradioApp.BuildAgentContext
# #region AgentChat.GradioApp.RepairBrokenThread [C:2] [TYPE Function] [SEMANTICS agent-chat,recovery,thread,repair]
# @ingroup AgentChat
# @BRIEF Repair a checkpoint whose AI messages have tool_calls without ToolMessages.
# @POST Answers every pending tool call with a synthetic error ToolMessage so the
# thread is consistent again and future sends/resumes stop raising
# LangGraph INVALID_CHAT_HISTORY. Returns the number of repaired calls.
# @RATIONALE A run that crashed after the LLM emitted a tool call (LLM provider
# flakiness, client disconnect) leaves the thread permanently broken: every
# subsequent message fails with INVALID_CHAT_HISTORY and there is no real tool
# result to replay. Synthetic error markers unblock the thread; the user retries.
async def _repair_pending_tool_calls(agent: Any, config: dict[str, Any]) -> int:
state = await agent.aget_state(config)
if state is None:
return 0
pending = pending_tool_calls_from_state(state)
if not pending:
return 0
current_msgs = list(state.values.get("messages", [])) if hasattr(state, "values") else []
repaired = [
ToolMessage(
content=(
f"Error: предыдущий запуск был прерван до выполнения инструмента "
f"{tool_name}; поток восстановлен автоматически. Повторите запрос."
),
tool_call_id=tcid,
name=tool_name,
)
for tool_name, _args, tcid in pending
]
await _await_if_async(agent.update_state(config, {"messages": [*current_msgs, *repaired]}))
return len(repaired)
# #endregion AgentChat.GradioApp.RepairBrokenThread
# #region AgentChat.GradioApp.TitleBestEffort [C:2] [TYPE Function] [SEMANTICS agent-chat,persistence,title]
# @ingroup AgentChat
# @BRIEF Run LLM title generation with a bounded timeout so request loops close cleanly.
@@ -1031,6 +1067,25 @@ async def agent_handler( # noqa: C901 — intentionally complex C4 orchestratio
_is_llm_error = any(p in str(exc).lower() for p in _llm_error_patterns)
if not _is_llm_error:
try:
# Broken-thread recovery: a checkpoint with AI messages whose
# tool_calls lack ToolMessages raises INVALID_CHAT_HISTORY on
# every send/resume (observed after LLM crashes). Repair it so
# the user can retry instead of hitting a dead end.
if await _repair_pending_tool_calls(agent, config) > 0:
logger.reason(
"Broken thread repaired in send path",
payload={"conv_id": conv_id},
extra={"src": "AgentChat.GradioApp.Handler"},
)
_request_result = "failed"
yield json.dumps(
{
"content": "⚠️ Предыдущий запуск был прерван. Поток восстановлен — отправьте запрос ещё раз.",
"metadata": {"type": "error", "code": "THREAD_REPAIRED", "detail": str(exc), "retryable": True},
}
)
await save_conversation(conv_id, visible_user_text, user_id, assistant_text="")
return
state = await agent.aget_state(config)
if getattr(state, "next", None):
yield confirmation_payload(conv_id, state, visible_user_text, user_role, env_id)

View File

@@ -0,0 +1,157 @@
# #region Test.Agent.Confirmation.Recovery [C:2] [TYPE Module] [SEMANTICS test,agent,confirmation,recovery,repair,run]
# @BRIEF Tests for broken-thread repair and lazy durable-run creation.
# @RELATION BINDS_TO -> [AgentChat.GradioApp.RepairBrokenThread]
# @RELATION BINDS_TO -> [AgentChat.Confirmation.EnsureScenarioRun]
# @TEST_EDGE broken_thread_repaired -> pending tool calls answered with ToolMessages
# @TEST_EDGE no_pending_calls -> repair is a no-op
# @TEST_EDGE lazy_run_from_compile_args -> run created with UIContextV2 scenario contract
# @TEST_EDGE lazy_run_from_scenario_json -> dashboard_context extracted for validate/resolve
# @TEST_EDGE lazy_run_no_identity -> returns "" without HTTP call
# @TEST_EDGE lazy_run_creation_failure -> returns "" and degrades gracefully
import os
from pathlib import Path
import sys
from unittest.mock import AsyncMock, MagicMock, patch
sys.path.append(str(Path(__file__).resolve().parent.parent.parent / "src"))
import pytest
os.environ.setdefault("BACKEND_URL", "http://test-backend:8000")
os.environ.setdefault("SERVICE_JWT", "test-service-jwt")
os.environ.setdefault("AUTH_SECRET_KEY", "test-secret-key-for-jwt-testing")
# ── Broken-thread repair (send path) ─────────────────────────────
class _FakeAgent:
def __init__(self, state):
self._state = state
self.updated = []
async def aget_state(self, config):
return self._state
async def update_state(self, config, values):
self.updated.append(values)
return {"configurable": config}
def _state_with_pending_tool_call():
from langchain_core.messages import AIMessage
state = MagicMock()
state.values.get.return_value = [
AIMessage(content="", tool_calls=[
{"id": "tc-1", "name": "scenario_compile", "args": {"dashboard_id": 5}, "type": "tool_call"},
]),
]
return state
@pytest.mark.anyio
async def test_repair_pending_tool_calls_answers_with_tool_messages():
from ss_tools.agent.app import _repair_pending_tool_calls
agent = _FakeAgent(_state_with_pending_tool_call())
count = await _repair_pending_tool_calls(agent, {"configurable": {"thread_id": "c1"}})
assert count == 1
assert len(agent.updated) == 1
msgs = agent.updated[0]["messages"]
from langchain_core.messages import ToolMessage
assert any(isinstance(m, ToolMessage) and m.tool_call_id == "tc-1" for m in msgs)
@pytest.mark.anyio
async def test_repair_no_pending_calls_is_noop():
from ss_tools.agent.app import _repair_pending_tool_calls
from langchain_core.messages import ToolMessage
state = MagicMock()
state.values.get.return_value = [
ToolMessage(content="ok", tool_call_id="tc-1", name="x"),
]
agent = _FakeAgent(state)
count = await _repair_pending_tool_calls(agent, {"configurable": {"thread_id": "c1"}})
assert count == 0
assert agent.updated == []
# ── Lazy durable-run creation (resume fallback) ──────────────────
class _FakeTracker:
def __init__(self, run_id="run-lazy-1", error=None):
self.run_id = run_id
self.error = error
self.created = []
async def create(self, context, conversation_id=None):
if self.error:
raise self.error
self.created.append((context, conversation_id))
return self.run_id
@pytest.mark.anyio
async def test_lazy_run_from_compile_args():
from ss_tools.agent._confirmation import _ensure_scenario_run
tracker = _FakeTracker()
with patch("ss_tools.agent._run_tracker.RunTracker", return_value=tracker), \
patch("ss_tools.agent.context.get_user_jwt", return_value="u-jwt"):
run_id = await _ensure_scenario_run(
{"dashboard_id": 5, "environment_id": "ss-dev", "dashboard_name": "COVID Vaccine Dashboard"},
"conv-1",
)
assert run_id == "run-lazy-1"
context, conversation_id = tracker.created[0]
assert context["objectId"] == "5"
assert context["envId"] == "ss-dev"
assert context["contextVersion"] == 2
assert context["intent"] == "build_dashboard_test_scenario"
assert context["route"] == "/dashboards/5"
assert conversation_id == "conv-1"
@pytest.mark.anyio
async def test_lazy_run_from_scenario_json_dashboard_context():
from ss_tools.agent._confirmation import _ensure_scenario_run
tracker = _FakeTracker()
scenario_json = (
'{"scenario_id": "d5-x", "dashboard_context": '
'{"environment_id": "ss-dev", "dashboard_id": 7, "dashboard_name": "Featured Charts"}, '
'"steps": []}'
)
with patch("ss_tools.agent._run_tracker.RunTracker", return_value=tracker), \
patch("ss_tools.agent.context.get_user_jwt", return_value="u-jwt"):
run_id = await _ensure_scenario_run({"scenario_json": scenario_json}, "conv-2")
assert run_id == "run-lazy-1"
assert tracker.created[0][0]["objectId"] == "7"
assert tracker.created[0][0]["envId"] == "ss-dev"
assert tracker.created[0][0]["objectName"] == "Featured Charts"
@pytest.mark.anyio
async def test_lazy_run_no_dashboard_identity_returns_empty():
from ss_tools.agent._confirmation import _ensure_scenario_run
tracker = _FakeTracker()
with patch("ss_tools.agent._run_tracker.RunTracker", return_value=tracker):
run_id = await _ensure_scenario_run({"objective_json": '{"goal": "x"}'}, "conv-3")
assert run_id == ""
assert tracker.created == []
@pytest.mark.anyio
async def test_lazy_run_creation_failure_returns_empty():
from ss_tools.agent._confirmation import _ensure_scenario_run
tracker = _FakeTracker(error=RuntimeError("backend down"))
with patch("ss_tools.agent._run_tracker.RunTracker", return_value=tracker), \
patch("ss_tools.agent.context.get_user_jwt", return_value="u-jwt"):
run_id = await _ensure_scenario_run({"dashboard_id": 5}, "conv-4")
assert run_id == ""
# #endregion Test.Agent.Confirmation.Recovery