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:
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
157
agent/tests/test_agent/test_confirmation_recovery.py
Normal file
157
agent/tests/test_agent/test_confirmation_recovery.py
Normal 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
|
||||
Reference in New Issue
Block a user